Merge pull request #2869 from robob27/moar-units-checks

Add isGeldable, isMarkedForGelding and isPet

docs/changelog.txt

@@ -45,6 +45,11 @@ changelog.txt uses a syntax similar to RST, with a few special sequences:
 
 ## API
 
+- Units module: added new predicates for:
+  - ``isGeldable()``
+  - ``isMarkedForGelding()``
+  - ``isPet()``
+
 ## Lua
 
 ## Removed

docs/dev/Lua API.rst

@@ -1343,6 +1343,8 @@ Units module
 * ``dfhack.units.isTamable(unit)``
 * ``dfhack.units.isDomesticated(unit)``
 * ``dfhack.units.isMarkedForSlaughter(unit)``
+* ``dfhack.units.isMarkedForGelding(unit)``
+* ``dfhack.units.isGeldable(unit)``
 * ``dfhack.units.isGelded(unit)``
 * ``dfhack.units.isEggLayer(unit)``
 * ``dfhack.units.isGrazer(unit)``
@@ -1362,7 +1364,7 @@ Units module
 
   The unit is available for adoption.
 
-
+* ``dfhack.units.isPet(unit)``
 * ``dfhack.units.isOpposedToLife(unit)``
 * ``dfhack.units.hasExtravision(unit)``
 * ``dfhack.units.isBloodsucker(unit)``

library/LuaApi.cpp

@@ -1756,6 +1756,8 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isTamable),
     WRAPM(Units, isDomesticated),
     WRAPM(Units, isMarkedForSlaughter),
+    WRAPM(Units, isMarkedForGelding),
+    WRAPM(Units, isGeldable),
     WRAPM(Units, isGelded),
     WRAPM(Units, isEggLayer),
     WRAPM(Units, isGrazer),
@@ -1763,6 +1765,7 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isForest),
     WRAPM(Units, isMischievous),
     WRAPM(Units, isAvailableForAdoption),
+    WRAPM(Units, isPet),
     WRAPM(Units, hasExtravision),
     WRAPM(Units, isOpposedToLife),
     WRAPM(Units, isBloodsucker),

library/include/modules/Units.h

@@ -110,6 +110,8 @@ DFHACK_EXPORT bool isTame(df::unit* unit);
 DFHACK_EXPORT bool isTamable(df::unit* unit);
 DFHACK_EXPORT bool isDomesticated(df::unit* unit);
 DFHACK_EXPORT bool isMarkedForSlaughter(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForGelding(df::unit* unit);
+DFHACK_EXPORT bool isGeldable(df::unit* unit);
 DFHACK_EXPORT bool isGelded(df::unit* unit);
 DFHACK_EXPORT bool isEggLayer(df::unit* unit);
 DFHACK_EXPORT bool isGrazer(df::unit* unit);
@@ -117,6 +119,7 @@ DFHACK_EXPORT bool isMilkable(df::unit* unit);
 DFHACK_EXPORT bool isForest(df::unit* unit);
 DFHACK_EXPORT bool isMischievous(df::unit *unit);
 DFHACK_EXPORT bool isAvailableForAdoption(df::unit* unit);
+DFHACK_EXPORT bool isPet(df::unit* unit);
 
 DFHACK_EXPORT bool hasExtravision(df::unit *unit);
 DFHACK_EXPORT bool isOpposedToLife(df::unit *unit);

library/modules/Units.cpp

@@ -511,6 +511,23 @@ bool Units::isMarkedForSlaughter(df::unit* unit)
     return unit->flags2.bits.slaughter == 1;
 }
 
+bool Units::isMarkedForGelding(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    return unit->flags3.bits.marked_for_gelding == 1;
+}
+
+bool Units::isGeldable(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+
+    if(world->raws.creatures.all[unit->race]->caste[unit->caste]->flags.is_set(caste_raw_flags::GELDABLE)) {
+        return true;
+    }
+
+    return false;
+}
+
 bool Units::isGelded(df::unit* unit)
 {
     CHECK_NULL_POINTER(unit);
@@ -587,6 +604,15 @@ bool Units::isAvailableForAdoption(df::unit* unit)
     return false;
 }
 
+bool Units::isPet(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+
+    if(unit->relationship_ids[df::unit_relationship_type::Pet] != -1)
+        return true;
+
+    return false;
+}
 
 bool Units::hasExtravision(df::unit *unit)
 {