Extract category extension policy

category_extensions.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+try:
+    from . import category_library as category_policy
+    from . import generate_prompt_batches as g
+    from . import row_item as row_item_policy
+except ImportError:  # Allows local smoke tests with top-level imports.
+    import category_library as category_policy
+    import generate_prompt_batches as g
+    import row_item as row_item_policy
+
+
+BUILTIN_CATEGORIES = [
+    "auto_weighted",
+    "auto_full",
+    "woman",
+    "man",
+    "couple",
+    "group_or_layout",
+    "custom_random",
+]
+
+_EXTENSIONS_APPLIED = False
+
+
+def list_from(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return value
+    return [value]
+
+
+def unique_extend(target: list[Any], additions: list[Any]) -> None:
+    seen = set()
+    for item in target:
+        try:
+            seen.add(json.dumps(item, sort_keys=True))
+        except TypeError:
+            seen.add(repr(item))
+    for item in additions:
+        try:
+            marker = json.dumps(item, sort_keys=True)
+        except TypeError:
+            marker = repr(item)
+        if marker not in seen:
+            target.append(item)
+            seen.add(marker)
+
+
+def extension_targets() -> dict[str, tuple[list[Any], bool]]:
+    return {
+        "women_clothes": (g.WOMEN_CLOTHES, False),
+        "women_clothes_minimal": (g.WOMEN_CLOTHES_MINIMAL, False),
+        "men_clothes": (g.MEN_CLOTHES, False),
+        "men_clothes_minimal": (g.MEN_CLOTHES_MINIMAL, False),
+        "couple_outfits": (g.COUPLE_OUTFITS, False),
+        "couple_outfits_minimal": (g.COUPLE_OUTFITS_MINIMAL, False),
+        "poses": (g.POSES, False),
+        "evocative_poses": (g.EVOCATIVE_POSES, False),
+        "backside_poses": (g.BACKSIDE_POSES, False),
+        "expressions": (g.EXPRESSIONS, False),
+        "compositions": (g.COMPOSITIONS, False),
+        "props": (g.PROPS, False),
+        "figure_curvy": (g.FIGURE_CURVY, False),
+        "figure_athletic": (g.FIGURE_ATHLETIC, False),
+        "figure_bombshell": (g.FIGURE_BOMBSHELL, False),
+        "scenes": (g.SCENES, True),
+        "group_scenes": (g.GROUP_SCENES, True),
+        "layouts_full": (g.LAYOUTS_FULL, True),
+        "layouts_minimal": (g.LAYOUTS_MINIMAL, True),
+        "group_compositions": (g.GROUP_COMPOSITIONS, False),
+        "group_ages": (g.GROUP_AGES, False),
+    }
+
+
+def apply_pool_extensions() -> None:
+    global _EXTENSIONS_APPLIED
+    if _EXTENSIONS_APPLIED:
+        return
+    targets = extension_targets()
+    for path in category_policy.category_json_files():
+        data = category_policy.read_category_json(path)
+        extensions = data.get("pool_extensions", {})
+        if not isinstance(extensions, dict):
+            raise ValueError(f"pool_extensions in {path} must be an object")
+        for target_name, additions in extensions.items():
+            if target_name not in targets:
+                known = ", ".join(sorted(targets))
+                raise ValueError(f"Unknown pool extension '{target_name}' in {path}. Known: {known}")
+            target, expects_pair = targets[target_name]
+            normalized = (
+                [row_item_policy.pair_from(item) for item in list_from(additions)]
+                if expects_pair
+                else [row_item_policy.item_text(item) for item in list_from(additions)]
+            )
+            unique_extend(target, normalized)
+    g.EVOCATIVE_ALL = g.EVOCATIVE_POSES + g.BACKSIDE_POSES
+    _EXTENSIONS_APPLIED = True
+
+
+def category_choices() -> list[str]:
+    apply_pool_extensions()
+    custom = [category["name"] for category in category_policy.load_category_library()]
+    return BUILTIN_CATEGORIES + [name for name in custom if name not in BUILTIN_CATEGORIES]
+
+
+def subcategory_choices() -> list[str]:
+    apply_pool_extensions()
+    choices = [category_policy.RANDOM_SUBCATEGORY]
+    for category in category_policy.load_category_library():
+        for subcategory in category["subcategories"]:
+            choices.append(f"{category['name']} / {subcategory['name']}")
+    return choices