Skip to content

Commit 0831e6e

Browse files
author
Jarkami
committed
Clean up/improve code readability in uniform-unstick
1 parent fb3f2d1 commit 0831e6e

File tree

1 file changed

+94
-68
lines changed

1 file changed

+94
-68
lines changed

uniform-unstick.lua

Lines changed: 94 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,15 @@ local validArgs = utils.invert({
1515

1616
-- Functions
1717

18+
-- @param item df.item
19+
-- @return string
1820
local function item_description(item)
1921
return "item #" .. item.id .. " '" .. dfhack.df2console(dfhack.items.getDescription(item, 0, true)) .. "'"
2022
end
2123

22-
local function get_item_pos(item)
24+
-- @param item df.item
25+
-- @return df.coord|nil
26+
local function get_visible_item_pos(item)
2327
local x, y, z = dfhack.items.getPosition(item)
2428
if not x or not y or not z then
2529
return
@@ -30,24 +34,30 @@ local function get_item_pos(item)
3034
end
3135
end
3236

33-
local function get_squad_position(unit, unit_name)
37+
-- @param unit df.unit
38+
-- @return df.squad_position|nil
39+
local function get_squad_position(unit)
3440
local squad = df.squad.find(unit.military.squad_id)
35-
if squad then
36-
if squad.entity_id ~= df.global.plotinfo.group_id then
37-
print("WARNING: Unit " .. unit_name .. " is a member of a squad from another site!" ..
38-
" This may be preventing them from doing any useful work." ..
39-
" You can fix this by assigning them to a local squad and then unassigning them.")
40-
print()
41-
return
42-
end
43-
else
41+
if not squad then
42+
return
43+
end
44+
45+
if squad.entity_id ~= df.global.plotinfo.group_id then
46+
print("WARNING: Unit " .. dfhack.df2console(dfhack.units.getReadableName(unit)) .. " is a member of a squad from another site!" ..
47+
" This may be preventing them from doing any useful work." ..
48+
" You can fix this by assigning them to a local squad and then unassigning them.")
49+
print()
4450
return
4551
end
52+
4653
if #squad.positions > unit.military.squad_position then
4754
return squad.positions[unit.military.squad_position]
4855
end
4956
end
5057

58+
-- @param unit df.unit
59+
-- @param item df.item
60+
-- @return number[] list of body part ids
5161
local function bodyparts_that_can_wear(unit, item)
5262
local bodyparts = {}
5363
local unitparts = dfhack.units.getCasteRaw(unit).body_info.body_parts
@@ -89,47 +99,61 @@ local function bodyparts_that_can_wear(unit, item)
8999
return bodyparts
90100
end
91101

92-
-- returns new value of need_newline
93-
local function print_line(text, need_newline)
94-
if need_newline then
95-
print()
96-
end
97-
print(text)
98-
return false
102+
-- @param unit_name string
103+
-- @param labor_name string
104+
local function print_bad_labor(unit_name, labor_name)
105+
return print("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
106+
" labor enabled, which conflicts with military uniforms.")
99107
end
100108

101-
local function print_bad_labor(unit_name, labor_name, need_newline)
102-
return print_line("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
103-
" labor enabled, which conflicts with military uniforms.", need_newline)
109+
-- @param squad_position df.squad_position
110+
-- @param item_id number
111+
local function remove_item_from_position(squad_position, item_id)
112+
for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
113+
for _, uniform_spec in ipairs(uniform_slot_specs) do
114+
for idx, assigned_item_id in ipairs(uniform_spec.assigned) do
115+
if assigned_item_id == item_id then
116+
uniform_spec.assigned:erase(idx)
117+
return
118+
end
119+
end
120+
end
121+
end
104122
end
105123

106124
-- Will figure out which items need to be moved to the floor, returns an item_id:item map
107-
local function process(unit, args, need_newline)
125+
local function process(unit, args)
108126
local silent = args.all -- Don't print details if we're iterating through all dwarves
109127
local unit_name = dfhack.df2console(dfhack.units.getReadableName(unit))
128+
local printed = false
110129

111130
if not silent then
112-
need_newline = print_line("Processing unit " .. unit_name, need_newline)
131+
print("Processing unit " .. unit_name)
132+
printed = true
113133
end
114134

115135
-- The return value
116136
local to_drop = {} -- item id to item object
117137

118138
-- First get squad position for an early-out for non-military dwarves
119-
local squad_position = get_squad_position(unit, unit_name)
139+
local squad_position = get_squad_position(unit)
120140
if not squad_position then
121141
if not silent then
122-
need_newline = print_line(unit_name .. " does not have a military uniform.", need_newline)
142+
print(unit_name .. " does not have a military uniform.")
143+
print()
123144
end
124145
return
125146
end
126147

127148
if unit.status.labors.MINE then
128-
need_newline = print_bad_labor(unit_name, "mining", need_newline)
149+
print_bad_labor(unit_name, "mining")
150+
printed = true
129151
elseif unit.status.labors.CUTWOOD then
130-
need_newline = print_bad_labor(unit_name, "woodcutting", need_newline)
152+
print_bad_labor(unit_name, "woodcutting")
153+
printed = true
131154
elseif unit.status.labors.HUNT then
132-
need_newline = print_bad_labor(unit_name, "hunting", need_newline)
155+
print_bad_labor(unit_name, "hunting")
156+
printed = true
133157
end
134158

135159
-- Find all worn items which may be at issue.
@@ -148,12 +172,12 @@ local function process(unit, args, need_newline)
148172
end
149173

150174
-- Now get info about which items have been assigned as part of the uniform
151-
local assigned_items = {} -- assigned item ids mapped to item objects
152-
for _, specs in ipairs(squad_position.equipment.uniform) do
153-
for _, spec in ipairs(specs) do
154-
for _, assigned in ipairs(spec.assigned) do
175+
local uniform_assigned_items = {} -- assigned item ids mapped to item objects
176+
for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
177+
for _, uniform_spec in ipairs(uniform_slot_specs) do
178+
for _, assigned_item_id in ipairs(uniform_spec.assigned) do
155179
-- Include weapon and shield so we can avoid dropping them, or pull them out of container/inventory later
156-
assigned_items[assigned] = df.item.find(assigned)
180+
uniform_assigned_items[assigned_item_id] = df.item.find(assigned_item_id)
157181
end
158182
end
159183
end
@@ -163,50 +187,48 @@ local function process(unit, args, need_newline)
163187

164188
local present_ids = {} -- map of item ID to item object
165189
local missing_ids = {} -- map of item ID to item object
166-
for u_id, item in pairs(assigned_items) do
167-
if not worn_items[u_id] then
190+
for item_id, item in pairs(uniform_assigned_items) do
191+
if not worn_items[item_id] then
168192
if not silent then
169-
need_newline = print_line(unit_name .. " is missing an assigned item, " .. item_description(item), need_newline)
193+
print(unit_name .. " is missing an assigned item, " .. item_description(item))
194+
printed = true
170195
end
171196
if dfhack.items.getGeneralRef(item, df.general_ref_type.UNIT_HOLDER) then
172-
need_newline = print_line(unit_name .. " cannot equip item: another unit has a claim on " .. item_description(item), need_newline)
197+
print(unit_name .. " cannot equip item: another unit has a claim on " .. item_description(item))
198+
printed = true
173199
if args.free then
174200
print(" Removing from uniform")
175-
assigned_items[u_id] = nil
176-
for _, specs in ipairs(squad_position.equipment.uniform) do
177-
for _, spec in ipairs(specs) do
178-
for idx, assigned in ipairs(spec.assigned) do
179-
if assigned == u_id then
180-
spec.assigned:erase(idx)
181-
break
182-
end
183-
end
184-
end
185-
end
201+
uniform_assigned_items[item_id] = nil
202+
remove_item_from_position(squad_position, item_id)
186203
end
187204
else
188-
missing_ids[u_id] = item
205+
missing_ids[item_id] = item
189206
if args.free then
190-
to_drop[u_id] = item
207+
to_drop[item_id] = item
191208
end
192209
end
193210
else
194-
present_ids[u_id] = item
211+
present_ids[item_id] = item
195212
end
196213
end
197214

198215
-- Make the equipment.assigned_items list consistent with what is present in equipment.uniform
199216
for i=#(squad_position.equipment.assigned_items)-1,0,-1 do
200-
local u_id = squad_position.equipment.assigned_items[i]
217+
local assigned_item_id = squad_position.equipment.assigned_items[i]
201218
-- Quiver, backpack, and flask are assigned in their own locations rather than in equipment.uniform, and thus need their own checks
202219
-- If more separately-assigned items are added in the future, this handling will need to be updated accordingly
203-
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
204-
local item = df.item.find(u_id)
220+
if uniform_assigned_items[assigned_item_id] == nil and
221+
assigned_item_id ~= squad_position.equipment.quiver and
222+
assigned_item_id ~= squad_position.equipment.backpack and
223+
assigned_item_id ~= squad_position.equipment.flask
224+
then
225+
local item = df.item.find(assigned_item_id)
205226
if item ~= nil then
206-
need_newline = print_line(unit_name .. " has an improperly assigned item, " .. item_description(item) .. '; removing it')
227+
print(unit_name .. " has an improperly assigned item, " .. item_description(item) .. "; removing it")
207228
else
208-
need_newline = print_line(unit_name .. " has a nonexistent item assigned, item # " .. u_id .. '; removing it')
229+
print(unit_name .. " has a nonexistent item assigned, item # " .. assigned_item_id .. "; removing it")
209230
end
231+
printed = true
210232
squad_position.equipment.assigned_items:erase(i)
211233
end
212234
end
@@ -217,10 +239,10 @@ local function process(unit, args, need_newline)
217239
-- unless --multi is specified, in which we don't care
218240
local covered = {} -- map of body part id to true/nil
219241
if not args.multi then
220-
for id, item in pairs(present_ids) do
242+
for item_id, item in pairs(present_ids) do
221243
-- weapons and shields don't "cover" the bodypart they're assigned to. (Needed to figure out if we're missing gloves.)
222244
if item._type ~= df.item_weaponst and item._type ~= df.item_shieldst then
223-
covered[worn_parts[id]] = true
245+
covered[worn_parts[item_id]] = true
224246
end
225247
end
226248
end
@@ -236,17 +258,23 @@ local function process(unit, args, need_newline)
236258
end
237259

238260
-- Drop everything (except uniform pieces) from body parts which should be covered but aren't
239-
for w_id, item in pairs(worn_items) do
240-
if assigned_items[w_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
241-
if uncovered[worn_parts[w_id]] then
242-
need_newline = print_line(unit_name .. " potentially has " .. item_description(item) .. " blocking a missing uniform item.", need_newline)
261+
for worn_item_id, item in pairs(worn_items) do
262+
if uniform_assigned_items[worn_item_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
263+
if uncovered[worn_parts[worn_item_id]] then
264+
print(unit_name .. " potentially has " .. item_description(item) .. " blocking a missing uniform item.")
265+
printed = true
243266
if args.drop then
244-
to_drop[w_id] = item
267+
to_drop[worn_item_id] = item
245268
end
246269
end
247270
end
248271
end
249272

273+
-- add a spacing line if there was any output
274+
if printed then
275+
print()
276+
end
277+
250278
return to_drop
251279
end
252280

@@ -255,8 +283,8 @@ local function do_drop(item_list)
255283
return
256284
end
257285

258-
for id, item in pairs(item_list) do
259-
local pos = get_item_pos(item)
286+
for _, item in pairs(item_list) do
287+
local pos = get_visible_item_pos(item)
260288
if not pos then
261289
dfhack.printerr("Could not find drop location for " .. item_description(item))
262290
else
@@ -278,10 +306,8 @@ local function main(args)
278306
end
279307

280308
if args.all then
281-
local need_newline = false
282309
for _, unit in ipairs(dfhack.units.getCitizens(true)) do
283-
do_drop(process(unit, args, need_newline))
284-
need_newline = true
310+
do_drop(process(unit, args))
285311
end
286312
else
287313
local unit = dfhack.gui.getSelectedUnit()

0 commit comments

Comments
 (0)