remove-stress: fix -all for soul-less units, indentation, module support

lethosor · lethosor · commit 7b30d55e97f2 · 2018-05-22T10:01:34.000-04:00
diff --git a/remove-stress.lua b/remove-stress.lua
@@ -1,5 +1,6 @@
 -- Sets stress to negative one million
 --By Putnam; http://www.bay12forums.com/smf/index.php?topic=139553.msg5820486#msg5820486
+--@module = true
 local help = [====[
 
 remove-stress
@@ -12,27 +13,39 @@ Applies to the selected unit, or use ``remove-stress -all`` to apply to all unit
 
 local utils = require 'utils'
 
-validArgs = validArgs or utils.invert({
- 'help',
- 'all'
+function removeStress(unit)
+    if unit.status.current_soul then
+        unit.status.current_soul.personality.stress_level = -1000000
+    end
+end
+
+validArgs = utils.invert({
+    'help',
+    'all'
 })
 
-local args = utils.processArgs({...}, validArgs)
+function main(...)
+    local args = utils.processArgs({...}, validArgs)
+
+    if args.help then
+        print(help)
+        return
+    end
 
-if args.help then
-    print(help)
-    return
+    if args.all then
+        for k,v in ipairs(df.global.world.units.active) do
+            removeStress(v)
+        end
+    else
+        local unit = dfhack.gui.getSelectedUnit()
+        if unit then
+            removeStress(unit)
+        else
+            error 'Invalid usage: No unit selected and -all argument not given.'
+        end
+    end
 end
 
-if args.all then
- for k,v in ipairs(df.global.world.units.active) do
-  v.status.current_soul.personality.stress_level=-1000000
- end
-else
- local unit = dfhack.gui.getSelectedUnit()
- if unit then
-  unit.status.current_soul.personality.stress_level=-1000000
- else
-  error 'Invalid usage: No unit selected and -all argument not given.'
- end
+if not dfhack_flags.module then
+    main(...)
 end
