use new unit module property functions

myk002 · myk002 · commit ee81c7396abf · 2024-08-09T21:33:47.000-07:00
add some annotations for units in gui/sitemap
clean up exterminate
diff --git a/agitation-rebalance.lua b/agitation-rebalance.lua
@@ -106,10 +106,6 @@ local function persist_state()
     dfhack.persistent.saveSiteData(GLOBAL_KEY, state)
 end
 
-local function is_agitated(unit)
-    return unit and unit.flags4.agitated_wilderness_creature
-end
-
 local world = df.global.world
 local map_features = world.features.map_features
 local plotinfo = df.global.plotinfo
@@ -237,7 +233,7 @@ end
 local function get_agitated_units()
     local agitators = {}
     for _, unit in ipairs(world.units.active) do
-        if is_unkilled(unit) and is_agitated(unit) then
+        if is_unkilled(unit) and dfhack.units.isAgitated(unit) then
             table.insert(agitators, unit)
         end
     end
@@ -250,7 +246,7 @@ local function check_new_unit(unit_id)
     if new_unit_min_frame_counter >= world.frame_counter then return end
     local unit = df.unit.find(unit_id)
     if not unit or not is_unkilled(unit) then return end
-    if state.features.surface and is_agitated(unit) then
+    if state.features.surface and dfhack.units.isAgitated(unit) then
         on_surface_attack()
         return
     end
diff --git a/changelog.txt b/changelog.txt
@@ -35,6 +35,8 @@ Template for new versions:
 - `timestream`: ensure child growth events (e.g. becoming an adult) are not skipped over
 
 ## Misc Improvements
+- `gui/sitemap`: show whether a unit is friendly, hostile, or wildlife
+- `gui/sitemap`: show whether a unit is caged
 
 ## Removed
 
diff --git a/docs/exterminate.rst b/docs/exterminate.rst
@@ -5,16 +5,15 @@ exterminate
     :summary: Kill things.
     :tags: fort armok units
 
-Kills any unit, or all undead, or all units of a given race. You can target any
-unit on a revealed tile of the map, including hidden ambushers, but caged or
-chained creatures cannot be killed with this tool.
+Kills any individual unit, or all undead, or all units of a given race. Caged
+and chained creatures are ignored.
 
 Usage
 -----
 
 ::
 
-    exterminate
+    exterminate [list]
     exterminate this [<options>]
     exterminate undead [<options>]
     exterminate all[:<caste>] [<options>]
@@ -25,10 +24,10 @@ Race and caste names are case insensitive.
 Examples
 --------
 
-``exterminate this``
-    Kill the selected unit.
 ``exterminate``
     List the targets on your map.
+``exterminate this``
+    Kill the selected unit.
 ``exterminate BIRD_RAVEN:MALE``
     Kill the ravens flying around the map (but only the male ones).
 ``exterminate goblin --method magma --only-visible``
@@ -65,18 +64,21 @@ Methods
 :drown: Drown the unit in water.
 :magma: Boil the unit in magma (not recommended for magma-safe creatures).
 :butcher: Will mark the units for butchering instead of killing them. This is
-    more useful for pets than armed enemies.
-:knockout: Will put units into an unconscious state for 30k ticks (about a month).
-:traumatize: Traumatizes all units, forcing them to stare off into space (catatonic state).
+    useful for pets and not useful for armed enemies.
+:knockout: Will put units into an unconscious state for 30k ticks (about a
+    month in fort mode).
+:traumatize: Traumatizes units, forcing them to stare off into space (catatonic
+    state).
 
 Technical details
 -----------------
 
-This tool kills by setting a unit's ``blood_count`` to 0, which means
-immediate death at the next game tick. For creatures where this is not enough,
-such as vampires, it also sets ``animal.vanish_countdown``, allowing the unit
-to vanish in a puff of smoke if the blood loss doesn't kill them.
+For the ``instant`` method, this tool kills by setting a unit's ``blood_count``
+to 0, which means immediate death at the next game tick. For creatures where
+this is not enough, such as vampires, it also sets ``animal.vanish_countdown``,
+allowing the unit to vanish in a puff of smoke if the blood loss doesn't kill
+them.
 
 If the method of choice involves liquids, the tile is filled with a liquid
 level of 7 every tick. If the target unit moves, the liquid moves along with
