Clean up/improve code readability in uniform-unstick

Author: Jarkami

uniform-unstick.lua

@@ -15,11 +15,15 @@ local validArgs = utils.invert({
 
 -- Functions
 
+-- @param item df.item
+-- @return string
 local function item_description(item)
     return "item #" .. item.id .. " '" .. dfhack.df2console(dfhack.items.getDescription(item, 0, true)) .. "'"
 end
 
-local function get_item_pos(item)
+-- @param item df.item
+-- @return df.coord|nil
+local function get_visible_item_pos(item)
     local x, y, z = dfhack.items.getPosition(item)
     if not x or not y or not z then
         return
@@ -30,24 +34,30 @@ local function get_item_pos(item)
     end
 end
 
-local function get_squad_position(unit, unit_name)
+-- @param unit df.unit
+-- @return df.squad_position|nil
+local function get_squad_position(unit)
     local squad = df.squad.find(unit.military.squad_id)
-    if squad then
-        if squad.entity_id ~= df.global.plotinfo.group_id then
-            print("WARNING: Unit " .. unit_name .. " is a member of a squad from another site!" ..
-                " This may be preventing them from doing any useful work." ..
-                " You can fix this by assigning them to a local squad and then unassigning them.")
-            print()
-            return
-        end
-    else
+    if not squad then
+        return
+    end
+
+    if squad.entity_id ~= df.global.plotinfo.group_id then
+        print("WARNING: Unit " .. dfhack.df2console(dfhack.units.getReadableName(unit)) .. " is a member of a squad from another site!" ..
+            " This may be preventing them from doing any useful work." ..
+            " You can fix this by assigning them to a local squad and then unassigning them.")
+        print()
         return
     end
+
     if #squad.positions > unit.military.squad_position then
         return squad.positions[unit.military.squad_position]
     end
 end
 
+-- @param unit df.unit
+-- @param item df.item
+-- @return number[] list of body part ids
 local function bodyparts_that_can_wear(unit, item)
     local bodyparts = {}
     local unitparts = dfhack.units.getCasteRaw(unit).body_info.body_parts
@@ -89,47 +99,61 @@ local function bodyparts_that_can_wear(unit, item)
     return bodyparts
 end
 
--- returns new value of need_newline
-local function print_line(text, need_newline)
-    if need_newline then
-        print()
-    end
-    print(text)
-    return false
+-- @param unit_name string
+-- @param labor_name string
+local function print_bad_labor(unit_name, labor_name)
+    return print("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
+        " labor enabled, which conflicts with military uniforms.")
 end
 
-local function print_bad_labor(unit_name, labor_name, need_newline)
-    return print_line("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
-        " labor enabled, which conflicts with military uniforms.", need_newline)
+-- @param squad_position df.squad_position
+-- @param item_id number
+local function remove_item_from_position(squad_position, item_id)
+    for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
+        for _, uniform_spec in ipairs(uniform_slot_specs) do
+            for idx, assigned_item_id in ipairs(uniform_spec.assigned) do
+                if assigned_item_id == item_id then
+                    uniform_spec.assigned:erase(idx)
+                    return
+                end
+            end
+        end
+    end
 end
 
 -- Will figure out which items need to be moved to the floor, returns an item_id:item map
-local function process(unit, args, need_newline)
+local function process(unit, args)
     local silent = args.all -- Don't print details if we're iterating through all dwarves
     local unit_name = dfhack.df2console(dfhack.units.getReadableName(unit))
+    local printed = false
 
     if not silent then
-        need_newline = print_line("Processing unit " .. unit_name, need_newline)
+        print("Processing unit " .. unit_name)
+        printed = true
     end
 
     -- The return value
     local to_drop = {} -- item id to item object
 
     -- First get squad position for an early-out for non-military dwarves
-    local squad_position = get_squad_position(unit, unit_name)
+    local squad_position = get_squad_position(unit)
     if not squad_position then
         if not silent then
-            need_newline = print_line(unit_name .. " does not have a military uniform.", need_newline)
+            print(unit_name .. " does not have a military uniform.")
+            print()
         end
         return
     end
 
     if unit.status.labors.MINE then
-        need_newline = print_bad_labor(unit_name, "mining", need_newline)
+        print_bad_labor(unit_name, "mining")
+        printed = true
     elseif unit.status.labors.CUTWOOD then
-        need_newline = print_bad_labor(unit_name, "woodcutting", need_newline)
+        print_bad_labor(unit_name, "woodcutting")
+        printed = true
     elseif unit.status.labors.HUNT then
-        need_newline = print_bad_labor(unit_name, "hunting", need_newline)
+        print_bad_labor(unit_name, "hunting")
+        printed = true
     end
 
     -- Find all worn items which may be at issue.
@@ -148,12 +172,12 @@ local function process(unit, args, need_newline)
     end
 
     -- Now get info about which items have been assigned as part of the uniform
-    local assigned_items = {} -- assigned item ids mapped to item objects
-    for _, specs in ipairs(squad_position.equipment.uniform) do
-        for _, spec in ipairs(specs) do
-            for _, assigned in ipairs(spec.assigned) do
+    local uniform_assigned_items = {} -- assigned item ids mapped to item objects
+    for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
+        for _, uniform_spec in ipairs(uniform_slot_specs) do
+            for _, assigned_item_id in ipairs(uniform_spec.assigned) do
                 -- Include weapon and shield so we can avoid dropping them, or pull them out of container/inventory later
-                assigned_items[assigned] = df.item.find(assigned)
+                uniform_assigned_items[assigned_item_id] = df.item.find(assigned_item_id)
             end
         end
     end
@@ -163,50 +187,48 @@ local function process(unit, args, need_newline)
 
     local present_ids = {} -- map of item ID to item object
     local missing_ids = {} -- map of item ID to item object
-    for u_id, item in pairs(assigned_items) do
-        if not worn_items[u_id] then
+    for item_id, item in pairs(uniform_assigned_items) do
+        if not worn_items[item_id] then
             if not silent then
-                need_newline = print_line(unit_name .. " is missing an assigned item, " .. item_description(item), need_newline)
+                print(unit_name .. " is missing an assigned item, " .. item_description(item))
+                printed = true
             end
             if dfhack.items.getGeneralRef(item, df.general_ref_type.UNIT_HOLDER) then
-                need_newline = print_line(unit_name .. " cannot equip item: another unit has a claim on " .. item_description(item), need_newline)
+                print(unit_name .. " cannot equip item: another unit has a claim on " .. item_description(item))
+                printed = true
                 if args.free then
                     print("  Removing from uniform")
-                    assigned_items[u_id] = nil
-                    for _, specs in ipairs(squad_position.equipment.uniform) do
-                        for _, spec in ipairs(specs) do
-                            for idx, assigned in ipairs(spec.assigned) do
-                                if assigned == u_id then
-                                    spec.assigned:erase(idx)
-                                    break
-                                end
-                            end
-                        end
-                    end
+                    uniform_assigned_items[item_id] = nil
+                    remove_item_from_position(squad_position, item_id)
                 end
             else
-                missing_ids[u_id] = item
+                missing_ids[item_id] = item
                 if args.free then
-                    to_drop[u_id] = item
+                    to_drop[item_id] = item
                 end
             end
         else
-            present_ids[u_id] = item
+            present_ids[item_id] = item
         end
     end
 
     -- Make the equipment.assigned_items list consistent with what is present in equipment.uniform
     for i=#(squad_position.equipment.assigned_items)-1,0,-1 do
-        local u_id = squad_position.equipment.assigned_items[i]
+        local assigned_item_id = squad_position.equipment.assigned_items[i]
         -- Quiver, backpack, and flask are assigned in their own locations rather than in equipment.uniform, and thus need their own checks
         -- If more separately-assigned items are added in the future, this handling will need to be updated accordingly
-        if assigned_items[u_id] == nil and u_id ~= squad_position.equipment.quiver and u_id ~= squad_position.equipment.backpack and u_id ~= squad_position.equipment.flask then
-            local item = df.item.find(u_id)
+        if uniform_assigned_items[assigned_item_id] == nil and
+            assigned_item_id ~= squad_position.equipment.quiver and
+            assigned_item_id ~= squad_position.equipment.backpack and
+            assigned_item_id ~= squad_position.equipment.flask
+        then
+            local item = df.item.find(assigned_item_id)
             if item ~= nil then
-                need_newline = print_line(unit_name .. " has an improperly assigned item, " .. item_description(item) .. '; removing it')
+                print(unit_name .. " has an improperly assigned item, " .. item_description(item) .. "; removing it")
             else
-                need_newline = print_line(unit_name .. " has a nonexistent item assigned, item # " .. u_id .. '; removing it')
+                print(unit_name .. " has a nonexistent item assigned, item # " .. assigned_item_id .. "; removing it")
             end
+            printed = true
             squad_position.equipment.assigned_items:erase(i)
         end
     end
@@ -217,10 +239,10 @@ local function process(unit, args, need_newline)
     -- unless --multi is specified, in which we don't care
     local covered = {} -- map of body part id to true/nil
     if not args.multi then
-        for id, item in pairs(present_ids) do
+        for item_id, item in pairs(present_ids) do
             -- weapons and shields don't "cover" the bodypart they're assigned to. (Needed to figure out if we're missing gloves.)
             if item._type ~= df.item_weaponst and item._type ~= df.item_shieldst then
-                covered[worn_parts[id]] = true
+                covered[worn_parts[item_id]] = true
             end
         end
     end
@@ -236,17 +258,23 @@ local function process(unit, args, need_newline)
     end
 
     -- Drop everything (except uniform pieces) from body parts which should be covered but aren't
-    for w_id, item in pairs(worn_items) do
-        if assigned_items[w_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
-            if uncovered[worn_parts[w_id]] then
-                need_newline = print_line(unit_name .. " potentially has " .. item_description(item) .. " blocking a missing uniform item.", need_newline)
+    for worn_item_id, item in pairs(worn_items) do
+        if uniform_assigned_items[worn_item_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
+            if uncovered[worn_parts[worn_item_id]] then
+                print(unit_name .. " potentially has " .. item_description(item) .. " blocking a missing uniform item.")
+                printed = true
                 if args.drop then
-                    to_drop[w_id] = item
+                    to_drop[worn_item_id] = item
                 end
             end
         end
     end
 
+    -- add a spacing line if there was any output
+    if printed then
+        print()
+    end
+
     return to_drop
 end
 
@@ -255,8 +283,8 @@ local function do_drop(item_list)
         return
     end
 
-    for id, item in pairs(item_list) do
-        local pos = get_item_pos(item)
+    for _, item in pairs(item_list) do
+        local pos = get_visible_item_pos(item)
         if not pos then
             dfhack.printerr("Could not find drop location for " .. item_description(item))
         else
@@ -278,10 +306,8 @@ local function main(args)
     end
 
     if args.all then
-        local need_newline = false
         for _, unit in ipairs(dfhack.units.getCitizens(true)) do
-            do_drop(process(unit, args, need_newline))
-            need_newline = true
+            do_drop(process(unit, args))
         end
     else
         local unit = dfhack.gui.getSelectedUnit()