sort: clickable stress icons in the top bar

docs/changelog.txt

@@ -61,6 +61,7 @@ Template for new versions:
 ## New Features
 - `autodump`: new ``undestroy`` option reverts pending item destruction while the game is still paused
 - `stocks`: the overlay's ``collapse all`` hotkey now toggles, expanding all categories again when everything is collapsed
+- `sort`: the stress/happiness icons in the top bar are now clickable, opening a list of the citizens in that stress category; selecting a citizen centers the map on them
 
 ## Fixes
 - Fix broken weather lookup in ``World::ReadCurrentWeather``

docs/plugins/sort.rst

@@ -185,3 +185,10 @@ World overlay
 
 Searching is supported for the Artifacts list when viewing the world map (where
 you can initiate raids).
+
+Stress icons overlay
+--------------------
+
+The stress/happiness icons in the top bar are clickable. Clicking an icon opens
+a searchable list of the citizens in that stress category, and selecting a
+citizen from the list centers the map on them.

plugins/lua/sort.lua

@@ -1326,6 +1326,7 @@ OVERLAY_WIDGETS = {
     places=require('plugins.sort.places').PlacesOverlay,
     elevate_barony=require('plugins.sort.diplomacy').DiplomacyOverlay,
     elevate_barony_preferences=require('plugins.sort.diplomacy').PreferenceOverlay,
+    stress_icons=require('plugins.sort.stressicons').StressIconsOverlay,
 }
 
 dfhack.onStateChange[GLOBAL_KEY] = function(sc)

plugins/lua/sort/stressicons.lua