-it, leaving the vacated tiles clean.
+it, leaving the vacated tiles clean (though possibly scorched).
diff --git a/emigration.lua b/emigration.lua
@@ -153,16 +153,13 @@ function canLeave(unit)
         return false
     end
 
-    return dfhack.units.isCitizen(unit) and
-           dfhack.units.isActive(unit) and
-           not dfhack.units.isOpposedToLife(unit) and
-           not unit.flags1.merchant and
-           not unit.flags1.diplomat and
-           not unit.flags1.chained and
-           dfhack.units.getNoblePositions(unit) == nil and
-           unit.military.squad_id == -1 and
-           not dfhack.units.isBaby(unit) and
-           not dfhack.units.isChild(unit)
+    return dfhack.units.isActive(unit) and
+        dfhack.units.isCitizen(unit) and
+        not dfhack.units.getNoblePositions(unit) and
+        not unit.flags1.chained and
+        unit.military.squad_id == -1 and
+        not dfhack.units.isBaby(unit) and
+        not dfhack.units.isChild(unit)
 end
 
 function checkForDeserters(method,civ_id)
diff --git a/exterminate.lua b/exterminate.lua
@@ -14,32 +14,24 @@ local function spawnLiquid(position, liquid_level, liquid_type, update_liquids)
   map_block.flags.update_liquid_twice = update_liquids
 end
 
-local function checkUnit(unit)
-    return (unit.body.blood_count ~= 0 or unit.body.blood_max == 0) and
-        not unit.flags1.inactive and
-        not unit.flags1.caged and
-        not unit.flags1.chained
-end
-
-local function isUnitFriendly(unit)
-    if dfhack.units.isDanger(unit) then
+local function checkUnit(opts, unit)
+    if not dfhack.units.isActive(unit) or
+        (unit.body.blood_max ~= 0 and unit.body.blood_count == 0) or
+        unit.flags1.caged or
+        unit.flags1.chained
+    then
         return false
     end
-    local adv = dfhack.world.getAdventurer()
-    if adv then
-        if adv == unit or
-            unit.relationship_ids.GroupLeader == adv.id or
-            unit.relationship_ids.PetOwner == adv.id
-        then
-            return true
-        end
+    if opts.only_visible and not dfhack.units.isVisible(unit) then
+        return false
     end
