From ee053b0aabb03497275598915c1a27aa9eb6f134 Mon Sep 17 00:00:00 2001
From: BRADSEC <7948876+bradsec@users.noreply.github.com>
Date: Sun, 21 Jun 2026 19:51:39 +1000
Subject: [PATCH] refactor: add keyed category collection and person_dict
 helper

---
 nodes.py              | 28 ++++++++++++++++++++++------
 test_random_person.py | 19 +++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/nodes.py b/nodes.py
index c9c215b..e97adb0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -436,12 +436,28 @@ def generate_person(seed, randomize, sex, category_args,
             dd("shoulders"), dd("chest"), out_bust, seed)
 
 
-def _collect_category_args(kwargs):
-    """Pull the (mode, allow_list, fixed) triple for each category from kwargs."""
-    return {
-        key: (kwargs[f"{key}_mode"], kwargs[f"{key}_allow_list"], kwargs[f"{key}_fixed"])
-        for key, _, _ in CATEGORY_SPECS
-    }
+def _collect_category_args(kwargs, keys=None):
+    """Pull the (mode, allow_list, fixed) triple for each category from kwargs.
+
+    When `keys` is given, only those categories are read from kwargs; every
+    other category is forced "off". This lets a segment node expose a subset of
+    widgets while the shared core still receives a full category map.
+    """
+    active = set(keys) if keys is not None else None
+    out = {}
+    for key, _, _ in CATEGORY_SPECS:
+        if (active is None or key in active) and f"{key}_mode" in kwargs:
+            out[key] = (kwargs[f"{key}_mode"],
+                        kwargs[f"{key}_allow_list"],
+                        kwargs[f"{key}_fixed"])
+        else:
+            out[key] = ("off", "", "(none)")
+    return out
+
+
+def person_dict(result):
+    """Map a generate_person result tuple to {return_name: value}."""
+    return dict(zip(RETURN_NAMES, result))
 
 
 # -- V1 node (legacy dict API) -------------------------------------------------
diff --git a/test_random_person.py b/test_random_person.py
index 1305988..b686e4c 100644
--- a/test_random_person.py
+++ b/test_random_person.py
@@ -169,6 +169,25 @@ class TestBuildDescription(unittest.TestCase):
         self.assertIn("year old female", s)
 
 
+class TestNodeHelpers(unittest.TestCase):
+
+    def test_collect_args_restricts_keys(self):
+        kwargs = {}
+        for key, _, _ in core.CATEGORY_SPECS:
+            kwargs[f"{key}_mode"] = "random"
+            kwargs[f"{key}_allow_list"] = ""
+            kwargs[f"{key}_fixed"] = "(none)"
+        args = core._collect_category_args(kwargs, keys=["eyes"])
+        self.assertEqual(args["eyes"][0], "random")
+        self.assertEqual(args["hair_color"], ("off", "", "(none)"))
+
+    def test_person_dict_maps_names(self):
+        result = run(seed=5, sex="female")
+        dmap = core.person_dict(result)
+        self.assertEqual(set(dmap.keys()), set(core.RETURN_NAMES))
+        self.assertEqual(dmap["seed"], result[-1])
+
+
 class TestBodyData(unittest.TestCase):
 
     BODY_KEYS = ("shoulders", "chest", "bust_size", "bust_shape")