From 0ec7eaf6060a14ef6d0bdf57b40e03371a281d4c Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sun, 26 Jul 2026 18:33:31 +0800 Subject: [PATCH] fix(wildcards): resolve weighted N::value syntax inside wildcard YAML lists (#1039) --- py/services/wildcard_service.py | 39 +++++++- tests/services/test_wildcard_service.py | 119 ++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 3 deletions(-) diff --git a/py/services/wildcard_service.py b/py/services/wildcard_service.py index b01eef15..8016febe 100644 --- a/py/services/wildcard_service.py +++ b/py/services/wildcard_service.py @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) _WILDCARD_PATTERN = re.compile(r"__([\w\s.\-+/*\\]+?)__") _OPTION_PATTERN = re.compile(r"{([^{}]*?)}") _TRIGGER_WORD_PATTERN = re.compile(r"^trigger_words\d+$") -_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*([0-9.]+)::") +_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*-?\d+(\.\d+)?::") _NUMERIC_PATTERN = re.compile(r"^-?\d+(\.\d+)?$") @@ -390,7 +390,7 @@ class WildcardService: ) -> str | None: keyword = _normalize_wildcard_key(raw_key) if keyword in wildcard_dict: - return rng.choice(wildcard_dict[keyword]) + return self._pick_weighted_or_plain(wildcard_dict[keyword], rng) if "*" in keyword: regex_pattern = keyword.replace("*", ".*").replace("+", r"\+") @@ -400,7 +400,7 @@ class WildcardService: if compiled.match(key): aggregated.extend(values) if aggregated: - return rng.choice(aggregated) + return self._pick_weighted_or_plain(aggregated, rng) if "/" not in keyword: fallback_keyword = _normalize_wildcard_key(f"*/{keyword}") @@ -409,6 +409,39 @@ class WildcardService: return None + def _pick_weighted_or_plain( + self, values: list[str], rng: random.Random + ) -> str: + """Pick a value from the list, respecting N::weight prefix if present. + + When any value in the list uses the ``N::value`` weighted syntax with a + weight different from 1, the pick uses weighted random selection. When + no such weighting is present, a plain ``rng.choice`` is used (preserving + backward compatibility for unweighted wildcard files). + + In either case the ``N::`` prefix is always stripped from the returned + value, matching the behaviour of ``{...}`` option groups. + """ + # Fast path: skip weighting logic entirely when no :: syntax exists + if not any("::" in v for v in values): + return rng.choice(values) + + weighted_options: list[tuple[float, str]] = [] + for value in values: + weight = 1.0 + parts = value.split("::", 1) + if len(parts) == 2 and _is_numeric_string(parts[0].strip()): + weight = float(parts[0].strip()) + weighted_options.append((weight, value)) + + any_weighted = any(w != 1.0 for w, _ in weighted_options) + if any_weighted: + picked = self._weighted_choice(weighted_options, rng) + else: + picked = rng.choice(values) + + return self._strip_weight_prefix(picked) + def is_trigger_words_input(name: str) -> bool: return bool(_TRIGGER_WORD_PATTERN.match(name)) diff --git a/tests/services/test_wildcard_service.py b/tests/services/test_wildcard_service.py index a5fc7f9f..047539c0 100644 --- a/tests/services/test_wildcard_service.py +++ b/tests/services/test_wildcard_service.py @@ -139,3 +139,122 @@ def test_contains_dynamic_syntax_detects_wildcards_and_options(): assert contains_dynamic_syntax("__flower__") is True assert contains_dynamic_syntax("{red|blue}") is True assert contains_dynamic_syntax("{2$$, $$red|blue|green}") is True + + +# --------------------------------------------------------------------------- +# _pick_weighted_or_plain +# --------------------------------------------------------------------------- + + +def test_pick_weighted_or_plain_plain_values(monkeypatch, tmp_path): + """Plain values without :: are picked via rng.choice (fast path).""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + rng = random.Random(42) + + result = service._pick_weighted_or_plain(["red", "green", "blue"], rng) + assert result in {"red", "green", "blue"} + assert "::" not in result + + +def test_pick_weighted_or_plain_deterministic_with_seed(monkeypatch, tmp_path): + """Same seed produces the same result for plain values.""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + first = service._pick_weighted_or_plain(["a", "b", "c"], random.Random(99)) + second = service._pick_weighted_or_plain(["a", "b", "c"], random.Random(99)) + assert first == second + + +def test_pick_weighted_or_plain_weighted_values(monkeypatch, tmp_path): + """Weighted values use weighted selection and strip the N:: prefix.""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + values = ["3::apple", "1::banana"] + results = {"apple": 0, "banana": 0} + for seed in range(4000): + result = service._pick_weighted_or_plain(values, random.Random(seed)) + assert result in results, f"Unexpected result: {result!r}" + assert "::" not in result + results[result] += 1 + + total = results["apple"] + results["banana"] + # 3:1 weight → apple ≈ 75%, banana ≈ 25% + assert 2700 < results["apple"] < 3300, f"apple count out of range: {results['apple']}" + assert 700 < results["banana"] < 1300, f"banana count out of range: {results['banana']}" + + +def test_pick_weighted_or_plain_weight_one_values(monkeypatch, tmp_path): + """Values with explicit 1:: prefix have prefix stripped but are not weighted.""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + # All weights are 1.0 → no actual weighting, but :: prefix is stripped + values = ["1::foo", "1::bar"] + rng = random.Random(42) + results = {service._pick_weighted_or_plain(values, rng) for _ in range(200)} + assert results == {"foo", "bar"} + # Ensure the prefix is always stripped + for result in results: + assert "::" not in result + + +def test_pick_weighted_or_plain_mixed_weighted_and_plain(monkeypatch, tmp_path): + """Mixed list with some weighted and some unweighted values.""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + values = ["5::x", "y", "z"] # x has weight 5, y/z have default weight 1 + results = {"x": 0, "y": 0, "z": 0} + for seed in range(4000): + result = service._pick_weighted_or_plain(values, random.Random(seed)) + assert result in results + assert "::" not in result + results[result] += 1 + + # x (5) vs combined y+z (1+1=2) → ~71% / ~29% + x_pct = results["x"] / sum(results.values()) + assert 0.65 < x_pct < 0.78, f"x proportion out of range: {x_pct:.3f}" + + +def test_pick_weighted_or_plain_invalid_weight_prefix(monkeypatch, tmp_path): + """Invalid numeric prefix (e.g. 1.2.3) is NOT treated as a weight and + the prefix is NOT stripped, matching the updated strict regex.""" + service, _ = _make_service(monkeypatch, tmp_path) + + import random + rng = random.Random(42) + + # "1.2.3::a" is not a valid number → treated as plain text value + result = service._pick_weighted_or_plain(["1.2.3::a", "b"], rng) + # It should keep the full text including :: because the prefix isn't a + # valid numeric weight according to the strict regex + assert result == "1.2.3::a" or result == "b" + + +def test_pick_weighted_or_plain_glob_aggregation(monkeypatch, tmp_path): + """Weighted wildcard resolution through glob aggregation (__*__).""" + service, wildcards_dir = _make_service(monkeypatch, tmp_path) + wildcards_dir.mkdir() + + (wildcards_dir / "animals").mkdir() + (wildcards_dir / "animals" / "cat.txt").write_text("3::tabby\n1::persian\n", encoding="utf-8") + (wildcards_dir / "animals" / "dog.txt").write_text("retriever\npoodle\n", encoding="utf-8") + + # __animals/*__ aggregates all values across both files + # Weighted values should have :: stripped + results = {"tabby": 0, "persian": 0, "retriever": 0, "poodle": 0} + for seed in range(4000): + expanded = service.expand_text("__animals/*__", seed=seed) + assert expanded in results, f"Unexpected result: {expanded!r}" + assert "::" not in expanded + results[expanded] += 1 + + # tabby (3) vs persian (1) → ~75% / ~25% within the cat subset + cat_total = results["tabby"] + results["persian"] + if cat_total > 0: + tabby_pct = results["tabby"] / cat_total + assert 0.65 < tabby_pct < 0.85, f"tabby proportion out of range: {tabby_pct:.3f}"