-
-    return dfhack.units.isOwnCiv(unit) or
-        dfhack.units.isOwnGroup(unit) or
-        dfhack.units.isVisiting(unit) or
-        dfhack.units.isTame(unit) or
-        dfhack.units.isDomesticated(unit)
+    if not opts.include_friendly and not dfhack.units.isDanger(unit) and not dfhack.units.isWildlife(unit) then
+        return false
+    end
+    if opts.selected_caste and opts.selected_caste ~= df.creature_raw.find(unit.race).caste[unit.caste].caste_id then
+        return false
+    end
+    return true
 end
 
 killMethod = {
@@ -138,26 +130,17 @@ local function getRaceCastes(race_id)
     return unit_castes
 end
 
-local function getMapRaces(only_visible, include_friendly)
+local function getMapRaces(opts)
     local map_races = {}
     for _, unit in pairs(df.global.world.units.active) do
-        if only_visible and not dfhack.units.isVisible(unit) then
-            goto skipunit
-        end
-        if not include_friendly and isUnitFriendly(unit) then
-            goto skipunit
-        end
-        if dfhack.units.isActive(unit) and checkUnit(unit) then
-            local unit_race_name = dfhack.units.isUndead(unit) and "UNDEAD" or df.creature_raw.find(unit.race).creature_id
-
-            local race = ensure_key(map_races, unit_race_name)
-            race.id = unit.race
-            race.name = unit_race_name
-            race.count = (race.count or 0) + 1
-        end
-        :: skipunit ::
+        if not checkUnit(opts, unit) then goto continue end
+        local unit_race_name = dfhack.units.isUndead(unit) and "UNDEAD" or df.creature_raw.find(unit.race).creature_id
+        local race = ensure_key(map_races, unit_race_name)
+        race.id = unit.race
+        race.name = unit_race_name
+        race.count = (race.count or 0) + 1
+        ::continue::
     end
-
     return map_races
 end
 
@@ -200,9 +183,9 @@ if positionals[1] == "this" then
     return
 end
 
-local map_races = getMapRaces(options.only_visible, options.include_friendly)
+local map_races = getMapRaces(options)
 
-if not positionals[1] then
+if not positionals[1] or positionals[1] == 'list' then
     local sorted_races = {}
     for race, value in pairs(map_races) do
         table.insert(sorted_races, { name = race, count = value.count })
@@ -224,28 +207,19 @@ if race_name:lower() == 'undead' then
         qerror("No undead found on the map.")
     end
     for _, unit in pairs(df.global.world.units.active) do
-        if dfhack.units.isUndead(unit) and checkUnit(unit) then
+        if dfhack.units.isUndead(unit) and checkUnit(options, unit) then
             killUnit(unit, options.method)
             count = count + 1
         end
     end
 elseif positionals[1]:split(':')[1] == "all" then
-    local selected_caste = positionals[1]:split(':')[2]
+    options.selected_caste = positionals[1]:split(':')[2]
 
     for _, unit in ipairs(df.global.world.units.active) do
         if options.limit > 0 and count >= options.limit then
             break
         end
-        if not checkUnit(unit) then
-            goto skipunit
-        end
-        if options.only_visible and not dfhack.units.isVisible(unit) then
-            goto skipunit
-        end
-        if not options.include_friendly and isUnitFriendly(unit) then
-            goto skipunit
-        end
-        if selected_caste and selected_caste ~= df.creature_raw.find(unit.race).caste[unit.caste].caste_id then
+        if not checkUnit(options, unit) then
             goto skipunit
         end
 
@@ -285,21 +259,13 @@ else
     end
 
     target = selected_race
+    options.selected_caste = selected_caste
 
     for _, unit in pairs(df.global.world.units.active) do
         if options.limit > 0 and count >= options.limit then
             break
         end
-        if not checkUnit(unit) then
-            goto skipunit
-        end
-        if options.only_visible and not dfhack.units.isVisible(unit) then
-            goto skipunit
-        end
-        if not options.include_friendly and isUnitFriendly(unit) then
-            goto skipunit
-        end
-        if selected_caste and selected_caste ~= df.creature_raw.find(unit.race).caste[unit.caste].caste_id then
+        if not checkUnit(options, unit) then
             goto skipunit
         end
 
diff --git a/gui/sitemap.lua b/gui/sitemap.lua
@@ -99,6 +99,27 @@ local function zoom_to_next_zone(_, choice)
     data.next_idx = data.next_idx % #data.zones + 1
 end
 
+local function get_unit_disposition_and_pen(unit)
+    local prefix = unit.flags1.caged and 'caged ' or ''
+    if dfhack.units.isDanger(unit) then
+        return prefix..'hostile', COLOR_LIGHTRED
+    end
+    if not dfhack.units.isFortControlled(unit) and dfhack.units.isWildlife(unit) then
+        return prefix..'wildlife', COLOR_GREEN
+    end
+    return prefix..'friendly', COLOR_LIGHTGREEN
+end
+
+local function get_unit_choice_text(unit)
+    local disposition, disposition_pen = get_unit_disposition_and_pen(unit)
+    return {
+        dfhack.units.getReadableName(unit),
+        ' (',
+        {text=disposition, pen=disposition_pen},
+        ')',
+    }
+end
+
 local function get_unit_choices()
     local is_fort = dfhack.world.isFortressMode()
     local choices = {}
@@ -110,7 +131,7 @@ local function get_unit_choices()
             goto continue
         end
         table.insert(choices, {
-            text=dfhack.units.getReadableName(unit),
+            text=get_unit_choice_text(unit),
             data={
                 unit_id=unit.id,
             },
diff --git a/internal/notify/notifications.lua b/internal/notify/notifications.lua
