add unassignTrainer to balance assignTrainer

DFHack · Jan 27, 2024 · 75b0a99 · 75b0a99
1 parent 5476185
commit 75b0a99
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 3 deletions.
diff --git a/docs/changelog.txt b/docs/changelog.txt
@@ -65,10 +65,12 @@ Template for new versions:
 
 ## API
 - ``Units::assignTrainer``: assign a trainer to a trainable animal
+- ``Units::unassignTrainer``: unassign a trainer from an animal
 
 ## Lua
 - ``dfhack.units.isTamable``: return false for invaders to match vanilla logic
 - ``dfhack.units.assignTrainer``: expose new API to Lua
+- ``dfhack.units.unassignTrainer``: expose new API to Lua
 
 ## Removed
 - ``nopause``: functionality has moved to `spectate`

diff --git a/docs/dev/Lua API.rst b/docs/dev/Lua API.rst
@@ -1486,10 +1486,11 @@ Units module
   tile occupancy flags appropriately. Returns true if successful.
 
 * ``dfhack.units.assignTrainer(unit[, trainer_id])``
+* ``dfhack.units.unassignTrainer(unit)``
 
-  Assignes a trainer to the specified trainable unit. The trainer ID can be
-  omitted if "any trainer" is desired. Returns a boolean indicating whether the
-  operation was successful.
+  Assignes (or unassigns) a trainer for the specified trainable unit. The
+  trainer ID can be omitted if "any trainer" is desired. Returns a boolean
+  indicating whether the operation was successful.
 
 * ``dfhack.units.getGeneralRef(unit, type)``
 

diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
@@ -1823,6 +1823,7 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, setActionTimers),
     WRAPM(Units, setGroupActionTimers),
     WRAPM(Units, getUnitByNobleRole),
+    WRAPM(Units, unassignTrainer),
     { NULL, NULL }
 };
 

diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
@@ -177,6 +177,7 @@ DFHACK_EXPORT void setNickname(df::unit *unit, std::string nick);
 DFHACK_EXPORT df::language_name *getVisibleName(df::unit *unit);
 
 DFHACK_EXPORT bool assignTrainer(df::unit *unit, int32_t trainer_id = -1);
+DFHACK_EXPORT bool unassignTrainer(df::unit *unit);
 
 DFHACK_EXPORT df::identity *getIdentity(df::unit *unit);
 DFHACK_EXPORT df::nemesis_record *getNemesis(df::unit *unit);

diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
@@ -507,6 +507,12 @@ bool Units::assignTrainer(df::unit* unit, int32_t trainer_id) {
     return true;
 }
 
+bool Units::unassignTrainer(df::unit* unit) {
+    CHECK_NULL_POINTER(unit);
+    return erase_from_vector(plotinfo->equipment.training_assignments,
+        &df::training_assignment::animal_id, unit->id);
+}
+
 // check if creature is domesticated
 // seems to be the only way to really tell if it's completely safe to autonestbox it (training can revert)
 bool Units::isDomesticated(df::unit* unit)