From ad31341a4d12096a968744cf07594f58f2fcb6e6 Mon Sep 17 00:00:00 2001 From: aladac Date: Mon, 18 May 2026 02:19:01 +0200 Subject: [PATCH] feat(scenes): add scene fragment library mirroring characters Extracts the per-name YAML-list storage logic out of tensors/characters.py into a new generic tensors.fragments.FragmentLibrary class, then layers two kind-specific function-style modules on top: - tensors.characters (unchanged public API, now thin wrapper) - tensors.scenes (new): save_scene / load_scene / list_scenes / delete_scene / scene_path / resolve_scene, stored in ~/.local/share/tensors/scenes/ CLI additions: - tsr scene save -o 'elem1, elem2, ...' - tsr scene list / show / delete Wired into generate + template + style-sweep: - tsr generate -S / --scene-prompt 'elems' injects scene elements into the positive prompt right after the character block (order: quality_prefix -> character -> scene -> rating -> user) - tsr template -S / --scene-prompt 'elems' embeds a 'scene' list field in the dumped JSON template (named + inline merged, deduped) - style-sweep templates accept 'scene' (str|list) and 'scene_prompt' keys - generate --input JSON honors both 'scene' (str=name or list=inline) and 'scene_prompt' keys FragmentLibrary uses kind-aware error messages (e.g. "Invalid scene name" vs "Invalid character name") so users see the right context. 28 new tests cover the generic library + scenes. Bumps version to 0.1.24 (supersedes v0.1.23 which failed to publish due to a transient Sigstore network error). --- tensors/__init__.py | 2 +- tensors/characters.py | 152 ++++++---------------- tensors/cli.py | 280 +++++++++++++++++++++++++++++++++++----- tensors/fragments.py | 178 +++++++++++++++++++++++++ tensors/scenes.py | 75 +++++++++++ tests/test_fragments.py | 91 +++++++++++++ tests/test_scenes.py | 96 ++++++++++++++ 7 files changed, 725 insertions(+), 149 deletions(-) create mode 100644 tensors/fragments.py create mode 100644 tensors/scenes.py create mode 100644 tests/test_fragments.py create mode 100644 tests/test_scenes.py diff --git a/tensors/__init__.py b/tensors/__init__.py index 8ca8bd6..2e82c02 100644 --- a/tensors/__init__.py +++ b/tensors/__init__.py @@ -1,6 +1,6 @@ """tsr: Read safetensor metadata, search and download CivitAI models.""" -__version__ = "0.1.23" +__version__ = "0.1.24" from tensors.cli import main from tensors.config import ( diff --git a/tensors/characters.py b/tensors/characters.py index d2056ea..c277bb0 100644 --- a/tensors/characters.py +++ b/tensors/characters.py @@ -1,126 +1,68 @@ -"""Character library: named lists of prompt elements, stored as YAML. +"""Character library: named lists of prompt elements for the *who*. Characters live in ``~/.local/share/tensors/characters/.yml`` and contain a flat YAML list of strings describing a subject (appearance, identity, outfit fragments). They are injected into a generation prompt via ``--character`` or -``--character-prompt`` on ``tsr generate``. +``--character-prompt`` on ``tsr generate`` and embedded in the ``character`` +field of ``tsr template`` JSON output. -Format is JSON-compatible YAML — each line is ``- "value"`` so the files round-trip -through both ``json`` and any YAML parser. Manual hand-editing with plain or -single-quoted scalars is also accepted on read. +This module exposes a function-style API on top of +:class:`tensors.fragments.FragmentLibrary`. Each call instantiates a library +rooted at the current module-level ``CHARACTERS_DIR`` so tests can monkeypatch +the directory without re-importing. """ from __future__ import annotations -import json -import re -from pathlib import Path # noqa: TC003 # used in runtime return annotations exposed to typer +from pathlib import Path # noqa: TC003 # used in runtime return annotations from tensors.config import DATA_DIR +from tensors.fragments import FragmentLibrary, parse_elements -# Storage location for character YAML files. +__all__ = [ + "CHARACTERS_DIR", + "character_path", + "delete_character", + "list_characters", + "load_character", + "parse_elements", + "resolve_character", + "save_character", +] + +# Default storage location. Tests may monkeypatch this attribute; every helper +# below dereferences it via globals() so the override is picked up live. CHARACTERS_DIR = DATA_DIR / "characters" -# Restrict character names to a safe subset so they can't escape CHARACTERS_DIR -# via path traversal and so file listings stay tidy. -_NAME_RE = re.compile(r"^[A-Za-z0-9_.-]+$") -# Minimum length for a quoted YAML scalar: opening + closing quote. -_MIN_QUOTED_SCALAR_LEN = 2 - - -def _validate_name(name: str) -> None: - if not name or not _NAME_RE.match(name): - raise ValueError(f"Invalid character name {name!r}: only letters, digits, '.', '_', '-' allowed") +def _lib() -> FragmentLibrary: + """Build a fresh library bound to the current ``CHARACTERS_DIR`` value.""" + return FragmentLibrary("characters", base_dir=globals()["CHARACTERS_DIR"]) def character_path(name: str) -> Path: - """Return the on-disk path for a character name (without ensuring it exists).""" - _validate_name(name) - return CHARACTERS_DIR / f"{name}.yml" - - -def parse_elements(value: str) -> list[str]: - """Split a comma-separated prompt fragment into clean, order-preserving elements. - - Empty pieces and duplicates are dropped. Used by ``tsr character save`` and - by the ``--character-prompt`` CLI flag so both share identical splitting - semantics. - """ - parts = [p.strip() for p in value.split(",")] - seen: set[str] = set() - out: list[str] = [] - for p in parts: - if p and p not in seen: - seen.add(p) - out.append(p) - return out + """Return the on-disk path for a character name (without ensuring existence).""" + return _lib().path(name) def save_character(name: str, elements: list[str]) -> Path: - """Persist a character's elements to disk and return the file path. - - Overwrites any existing file. Each element is written on its own line as a - JSON-encoded YAML scalar (``- "value"``), which keeps embedded commas, - quotes, and unicode safe. - """ - CHARACTERS_DIR.mkdir(parents=True, exist_ok=True) - path = character_path(name) - body = "\n".join(f"- {json.dumps(e, ensure_ascii=False)}" for e in elements) - path.write_text(body + "\n" if body else "") - return path + """Persist a character's elements to disk and return the file path.""" + return _lib().save(name, elements) def load_character(name: str) -> list[str]: - """Load a character's elements. Raises ``FileNotFoundError`` if missing. - - Accepts both JSON-quoted scalars (``- "value"``), single-quoted YAML scalars - (``- 'value'``) and bare scalars (``- value``). Blank lines and ``#`` comments - are ignored. - """ - path = character_path(name) - if not path.is_file(): - raise FileNotFoundError(f"Character {name!r} not found at {path}") - - elements: list[str] = [] - for raw in path.read_text().splitlines(): - line = raw.strip() - if not line or line.startswith("#"): - continue - if not line.startswith("-"): - # Skip any non-list-item lines (e.g. a YAML document header users - # might add manually); we only consume flat list entries. - continue - item = line[1:].lstrip() - if not item: - continue - # Prefer strict JSON decode (covers our own writer output and any - # double-quoted YAML scalar). Fall back to single-quoted YAML, then - # bare scalar. - try: - value = json.loads(item) - if not isinstance(value, str): - value = str(value) - except json.JSONDecodeError: - value = item[1:-1].replace("''", "'") if len(item) >= _MIN_QUOTED_SCALAR_LEN and item[0] == item[-1] == "'" else item - elements.append(value) - return elements + """Load a character's elements. Raises ``FileNotFoundError`` if missing.""" + return _lib().load(name) def list_characters() -> list[str]: - """Return sorted character names. Empty list if no characters dir exists yet.""" - if not CHARACTERS_DIR.is_dir(): - return [] - return sorted(p.stem for p in CHARACTERS_DIR.glob("*.yml") if p.is_file()) + """Return sorted character names.""" + return _lib().list() def delete_character(name: str) -> bool: - """Delete a character file. Returns True on success, False if it was missing.""" - path = character_path(name) - if not path.is_file(): - return False - path.unlink() - return True + """Delete a character file. Returns True on success, False if missing.""" + return _lib().delete(name) def resolve_character( @@ -129,25 +71,5 @@ def resolve_character( character_prompt: str | None = None, extra: list[str] | None = None, ) -> list[str]: - """Merge a named character with an inline ``--character-prompt`` and optional extras. - - Resolution order (first match wins per duplicate): named character → - ``--character-prompt`` elements → ``extra`` list. The result preserves order - and drops duplicates, mirroring ``parse_elements``. - """ - merged: list[str] = [] - seen: set[str] = set() - - def _push(items: list[str]) -> None: - for item in items: - if item and item not in seen: - seen.add(item) - merged.append(item) - - if character: - _push(load_character(character)) - if character_prompt: - _push(parse_elements(character_prompt)) - if extra: - _push(extra) - return merged + """Merge a named character with an inline ``--character-prompt`` and extras.""" + return _lib().resolve(name=character, inline=character_prompt, extra=extra) diff --git a/tensors/cli.py b/tensors/cli.py index 7baab91..f1e63b7 100644 --- a/tensors/cli.py +++ b/tensors/cli.py @@ -788,6 +788,14 @@ def generate( # noqa: PLR0915 str | None, typer.Option("--character-prompt", help='Inline character fragment, comma-separated (e.g. "blond hair, blue eyes")'), ] = None, + scene: Annotated[ + str | None, + typer.Option("-S", "--scene", help="Saved scene name (loaded from ~/.local/share/tensors/scenes/)"), + ] = None, + scene_prompt: Annotated[ + str | None, + typer.Option("--scene-prompt", help='Inline scene fragment, comma-separated (e.g. "luxury penthouse, volumetric lighting")'), + ] = None, family: Annotated[ str | None, typer.Option( @@ -912,7 +920,20 @@ def generate( # noqa: PLR0915 character_prompt = ", ".join(str(x) for x in val if str(x).strip()) if "character_prompt" in mapped and "character_prompt" not in explicit: cp_val = mapped["character_prompt"] - character_prompt = cp_val if isinstance(cp_val, str) else ", ".join(str(x) for x in cp_val if str(x).strip()) + character_prompt = ( + cp_val if isinstance(cp_val, str) else ", ".join(str(x) for x in cp_val if str(x).strip()) + ) + if "scene" in mapped and "scene" not in explicit: + sv = mapped["scene"] + if isinstance(sv, str): + scene = sv + elif isinstance(sv, (list, tuple)): + scene_prompt = ", ".join(str(x) for x in sv if str(x).strip()) + if "scene_prompt" in mapped and "scene_prompt" not in explicit: + sp_val = mapped["scene_prompt"] + scene_prompt = ( + sp_val if isinstance(sp_val, str) else ", ".join(str(x) for x in sp_val if str(x).strip()) + ) if "rating" in mapped and "rating" not in explicit: rating = mapped["rating"] @@ -942,6 +963,8 @@ def generate( # noqa: PLR0915 no_negative=no_negative, character=character, character_prompt=character_prompt, + scene=scene, + scene_prompt=scene_prompt, family=family, output=output, remote=remote, @@ -1032,6 +1055,8 @@ def _run_generation( # noqa: PLR0915 no_negative: bool = False, character: str | None = None, character_prompt: str | None = None, + scene: str | None = None, + scene_prompt: str | None = None, family: str | None = None, output: Path | None = None, remote: str | None = None, @@ -1105,6 +1130,31 @@ def _run_generation( # noqa: PLR0915 f"[dim]Character ({origin}, {len(character_elements)} elements): {', '.join(character_elements)}[/dim]" ) + # Resolve scene (named lookup + inline --scene-prompt, merged + deduped). + # Scene sits between character (who) and rating/user prompt (what's happening) + # so the natural reading order is: quality → character → scene → rating → user. + scene_elements: list[str] = [] + if scene or scene_prompt: + from tensors.scenes import resolve_scene # noqa: PLC0415 + + try: + scene_elements = resolve_scene(scene=scene, scene_prompt=scene_prompt) + except FileNotFoundError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + except ValueError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + + if scene_elements: + prompt_parts.extend(scene_elements) + if not json_output: + origin = f"'{scene}'" if scene else "inline" + console.print( + f"[dim]Scene ({origin}, {len(scene_elements)} elements): " + f"{', '.join(scene_elements)}[/dim]" + ) + # Add rating tag based on model family (Pony/Illustrious) if rating: from tensors.config import get_rating_tag # noqa: PLC0415 @@ -1350,6 +1400,8 @@ _STYLE_SWEEP_TEMPLATE_KEYS = { "no_negative", "character", "character_prompt", + "scene", + "scene_prompt", "rating", "family", "remote", @@ -1680,22 +1732,27 @@ def style_sweep( # noqa: PLR0915 pending_tasks.append((i, entry, result, out_path)) - # Character resolution: templates may carry either a name string (look up - # at run-time) or an inline list of resolved elements (e.g. produced by - # `tsr template -C ...`). Lists are joined into `character_prompt` so - # _run_generation sees a uniform CSV string and skips the disk lookup. - char_val = tpl_data.get("character") - char_prompt_val = tpl_data.get("character_prompt") - char_name: str | None = None - char_inline: str | None = None - if isinstance(char_val, str): - char_name = char_val - elif isinstance(char_val, (list, tuple)): - char_inline = ", ".join(str(x) for x in char_val if str(x).strip()) - if char_prompt_val is not None: - char_inline = ( - char_prompt_val if isinstance(char_prompt_val, str) else ", ".join(str(x) for x in char_prompt_val if str(x).strip()) - ) + # Character / scene resolution: templates may carry either a name string + # (look up at run-time) or an inline list of resolved elements (e.g. produced + # by `tsr template -C ... -S ...`). Lists are joined into the *_prompt arg + # so _run_generation sees a uniform CSV string and skips the disk lookup. + def _split_fragment(name_val: Any, prompt_val: Any) -> tuple[str | None, str | None]: + name_out: str | None = None + inline_out: str | None = None + if isinstance(name_val, str): + name_out = name_val + elif isinstance(name_val, (list, tuple)): + inline_out = ", ".join(str(x) for x in name_val if str(x).strip()) + if prompt_val is not None: + inline_out = ( + prompt_val + if isinstance(prompt_val, str) + else ", ".join(str(x) for x in prompt_val if str(x).strip()) + ) + return name_out, inline_out + + char_name, char_inline = _split_fragment(tpl_data.get("character"), tpl_data.get("character_prompt")) + scene_name, scene_inline = _split_fragment(tpl_data.get("scene"), tpl_data.get("scene_prompt")) # Common kwargs for every _run_generation call — extracted from the # template once, reused across sequential and parallel paths. @@ -1720,6 +1777,8 @@ def style_sweep( # noqa: PLR0915 "no_negative": bool(_t("no_negative", default=False)), "character": char_name, "character_prompt": char_inline, + "scene": scene_name, + "scene_prompt": scene_inline, "family": _t("family"), "remote": gen_remote, "json_output": False, @@ -1883,6 +1942,17 @@ def template( help="Inline character fragment, comma-separated (merged with --character into `character`)", ), ] = None, + scene: Annotated[ + str | None, + typer.Option("-S", "--scene", help="Saved scene name (resolved into the `scene` list field)"), + ] = None, + scene_prompt: Annotated[ + str | None, + typer.Option( + "--scene-prompt", + help='Inline scene fragment, comma-separated (merged with --scene into `scene`)', + ), + ] = None, output: Annotated[Path | None, typer.Option("-o", "--output", help="Save template to file")] = None, ) -> None: """Dump a JSON generation template with resolved defaults for a model. @@ -1890,17 +1960,18 @@ def template( Outputs a ready-to-use JSON object with all parameters auto-resolved from the checkpoint family. Pipe to 'tsr generate --input' or save to a file for reuse. - ``--character`` and ``--character-prompt`` append a ``character`` list to the - template (saved-name elements first, inline elements appended, deduped). + ``--character`` / ``--character-prompt`` append a ``character`` list to the + template; ``--scene`` / ``--scene-prompt`` append a ``scene`` list (named + elements first, inline elements appended, deduped within each list). Examples: tsr template -m ponyDiffusionV6XL_v6StartWithThisOne.safetensors tsr template -m beautifulRealistic_v7.safetensors -O portrait tsr template -m waiIllustriousSDXL_v160.safetensors -l "Elvira iIlluLoRA.safetensors" tsr template -m ponyRealism_V22.safetensors -o pony_preset.json - tsr template -m flux1-dev.safetensors -C cassie_cage # embeds saved character - tsr template -m flux1-dev.safetensors --character-prompt "blond hair, blue eyes" - tsr template -m flux1-dev.safetensors -C cassie --character-prompt "wet skin" + tsr template -m flux1-dev.safetensors -C cassie_cage -S penthouse + tsr template -m flux1-dev.safetensors --character-prompt "blond hair, blue eyes" \\ + --scene-prompt "luxury penthouse, volumetric lighting, Canon R5" tsr generate --input "$(tsr template -m ponyRealism_V22.safetensors)" "a portrait" """ from tensors.config import get_model_generation_defaults, resolve_orientation # noqa: PLC0415 @@ -1954,11 +2025,10 @@ def template( tpl["lora"] = lora tpl["lora_strength"] = lora_strength - # Resolve character into a flat list embedded in the template. When the - # template is later fed to `tsr generate --input`, _run_generation will - # treat a list under `character` as inline elements (no re-lookup needed). - # --character and --character-prompt merge in that order (named first, - # inline appended, duplicates dropped). + # Resolve character / scene into flat lists embedded in the template. When + # the template is later fed to `tsr generate --input`, _run_generation will + # treat the lists under `character` / `scene` as inline elements (no + # re-lookup needed). The `_*_name` fields are informational only. if character or character_prompt: from tensors.characters import resolve_character # noqa: PLC0415 @@ -1976,6 +2046,23 @@ def template( if character: tpl["_character_name"] = character + if scene or scene_prompt: + from tensors.scenes import resolve_scene # noqa: PLC0415 + + try: + resolved_scene = resolve_scene(scene=scene, scene_prompt=scene_prompt) + except FileNotFoundError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + except ValueError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + + if resolved_scene: + tpl["scene"] = resolved_scene + if scene: + tpl["_scene_name"] = scene + # Add metadata (not used by generate, but informational) tpl["_family"] = family or "unknown" if base_model_str: @@ -2460,12 +2547,19 @@ def hf_download( # ============================================================================= -# Character Commands +# Fragment Commands (character + scene) # ============================================================================= -# Characters are named, comma-split prompt fragments stored as YAML lists in -# ~/.local/share/tensors/characters/.yml. They are injected into the -# positive prompt by `tsr generate --character ` (or inline via -# `--character-prompt "elem1, elem2"`). +# Characters and scenes are named, comma-split prompt fragments stored as YAML +# lists in ~/.local/share/tensors//.yml. They are injected into the +# positive prompt by `tsr generate --character ` / `--scene ` (or +# inline via `--character-prompt "..."` / `--scene-prompt`). +# +# Both subcommand groups share the underlying tensors.fragments.FragmentLibrary, +# but the CLI commands are spelled out per-kind to keep Typer's signature +# introspection happy under `from __future__ import annotations` (closures +# referencing per-kind labels break typer's eval_str=True resolution). + +# ---- character ---- character_app = typer.Typer( name="character", @@ -2501,7 +2595,7 @@ def character_save( raise typer.Exit(1) from e if json_output: - console.print_json(data={"name": name, "path": str(path), "elements": parsed}) + console.print_json(data={"name": name, "path": str(path), "elements": parsed, "kind": "characters"}) return console.print(f"[green]Saved character '{name}' ({len(parsed)} elements):[/green] {path}") @@ -2582,6 +2676,126 @@ def character_delete( raise typer.Exit(1) +# ---- scene ---- + +scene_app = typer.Typer( + name="scene", + help="Manage saved scene prompts (~/.local/share/tensors/scenes/).", + no_args_is_help=True, +) +app.add_typer(scene_app) + + +@scene_app.command("save") +def scene_save( + elements: Annotated[ + str, typer.Argument(help='Comma-separated prompt elements (e.g. "luxury penthouse, volumetric lighting")') + ], + name: Annotated[str, typer.Option("-o", "--output", help="Scene name (used as filename)")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Save a scene as a YAML list of prompt elements. + + Examples: + tsr scene save -o penthouse "luxury penthouse, volumetric lighting, Canon R5, 85mm" + tsr scene save -o forest "deep forest, dappled sunlight, moss-covered rocks" + """ + from tensors.fragments import parse_elements # noqa: PLC0415 + from tensors.scenes import save_scene # noqa: PLC0415 + + parsed = parse_elements(elements) + if not parsed: + console.print("[red]No usable elements after splitting on commas[/red]") + raise typer.Exit(1) + + try: + path = save_scene(name, parsed) + except ValueError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + + if json_output: + console.print_json(data={"name": name, "path": str(path), "elements": parsed, "kind": "scenes"}) + return + + console.print(f"[green]Saved scene '{name}' ({len(parsed)} elements):[/green] {path}") + for elem in parsed: + console.print(f" • {elem}") + + +@scene_app.command("list") +def scene_list( + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """List saved scenes.""" + from tensors.scenes import SCENES_DIR, list_scenes # noqa: PLC0415 + + names = list_scenes() + if json_output: + console.print_json(data={"dir": str(SCENES_DIR), "scenes": names}) + return + + if not names: + console.print(f"[yellow]No scenes saved in {SCENES_DIR}.[/yellow]") + console.print('[dim]Create one with: tsr scene save -o "elem1, elem2"[/dim]') + return + + console.print(f"[bold]Scenes[/bold] ({len(names)}) [dim]in {SCENES_DIR}[/dim]") + for n in names: + console.print(f" • {n}") + + +@scene_app.command("show") +def scene_show( + name: Annotated[str, typer.Argument(help="Scene name")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Show a scene's elements.""" + from tensors.scenes import load_scene, scene_path # noqa: PLC0415 + + try: + elements = load_scene(name) + except FileNotFoundError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + except ValueError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + + if json_output: + console.print_json(data={"name": name, "path": str(scene_path(name)), "elements": elements}) + return + + console.print(f"[bold]{name}[/bold] [dim]({scene_path(name)})[/dim]") + for elem in elements: + console.print(f" • {elem}") + + +@scene_app.command("delete") +def scene_delete( + name: Annotated[str, typer.Argument(help="Scene name")], + json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False, +) -> None: + """Delete a saved scene.""" + from tensors.scenes import delete_scene # noqa: PLC0415 + + try: + deleted = delete_scene(name) + except ValueError as e: + console.print(f"[red]{e}[/red]") + raise typer.Exit(1) from e + + if json_output: + console.print_json(data={"name": name, "deleted": deleted}) + return + + if deleted: + console.print(f"[green]Deleted scene '{name}'[/green]") + else: + console.print(f"[yellow]Scene '{name}' does not exist[/yellow]") + raise typer.Exit(1) + + # ============================================================================= # ComfyUI Commands # ============================================================================= diff --git a/tensors/fragments.py b/tensors/fragments.py new file mode 100644 index 0000000..4ce11e0 --- /dev/null +++ b/tensors/fragments.py @@ -0,0 +1,178 @@ +"""Generic prompt-fragment library. + +A *fragment* is a named, ordered list of comma-style prompt elements +(e.g. ``["blond hair", "broad chin"]``) stored as a flat YAML list on disk. +Different *kinds* of fragments (characters, scenes, …) each get their own +subdirectory under ``~/.local/share/tensors//`` and their own +``FragmentLibrary`` instance. + +Files are written as JSON-encoded YAML scalars (``- "value"``) so they round-trip +through both ``json`` and any YAML parser. Hand-edited single-quoted or bare +scalars are accepted on read. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path # noqa: TC003 # used in runtime return annotations exposed to typer + +from tensors.config import DATA_DIR + +# Restrict fragment names to a safe subset so they can't escape the storage dir +# via path traversal and so file listings stay tidy. +_NAME_RE = re.compile(r"^[A-Za-z0-9_.-]+$") + +# Minimum length for a quoted YAML scalar: opening + closing quote. +_MIN_QUOTED_SCALAR_LEN = 2 + + +class FragmentLibrary: + """A named collection of prompt-fragment YAML files of a single *kind*. + + Each ``FragmentLibrary`` is rooted at ``//`` (overridable for + tests). Instance methods are stateless wrappers around that directory. + """ + + def __init__(self, kind: str, base_dir: Path | None = None) -> None: + """Create a library for ``kind`` (e.g. ``"characters"`` or ``"scenes"``). + + ``base_dir`` defaults to ``DATA_DIR / kind`` and is recomputed lazily so + tests can monkeypatch ``DATA_DIR`` *or* the ``base_dir`` attribute + directly without re-importing. + """ + if not kind or not _NAME_RE.match(kind): + raise ValueError(f"Invalid library kind {kind!r}") + self.kind = kind + self.base_dir = base_dir if base_dir is not None else DATA_DIR / kind + + # ---------- internals ---------- + + @property + def _singular(self) -> str: + """Human-readable singular form of ``kind`` used in error messages.""" + return self.kind[:-1] if self.kind.endswith("s") else self.kind + + def _validate_name(self, name: str) -> None: + if not name or not _NAME_RE.match(name): + raise ValueError( + f"Invalid {self._singular} name {name!r}: only letters, digits, '.', '_', '-' allowed" + ) + + def path(self, name: str) -> Path: + """Return the on-disk path for ``name`` (without ensuring it exists).""" + self._validate_name(name) + return self.base_dir / f"{name}.yml" + + # ---------- CRUD ---------- + + def save(self, name: str, elements: list[str]) -> Path: + """Persist ``elements`` to disk and return the file path. + + Overwrites any existing file. Each element is written on its own line as + a JSON-encoded YAML scalar so embedded commas, quotes and unicode are + safe. + """ + self.base_dir.mkdir(parents=True, exist_ok=True) + path = self.path(name) + body = "\n".join(f"- {json.dumps(e, ensure_ascii=False)}" for e in elements) + path.write_text(body + "\n" if body else "") + return path + + def load(self, name: str) -> list[str]: + """Load a fragment. Raises ``FileNotFoundError`` if missing. + + Accepts JSON-quoted scalars (``- "value"``), single-quoted YAML scalars + (``- 'value'``) and bare scalars (``- value``). Blank lines and ``#`` + comments are ignored. + """ + path = self.path(name) + if not path.is_file(): + raise FileNotFoundError(f"{self._singular.capitalize()} {name!r} not found at {path}") + + elements: list[str] = [] + for raw in path.read_text().splitlines(): + line = raw.strip() + if not line or line.startswith("#"): + continue + if not line.startswith("-"): + # Skip non-list lines (e.g. a YAML document header users might + # add manually); we only consume flat list entries. + continue + item = line[1:].lstrip() + if not item: + continue + try: + value = json.loads(item) + if not isinstance(value, str): + value = str(value) + except json.JSONDecodeError: + value = ( + item[1:-1].replace("''", "'") + if len(item) >= _MIN_QUOTED_SCALAR_LEN and item[0] == item[-1] == "'" + else item + ) + elements.append(value) + return elements + + def list(self) -> list[str]: + """Return sorted fragment names. Empty list if the dir doesn't exist yet.""" + if not self.base_dir.is_dir(): + return [] + return sorted(p.stem for p in self.base_dir.glob("*.yml") if p.is_file()) + + def delete(self, name: str) -> bool: + """Delete a fragment. Returns True on success, False if it was missing.""" + path = self.path(name) + if not path.is_file(): + return False + path.unlink() + return True + + # ---------- helpers ---------- + + def resolve( + self, + *, + name: str | None = None, + inline: str | None = None, + extra: list[str] | None = None, + ) -> list[str]: + """Merge a named fragment with an inline CSV string and optional extras. + + Resolution order (first match wins per duplicate): named → inline → extra. + Result preserves order and drops duplicates and empty pieces. + """ + merged: list[str] = [] + seen: set[str] = set() + + def _push(items: list[str]) -> None: + for item in items: + if item and item not in seen: + seen.add(item) + merged.append(item) + + if name: + _push(self.load(name)) + if inline: + _push(parse_elements(inline)) + if extra: + _push(extra) + return merged + + +def parse_elements(value: str) -> list[str]: + """Split a comma-separated prompt fragment into clean, order-preserving elements. + + Empty pieces and duplicates are dropped. Shared by ``tsr character|scene save`` + and the ``--character-prompt`` / ``--scene-prompt`` CLI flags so the splitting + semantics stay identical across surfaces. + """ + parts = [p.strip() for p in value.split(",")] + seen: set[str] = set() + out: list[str] = [] + for p in parts: + if p and p not in seen: + seen.add(p) + out.append(p) + return out diff --git a/tensors/scenes.py b/tensors/scenes.py new file mode 100644 index 0000000..f234f7a --- /dev/null +++ b/tensors/scenes.py @@ -0,0 +1,75 @@ +"""Scene library: named lists of prompt elements for the *where*. + +Scenes live in ``~/.local/share/tensors/scenes/.yml`` and contain a flat +YAML list of strings describing a setting (location, lighting, camera, etc.). +They are injected into a generation prompt via ``--scene`` or ``--scene-prompt`` +on ``tsr generate`` and embedded in the ``scene`` field of ``tsr template`` +JSON output. + +This module exposes a function-style API on top of +:class:`tensors.fragments.FragmentLibrary`. Each call instantiates a library +rooted at the current module-level ``SCENES_DIR`` so tests can monkeypatch the +directory without re-importing. +""" + +from __future__ import annotations + +from pathlib import Path # noqa: TC003 # used in runtime return annotations + +from tensors.config import DATA_DIR +from tensors.fragments import FragmentLibrary, parse_elements + +__all__ = [ + "SCENES_DIR", + "delete_scene", + "list_scenes", + "load_scene", + "parse_elements", + "resolve_scene", + "save_scene", + "scene_path", +] + +# Default storage location. Tests may monkeypatch this attribute; every helper +# below dereferences it via globals() so the override is picked up live. +SCENES_DIR = DATA_DIR / "scenes" + + +def _lib() -> FragmentLibrary: + """Build a fresh library bound to the current ``SCENES_DIR`` value.""" + return FragmentLibrary("scenes", base_dir=globals()["SCENES_DIR"]) + + +def scene_path(name: str) -> Path: + """Return the on-disk path for a scene name (without ensuring existence).""" + return _lib().path(name) + + +def save_scene(name: str, elements: list[str]) -> Path: + """Persist a scene's elements to disk and return the file path.""" + return _lib().save(name, elements) + + +def load_scene(name: str) -> list[str]: + """Load a scene's elements. Raises ``FileNotFoundError`` if missing.""" + return _lib().load(name) + + +def list_scenes() -> list[str]: + """Return sorted scene names.""" + return _lib().list() + + +def delete_scene(name: str) -> bool: + """Delete a scene file. Returns True on success, False if missing.""" + return _lib().delete(name) + + +def resolve_scene( + *, + scene: str | None = None, + scene_prompt: str | None = None, + extra: list[str] | None = None, +) -> list[str]: + """Merge a named scene with an inline ``--scene-prompt`` and extras.""" + return _lib().resolve(name=scene, inline=scene_prompt, extra=extra) diff --git a/tests/test_fragments.py b/tests/test_fragments.py new file mode 100644 index 0000000..57c0afd --- /dev/null +++ b/tests/test_fragments.py @@ -0,0 +1,91 @@ +"""Tests for the generic FragmentLibrary in tensors.fragments. + +The character/scene-specific modules are thin wrappers over this class; their +behavioral coverage lives in test_characters.py and test_scenes.py. Here we +focus on cross-kind invariants (kind-aware error messages, name validation, +parse helper). +""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture +def lib_factory(tmp_path): + """Return a callable that builds an isolated FragmentLibrary per kind.""" + from tensors.fragments import FragmentLibrary + + def _make(kind: str = "characters"): + return FragmentLibrary(kind, base_dir=tmp_path / kind) + + return _make + + +class TestParseElements: + def test_splits_trims_and_dedupes(self): + from tensors.fragments import parse_elements + + assert parse_elements("a, b , a, c, b") == ["a", "b", "c"] + + def test_empty_returns_empty(self): + from tensors.fragments import parse_elements + + assert parse_elements("") == [] + assert parse_elements(" , , ") == [] + + +class TestKindAwareMessages: + def test_singular_in_error_for_plural_kind(self, lib_factory): + lib = lib_factory("scenes") + with pytest.raises(ValueError, match=r"Invalid scene name"): + lib.save("bad name", ["x"]) + + def test_singular_in_error_for_already_singular_kind(self, lib_factory): + lib = lib_factory("style") # already singular — should NOT strip the 'e' + with pytest.raises(ValueError, match=r"Invalid style name"): + lib.save("bad name", ["x"]) + + def test_load_missing_uses_kind_singular(self, lib_factory): + lib = lib_factory("scenes") + with pytest.raises(FileNotFoundError, match=r"Scene 'ghost' not found"): + lib.load("ghost") + + +class TestLibraryConstructor: + def test_rejects_invalid_kind(self): + from tensors.fragments import FragmentLibrary + + with pytest.raises(ValueError, match=r"Invalid library kind"): + FragmentLibrary("bad kind") + + def test_accepts_dotted_alphanumeric_kind(self): + from tensors.fragments import FragmentLibrary + + lib = FragmentLibrary("custom_v2.beta") + assert lib.kind == "custom_v2.beta" + + +class TestResolve: + def test_only_named(self, lib_factory): + lib = lib_factory() + lib.save("base", ["x", "y"]) + assert lib.resolve(name="base") == ["x", "y"] + + def test_only_inline(self, lib_factory): + lib = lib_factory() + assert lib.resolve(inline="a, b, c") == ["a", "b", "c"] + + def test_named_inline_extras_merge_deduped_in_order(self, lib_factory): + lib = lib_factory() + lib.save("base", ["a", "b"]) + assert lib.resolve(name="base", inline="b, c", extra=["c", "d"]) == ["a", "b", "c", "d"] + + def test_no_args_empty(self, lib_factory): + lib = lib_factory() + assert lib.resolve() == [] + + def test_named_missing_raises(self, lib_factory): + lib = lib_factory() + with pytest.raises(FileNotFoundError): + lib.resolve(name="absent") diff --git a/tests/test_scenes.py b/tests/test_scenes.py new file mode 100644 index 0000000..b90a902 --- /dev/null +++ b/tests/test_scenes.py @@ -0,0 +1,96 @@ +"""Tests for the scene library (tensors.scenes). + +Mirrors test_characters.py — the underlying implementation is shared so the +focus here is verifying that scene-specific exports route through correctly +and that monkeypatching SCENES_DIR works the same way. +""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture +def scene_env(tmp_path, monkeypatch): + """Redirect SCENES_DIR at fresh tmp_path per test.""" + from tensors import scenes as scene_mod + + scene_dir = tmp_path / "scenes" + monkeypatch.setattr(scene_mod, "SCENES_DIR", scene_dir) + return scene_mod, scene_dir + + +class TestSaveLoad: + def test_save_creates_dir_and_file(self, scene_env): + scene_mod, scene_dir = scene_env + path = scene_mod.save_scene("penthouse", ["luxury penthouse", "volumetric lighting"]) + + assert path == scene_dir / "penthouse.yml" + assert path.is_file() + assert path.read_text() == '- "luxury penthouse"\n- "volumetric lighting"\n' + + def test_load_roundtrip(self, scene_env): + scene_mod, _ = scene_env + elements = ["a", "b with spaces", "c, with, commas", 'd "quotes"'] + scene_mod.save_scene("mixed", elements) + + assert scene_mod.load_scene("mixed") == elements + + def test_load_missing_raises(self, scene_env): + scene_mod, _ = scene_env + with pytest.raises(FileNotFoundError, match=r"Scene 'nope' not found"): + scene_mod.load_scene("nope") + + +class TestListAndDelete: + def test_list_empty_when_dir_missing(self, scene_env): + scene_mod, _ = scene_env + assert scene_mod.list_scenes() == [] + + def test_list_sorted_names(self, scene_env): + scene_mod, _ = scene_env + scene_mod.save_scene("zeta", ["x"]) + scene_mod.save_scene("alpha", ["y"]) + + assert scene_mod.list_scenes() == ["alpha", "zeta"] + + def test_delete_existing(self, scene_env): + scene_mod, _ = scene_env + scene_mod.save_scene("doomed", ["x"]) + + assert scene_mod.delete_scene("doomed") is True + assert scene_mod.list_scenes() == [] + + def test_delete_missing(self, scene_env): + scene_mod, _ = scene_env + assert scene_mod.delete_scene("ghost") is False + + +class TestResolveScene: + def test_named_plus_inline_dedup(self, scene_env): + scene_mod, _ = scene_env + scene_mod.save_scene("base", ["a", "b"]) + + assert scene_mod.resolve_scene(scene="base", scene_prompt="b, c") == ["a", "b", "c"] + + def test_extra_appends_last(self, scene_env): + scene_mod, _ = scene_env + + assert scene_mod.resolve_scene(scene_prompt="a, b", extra=["b", "c"]) == ["a", "b", "c"] + + def test_no_args_empty(self, scene_env): + scene_mod, _ = scene_env + assert scene_mod.resolve_scene() == [] + + +class TestNameValidation: + @pytest.mark.parametrize("bad_name", ["", "foo/bar", "../etc", "with space", "a$b"]) + def test_invalid_names_rejected(self, scene_env, bad_name): + scene_mod, _ = scene_env + with pytest.raises(ValueError, match=r"Invalid scene name"): + scene_mod.save_scene(bad_name, ["x"]) + + def test_scene_path_validates(self, scene_env): + scene_mod, _ = scene_env + with pytest.raises(ValueError, match=r"Invalid scene name"): + scene_mod.scene_path("bad name")