show how many work animals a unit has when assigning

docs/changelog.txt

@@ -69,6 +69,7 @@ Template for new versions:
 - `buildingplan`: clarify interface when building single-tile staircases
 - `sort`: allow searching by profession on the squad assignment page
 - `sort`: add search for work animal assignment screen; allow filtering by miltary/squad/civilian/burrow
+- `sort`: new overlay on the animal assignment screen that shows how many work animals each visible unit already has assigned to them
 - `dreamfort`: Inside+ and Clearcutting burrows now automatically created and managed
 
 ## Documentation

library/modules/Gui.cpp

@@ -252,6 +252,8 @@ DEFINE_GET_FOCUS_STRING_HANDLER(dwarfmode)
                 newFocusString += "/ActivityDetails";
             else if (game->main_interface.info.creatures.adding_trainer)
                 newFocusString += "/AddingTrainer";
+            else if (game->main_interface.info.creatures.assign_work_animal)
+                newFocusString += "/AssignWorkAnimal";
             else
                 newFocusString += '/' + enum_item_key(game->main_interface.info.creatures.current_mode);
             break;

plugins/lua/sort.lua

@@ -1279,6 +1279,7 @@ OVERLAY_WIDGETS = {
     squad_assignment=SquadAssignmentOverlay,
     squad_annotation=SquadAnnotationOverlay,
     info=require('plugins.sort.info').InfoOverlay,
+    workanimals=require('plugins.sort.info').WorkAnimalOverlay,
     candidates=require('plugins.sort.info').CandidatesOverlay,
     interrogation=require('plugins.sort.info').InterrogationOverlay,
     location_selector=require('plugins.sort.locationselector').LocationSelectorOverlay,

plugins/lua/sort/info.lua

@@ -1,6 +1,7 @@
 local _ENV = mkmodule('plugins.sort.info')
 
 local gui = require('gui')
+local overlay = require('plugins.overlay')
 local sortoverlay = require('plugins.sort.sortoverlay')
 local widgets = require('gui.widgets')
 local utils = require('utils')
@@ -469,6 +470,73 @@ function CandidatesOverlay:onRenderBody(dc)
     end
 end
 
+-- ----------------------
+-- WorkAnimalOverlay
+--
+
+WorkAnimalOverlay = defclass(WorkAnimalOverlay, overlay.OverlayWidget)
+WorkAnimalOverlay.ATTRS{
+    default_pos={x=-33, y=12},
+    viewscreens='dwarfmode/Info/CREATURES/AssignWorkAnimal',
+    default_enabled=true,
+    frame={w=29, h=1},
+}
+
+function WorkAnimalOverlay:init()
+    self:addviews{
+        widgets.Label{
+            view_id='annotations',
+            frame={t=0, l=0},
+            text='',
+        }
+    }
+end
+
+local function get_work_animal_counts()
+    local counts = {}
+    for _,unit in ipairs(df.global.world.units.active) do
+        if not dfhack.units.isOwnCiv(unit) or
+            (not dfhack.units.isWar(unit) and not dfhack.units.isHunter(unit))
+        then
+            goto continue
+        end
+        local owner_id = unit.relationship_ids.Pet
+        if owner_id == -1 then goto continue end
+        counts[owner_id] = (counts[owner_id] or 0) + 1
+        ::continue::
+    end
+    return counts
+end
+
+function WorkAnimalOverlay:onRenderFrame(dc, rect)
+    local _, sh = dfhack.screen.getWindowSize()
+    local _, t = get_panel_offsets()
+    local list_height = sh - (17 + t)
+    local num_elems = list_height // 3
+    local max_elem = math.min(#creatures.work_animal_recipient-1,
+        creatures.scroll_position_work_animal+num_elems-1)
+
+    local annotations = {}
+    local counts = get_work_animal_counts()
+    for idx=creatures.scroll_position_work_animal,max_elem do
+        table.insert(annotations, NEWLINE)
+        table.insert(annotations, NEWLINE)
+        local animal_count = counts[creatures.work_animal_recipient[idx].id]
+        if animal_count and animal_count > 0 then
+            table.insert(annotations, {text='[', pen=COLOR_RED})
+            table.insert(annotations, ('Assigned work animals: %d'):format(animal_count))
+            table.insert(annotations, {text=']', pen=COLOR_RED})
+        end
+        table.insert(annotations, NEWLINE)
+    end
+
+    self.subviews.annotations.frame.t = t
+    self.subviews.annotations:setText(annotations)
+    self.frame.h = list_height + t
+
+    WorkAnimalOverlay.super.onRenderFrame(self, dc, rect)
+end
+
 -- ----------------------
 -- InterrogationOverlay
 --