add training check functions to Units

docs/dev/Lua API.rst

@@ -1380,6 +1380,10 @@ Units module
 * ``dfhack.units.isTame(unit)``
 * ``dfhack.units.isTamable(unit)``
 * ``dfhack.units.isDomesticated(unit)``
+* ``dfhack.units.isMarkedForTraining(unit)``
+* ``dfhack.units.isMarkedForTaming(unit)``
+* ``dfhack.units.isMarkedForWarTraining(unit)``
+* ``dfhack.units.isMarkedForHuntTraining(unit)``
 * ``dfhack.units.isMarkedForSlaughter(unit)``
 * ``dfhack.units.isMarkedForGelding(unit)``
 * ``dfhack.units.isGeldable(unit)``

library/LuaApi.cpp

@@ -1777,6 +1777,10 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isTame),
     WRAPM(Units, isTamable),
     WRAPM(Units, isDomesticated),
+    WRAPM(Units, isMarkedForTraining),
+    WRAPM(Units, isMarkedForTaming),
+    WRAPM(Units, isMarkedForWarTraining),
+    WRAPM(Units, isMarkedForHuntTraining),
     WRAPM(Units, isMarkedForSlaughter),
     WRAPM(Units, isMarkedForGelding),
     WRAPM(Units, isGeldable),

library/include/modules/Units.h

@@ -109,6 +109,10 @@ DFHACK_EXPORT bool isWar(df::unit* unit);
 DFHACK_EXPORT bool isTame(df::unit* unit);
 DFHACK_EXPORT bool isTamable(df::unit* unit);
 DFHACK_EXPORT bool isDomesticated(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForTraining(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForTaming(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForWarTraining(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForHuntTraining(df::unit* unit);
 DFHACK_EXPORT bool isMarkedForSlaughter(df::unit* unit);
 DFHACK_EXPORT bool isMarkedForGelding(df::unit* unit);
 DFHACK_EXPORT bool isGeldable(df::unit* unit);

library/modules/Units.cpp

@@ -76,6 +76,7 @@ using namespace std;
 #include "df/tile_occupancy.h"
 #include "df/plotinfost.h"
 #include "df/syndrome.h"
+#include "df/training_assignment.h"
 #include "df/unit_inventory_item.h"
 #include "df/unit_misc_trait.h"
 #include "df/unit_relationship_type.h"
@@ -510,6 +511,38 @@ bool Units::isDomesticated(df::unit* unit)
     return tame;
 }
 
+static df::training_assignment * get_training_assignment(df::unit* unit) {
+    return binsearch_in_vector(df::global::plotinfo->equipment.training_assignments,
+        &df::training_assignment::animal_id, unit->id);
+}
+
+bool Units::isMarkedForTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    return !!get_training_assignment(unit);
+}
+
+bool Units::isMarkedForTaming(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && !assignment->flags.bits.train_war && !assignment->flags.bits.train_hunt;
+}
+
+bool Units::isMarkedForWarTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && assignment->flags.bits.train_war;
+}
+
+bool Units::isMarkedForHuntTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && assignment->flags.bits.train_hunt;
+}
+
 bool Units::isMarkedForSlaughter(df::unit* unit)
 {
     CHECK_NULL_POINTER(unit);