DFHack script train-caged: Marks all caged, trainable creatures for training

train-caged.lua

-- Marks all caged, trainable creatures for training.
-- Author: Joel Verhagen (@joelverhagen on GitHub)

local utils = require('utils')

function getTrainingAssignment(unit)
    local training_assignment, _, _ = utils.binsearch(
        df.global.plotinfo.equipment.training_assignments,
        unit.id,
        'animal_id')
    return training_assignment
end

function isMarkedForTraining(unit)
    return getTrainingAssignment(unit) ~= nil
end

for i, unit in pairs(df.global.world.units.active) do
    -- Skip uncaged units
    if not unit.flags1.caged then
        goto continue
    end

    -- Skip units that cannot be trained or are already tame
    -- Source: https://github.com/DFHack/dfhack/blob/50.09-r3rc2/plugins/logistics.cpp#L403C6-L403C6
    if dfhack.units.isInvader(unit) or
        not dfhack.units.isTamable(unit) or
        dfhack.units.isTame(unit) then
        goto continue
    end

    local creature_raw = df.creature_raw.find(unit.race)
    local unit_name = creature_raw.creature_id

    if isMarkedForTraining(unit) then
        print(([[Already marked for training: %s %s]]):format(unit_name, unit.id))
        goto continue
    end

    local assignment = df.training_assignment:new()
    assignment.animal_id = unit.id
    assignment.trainer_id = -1
    assignment.flags.any_trainer = true
    utils.insert_sorted(df.global.plotinfo.equipment.training_assignments, assignment, 'animal_id')

    print(([[Marked for training: %s %s]]):format(unit_name, unit.id))

    ::continue::
end