@@ -0,0 +1,95 @@
+local _ENV = mkmodule('plugins.sort.stressicons')
+
+local dialogs = require('gui.dialogs')
+local overlay = require('plugins.overlay')
+
+-- indexed by stress category (0-6, where 0 is most stressed)
+local STRESS_CATEGORY_NAMES = {
+    [0]='Miserable', 'Unhappy', 'Displeased', 'Content', 'Pleased', 'Happy',
+    'Ecstatic',
+}
+
+local function show_stress_units(category)
+    local choices = {}
+    for _,unit in ipairs(df.global.world.units.active) do
+        if dfhack.units.isCitizen(unit, true) and
+                dfhack.units.getStressCategory(unit) == category then
+            table.insert(choices, {
+                text=dfhack.units.getReadableName(unit),
+                unit_id=unit.id,
+            })
+        end
+    end
+    local name = STRESS_CATEGORY_NAMES[category]
+    if #choices == 0 then
+        dialogs.showMessage('Citizens: ' .. name,
+            ('No citizens are currently %s.'):format(name:lower()))
+        return
+    end
+    dialogs.ListBox{
+        frame_title='Citizens: ' .. name,
+        with_filter=true,
+        choices=choices,
+        on_select=function(_, choice)
+            -- re-resolve in case the unit died while the list was open
+            local unit = df.unit.find(choice.unit_id)
+            if not unit then return end
+            dfhack.gui.revealInDwarfmodeMap(unit.pos.x, unit.pos.y, unit.pos.z, true, true)
+        end,
+    }:show()
+end
+
+-- invisible hotspot over the stress icons in the top bar; clicking an icon
+-- opens a list of the citizens in that stress category
+StressIconsOverlay = defclass(StressIconsOverlay, overlay.OverlayWidget)
+StressIconsOverlay.ATTRS{
+    desc='Click a stress icon in the top bar to list the citizens in that stress category.',
+    default_pos={x=1, y=1},
+    default_enabled=true,
+    viewscreens={'dwarfmode/Default'},
+    overlay_onupdate_max_freq_seconds=1,
+    frame={w=1, h=1},
+}
+
+function StressIconsOverlay:overlay_onupdate()
+    self.icon_rects = nil
+    local dimx = df.global.gps.dimx
+    -- the 'Pop' label immediately precedes the stress icons in the top bar;
+    -- scan right to left so a fortress name containing 'Pop' doesn't win
+    local pop_x
+    for x = math.min(dimx-3, 200), 0, -1 do
+        if dfhack.screen.readTile(x, 0).ch == 80 and         -- 'P'
+                dfhack.screen.readTile(x+1, 0).ch == 111 and -- 'o'
+                dfhack.screen.readTile(x+2, 0).ch == 112 then -- 'p'
+            pop_x = x
+            break
+        end
+    end
+    if not pop_x then return end
+    -- the icons are drawn in 3-column cells right of the 'Pop' block, with the
+    -- happiest category first
+    local rects = {}
+    for i = 0, 6 do
+        local x1 = pop_x + 4 + i * 3
+        if x1 + 1 >= dimx then break end
+        table.insert(rects, {x1=x1, y1=0, x2=x1+2, y2=2, category=6-i})
+    end
+    self.icon_rects = rects
+    self.frame = {l=pop_x+4, t=0, w=#rects*3, h=3}
+end
+
+function StressIconsOverlay:onInput(keys)
+    if not keys._MOUSE_L or not self.icon_rects then return false end
+    local mx, my = dfhack.screen.getMousePos()
+    if not mx then return false end
+    for _,rect in ipairs(self.icon_rects) do
+        if my >= rect.y1 and my <= rect.y2 and
+                mx >= rect.x1 and mx <= rect.x2 then
+            show_stress_units(rect.category)
+            return true
+        end
+    end
+    return false
+end
+
+return _ENV

test/plugins/sort.lua

@@ -0,0 +1,166 @@
+config.mode = 'fortress'
+config.target = 'sort'
+
+local gui = require('gui')
+local overlay = require('plugins.overlay')
+
+local function get_widget()
+    return overlay.get_state().db['sort.stress_icons'].widget
+end
+
+local function get_citizen_counts()
+    local counts = {}
+    for _,unit in ipairs(df.global.world.units.active) do
+        if dfhack.units.isCitizen(unit, true) then
+            local cat = dfhack.units.getStressCategory(unit)
+            counts[cat] = (counts[cat] or 0) + 1
+        end
+    end
+    return counts
+end
+
+-- rects are stored happiest-first, so category c is at index 7-c
+local function rect_for_category(w, cat)
+    if not w.icon_rects then qerror('icon_rects not populated') end
+    return w.icon_rects[7 - cat]
+end
+
+local function click_at(x, y)
+    local handled
+    mock.patch(dfhack.screen, 'getMousePos', mock.func(x, y), function()
+        handled = get_widget():onInput{_MOUSE_L=true}
+    end)
+    return handled
+end
+
+local function dismiss_top()
+    gui.simulateInput(dfhack.gui.getCurViewscreen(true), 'LEAVESCREEN')
+end
+
+local function synth_rects()
+    local rects = {}
+    for i = 0, 6 do
+        local x1 = 42 + i * 3
+        table.insert(rects, {x1=x1, y1=0, x2=x1+2, y2=2, category=6-i})
+    end
+    return rects
+end
+
+config.wrapper = function(test_fn)
+    overlay.rescan()
+    local w = get_widget()
+    -- give the top bar a few frames to render so the anchor scan finds 'Pop';
+    -- environments where the vanilla screen is never rendered (e.g. headless
+    -- Windows CI) fall back to synthetic rects so the click/list behavior is
+    -- still exercised
+    for _ = 1, 20 do
+        w:overlay_onupdate()
+        if w.icon_rects then break end
+        delay()
+    end
+    w.scan_unavailable = not w.icon_rects
+    if w.scan_unavailable then
+        print('NOTE: top bar not rendered; testing with synthetic icon rects')
+        w.icon_rects = synth_rects()
+    end
+    return dfhack.with_finalize(function()
+        -- dismiss any screens the test left open
+        while (dfhack.gui.getCurFocus(true)[1] or ''):match('^dfhack/') do
+            dismiss_top()
+        end
+    end, test_fn)
+end
+
+function test.icon_rects_detected()
+    local w = get_widget()
+    if w.scan_unavailable then
+        -- the vanilla top bar is never rendered in this environment, so the
+        -- screen scan cannot be verified here (covered by other platforms)
+        return
+    end
+    expect.eq(7, #w.icon_rects)
+    expect.eq(6, w.icon_rects[1].category)
+    expect.eq(0, w.icon_rects[7].category)
+    -- each rect covers a 3-wide, 3-tall cell in the top bar
+    for _,rect in ipairs(w.icon_rects) do
+        expect.eq(0, rect.y1)
+        expect.eq(2, rect.y2)
+        expect.eq(2, rect.x2 - rect.x1)
+    end
+end
+
+function test.click_outside_icons_does_nothing()
+    local w = get_widget()
+    local before = dfhack.gui.getCurFocus(true)[1]
+    expect.false_(click_at(0, 30))
+    expect.eq(before, dfhack.gui.getCurFocus(true)[1])
+end
+
+function test.click_icon_opens_filtered_list()
+    local w = get_widget()
+    local counts = get_citizen_counts()
+    local cat
+    for c = 0, 6 do
+        if (counts[c] or 0) > 0 then cat = c break end
+    end
+    expect.ne(nil, cat, 'no citizens with a stress category')
+
+    local rect = rect_for_category(w, cat)
+    expect.true_(click_at(rect.x1 + 1, 1))
+    expect.eq('dfhack/lua/ListBox', dfhack.gui.getCurFocus(true)[1])
+    dismiss_top()
+end
+
+function test.empty_category_shows_message()
+    local w = get_widget()
+    local counts = get_citizen_counts()
+    local cat
+    for c = 0, 6 do
+        if not counts[c] or counts[c] == 0 then cat = c break end
+    end
+    if not cat then
+        -- all categories populated; nothing to test
+        return
+    end
+    local rect = rect_for_category(w, cat)
+    expect.true_(click_at(rect.x1 + 1, 1))
+    expect.true_((dfhack.gui.getCurFocus(true)[1] or ''):match('^dfhack/lua/') ~= nil)
+    dismiss_top()
+end
+
+function test.selecting_unit_centers_map()
+    local w = get_widget()
+    local counts = get_citizen_counts()
+    local cat
+    for c = 0, 6 do
+        if (counts[c] or 0) > 0 then cat = c break end
+    end
+    expect.ne(nil, cat, 'no citizens with a stress category')
+
+    local rect = rect_for_category(w, cat)
+    expect.true_(click_at(rect.x1 + 1, 1))
+    expect.eq('dfhack/lua/ListBox', dfhack.gui.getCurFocus(true)[1])
+
+    local reveal_spy = mock.observe_func(function() end)
+    mock.patch(dfhack.gui, 'revealInDwarfmodeMap', reveal_spy,
+        function()
+            gui.simulateInput(dfhack.gui.getCurViewscreen(true), 'SELECT')
+        end)
+    expect.eq(1, reveal_spy.call_count)
+    local args = reveal_spy.call_args[1]
+    local called_with = {x=args[1], y=args[2], z=args[3]}
+
+    expect.ne(nil, called_with)
+    -- the revealed position must be a real citizen's position in that bucket
+    local matched = false
+    for _,unit in ipairs(df.global.world.units.active) do
+        if dfhack.units.isCitizen(unit, true) and
+                dfhack.units.getStressCategory(unit) == cat and
+                unit.pos.x == called_with.x and unit.pos.y == called_with.y and
+                unit.pos.z == called_with.z then
+            matched = true
+            break
+        end
+    end
+    expect.true_(matched)
+end