From 552946ff737d424ba98d5689a29f70730c030d6a Mon Sep 17 00:00:00 2001 From: imi Date: Mon, 2 Feb 2026 23:15:23 +0100 Subject: [PATCH] v1.0.0 --- .github/workflows/publish.yml | 28 ++++++++ LICENSE | 21 ++++++ api/autocomplete.py | 6 +- api/tokenize.py | 27 ++++---- pyproject.toml | 30 ++++++++ requirements.txt | 4 +- tests/__init__.py | 0 tests/mock_tags.csv | 7 ++ tests/test_autocomplete.py | 64 +++++++++++++++++ tests/test_hooks.py | 125 ++++++++++++++++++++++++++++++++++ tests/test_nodes_logic.py | 91 +++++++++++++++++++++++++ tests/test_scheduler.py | 94 +++++++++++++++++++++++++ tests/test_tokenizer.py | 70 +++++++++++++++++++ 13 files changed, 551 insertions(+), 16 deletions(-) create mode 100644 .github/workflows/publish.yml create mode 100644 LICENSE create mode 100644 pyproject.toml create mode 100644 tests/__init__.py create mode 100644 tests/mock_tags.csv create mode 100644 tests/test_autocomplete.py create mode 100644 tests/test_hooks.py create mode 100644 tests/test_nodes_logic.py create mode 100644 tests/test_scheduler.py create mode 100644 tests/test_tokenizer.py diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..3529040 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,28 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + +permissions: + contents: read + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'Enferlain' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + submodules: true + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..6443d7a --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Enferlain + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/api/autocomplete.py b/api/autocomplete.py index 7dd009c..d218294 100644 --- a/api/autocomplete.py +++ b/api/autocomplete.py @@ -187,7 +187,7 @@ class TagDatabase: return [] query_lower = query.lower().strip() - if not query_lower: + if not query_lower or len(query_lower) > 255: return [] results: List[Tuple[int, TagEntry, Optional[str]]] = [] @@ -339,6 +339,10 @@ if _HAS_SERVER and PromptServer: return web.json_response({"error": "Invalid JSON"}, status=400) query = data.get("query", "") + if len(query) > 255: + return web.json_response( + {"results": [], "tag_count": 0, "error": "Query too long"} + ) limit = min(data.get("limit", 20), 100) # Cap at 100 tag_file = data.get( "tag_file" diff --git a/api/tokenize.py b/api/tokenize.py index 45bcd66..4e518fe 100644 --- a/api/tokenize.py +++ b/api/tokenize.py @@ -46,6 +46,13 @@ def strip_a1111_syntax(text: str) -> str: text = text.replace("\\[", "\x00LBRACK\x00") text = text.replace("\\]", "\x00RBRACK\x00") + # Remove weight specifications: :1.2) at end of emphasis + # This handles (word:1.2) + text = re.sub(r":[\d.]+(?=\))", "", text) + + # Remove any remaining weight specs that might be floating + text = re.sub(r":[\d.]+(?=\s|$)", "", text) + # Handle bracket expressions: [A|B|C] or [from:to:when] # Keep the LONGEST option for worst-case token counting def keep_longest_option(match): @@ -55,10 +62,8 @@ def strip_a1111_syntax(text: str) -> str: # It's alternation [A|B|C] - find longest option options = content.split("|") # Remove any trailing :number from last option (scheduled alternation) - # e.g., [A|B:0.5] -> options = ["A", "B:0.5"] last = options[-1] if ":" in last: - # Check if it ends with :number colon_match = re.match(r"^(.+?)::?[\d.]+$", last) if colon_match: options[-1] = colon_match.group(1) @@ -70,31 +75,25 @@ def strip_a1111_syntax(text: str) -> str: # It's scheduling [from:to:when] - keep longer of from/to parts = content.split(":") if len(parts) >= 2: - # parts[0] = from, parts[1] = to, parts[2+] = when (numbers) from_part = parts[0] to_part = parts[1] if len(parts) > 1 else "" - # Return the longer of from/to - return from_part if len(from_part) >= len(to_part) else to_part + + # Compare based on tokenizable length (ignoring parens) for better worst-case estimation + clean_from = from_part.replace("(", "").replace(")", "").strip() + clean_to = to_part.replace("(", "").replace(")", "").strip() + + return from_part if len(clean_from) >= len(clean_to) else to_part return content else: - # Simple bracket emphasis [word] - just return content return content # Process bracket expressions - non-greedy match for innermost brackets - # Repeat to handle nested structures prev_text = None while prev_text != text: prev_text = text text = re.sub(r"\[([^\[\]]*)\]", keep_longest_option, text) - # Remove weight specifications: :1.2) at end of emphasis - # This handles (word:1.2) - text = re.sub(r":[\d.]+(?=\))", "", text) - - # Remove any remaining weight specs that might be floating - text = re.sub(r":[\d.]+(?=\s|$)", "", text) - # Remove parentheses (they're just syntax markers now) text = text.replace("(", " ") text = text.replace(")", " ") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..08a87da --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[project] +name = "a1111-prompt-node" +description = "A custom ComfyUI node that implements A1111-style prompt handling with proper isolation and emphasis math." +version = "1.0.0" +license = {file = "LICENSE"} +authors = [ + { name = "Enferlain" } +] +dependencies = ["lark", "torch", "aiohttp"] + +[project.urls] +Repository = "https://github.com/Enferlain/ComfyUI-A1111-cond" +"Bug Tracker" = "https://github.com/Enferlain/ComfyUI-A1111-cond/issues" + +[tool.comfy] +PublisherId = "Enferlain" +DisplayName = "A1111 Prompt Node" +Icon = "" +includes = [ + "nodes/", + "parser/", + "api/", + "js/", + "data/", + "__init__.py", + "hooks.py", + "LICENSE", + "README.md", + "requirements.txt" +] diff --git a/requirements.txt b/requirements.txt index ade3771..53d8401 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ -lark \ No newline at end of file +lark +torch +aiohttp \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/mock_tags.csv b/tests/mock_tags.csv new file mode 100644 index 0000000..ad97cec --- /dev/null +++ b/tests/mock_tags.csv @@ -0,0 +1,7 @@ +1girl,0,6000000,"1girls,sole_female" +1boy,0,1000000,"1boys,sole_male" +solo,0,8000000, +highres,5,5000000,"high_res" +hat,0,500000,"headwear" +artist_name,1,100, +character_name,4,1000, diff --git a/tests/test_autocomplete.py b/tests/test_autocomplete.py new file mode 100644 index 0000000..585f254 --- /dev/null +++ b/tests/test_autocomplete.py @@ -0,0 +1,64 @@ +import unittest +import sys +import os +from pathlib import Path + +# Add the project root to sys.path +sys.path.append(str(Path(__file__).parent.parent)) + +from api.autocomplete import TagDatabase, TagEntry + + +class TestAutocomplete(unittest.TestCase): + def setUp(self): + self.db = TagDatabase() + self.mock_csv = Path(__file__).parent / "mock_tags.csv" + self.db.load_csv(self.mock_csv) + + def test_prefix_search(self): + # Search for "1g" -> should find "1girl" + results = self.db.search("1g") + self.assertGreater(len(results), 0) + self.assertEqual(results[0]["name"], "1girl") + + def test_alias_matching(self): + # Search for "sole_female" -> should find "1girl" (via alias) + results = self.db.search("sole_female") + self.assertGreater(len(results), 0) + self.assertEqual(results[0]["name"], "1girl") + self.assertEqual(results[0]["matched_alias"], "sole_female") + + def test_sorting_by_count(self): + # "1girl" has 6M, "1boy" has 1M + # Searching for "1" should return both, girl first + results = self.db.search("1") + names = [r["name"] for r in results] + self.assertIn("1girl", names) + self.assertIn("1boy", names) + # Verify 1girl appears before 1boy due to count + girl_idx = names.index("1girl") + boy_idx = names.index("1boy") + self.assertLess(girl_idx, boy_idx) + + def test_tag_types(self): + # artist_name is type 1 + results = self.db.search("artist_name") + self.assertEqual(results[0]["type"], 1) + + # character_name is type 4 + results = self.db.search("character_name") + self.assertEqual(results[0]["type"], 4) + + def test_empty_query(self): + results = self.db.search("") + self.assertEqual(len(results), 0) + + def test_long_query(self): + # Even if the API handles it, the search method should be robust + long_query = "a" * 300 + results = self.db.search(long_query) + self.assertEqual(len(results), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_hooks.py b/tests/test_hooks.py new file mode 100644 index 0000000..865fe92 --- /dev/null +++ b/tests/test_hooks.py @@ -0,0 +1,125 @@ +import unittest +import sys +from pathlib import Path +import torch +from unittest.mock import MagicMock + +# Add project root to sys.path +sys.path.append(str(Path(__file__).parent.parent)) + +# Mock ComfyUI dependencies +from unittest.mock import MagicMock + +mock_hooks = MagicMock() + + +class DummyHook: + def __init__(self, *args, **kwargs): + pass + + def clone(self): + return self + + +mock_hooks.TransformerOptionsHook = DummyHook +mock_hooks.EnumHookScope = MagicMock() + +sys.modules["comfy.hooks"] = mock_hooks +sys.modules["server"] = MagicMock() + + +from hooks import A1111StepConditioningHook + + +class TestHooks(unittest.TestCase): + def test_get_step_from_sigma(self): + # mock sample_sigmas [10.0, 8.0, 5.0, 0.0] + # step 0: 10.0 -> 8.0 + # step 1: 8.0 -> 5.0 + # step 2: 5.0 -> 0.0 + sigmas = torch.tensor([10.0, 8.0, 5.0, 0.0]) + hook = A1111StepConditioningHook(step_embeddings=[]) + + # Exact match start + self.assertEqual(hook.get_step_from_sigma(10.0, sigmas), 0) + # Inside first step + self.assertEqual(hook.get_step_from_sigma(9.0, sigmas), 0) + # Boundary: 8.0 is the START of step 1 + self.assertEqual(hook.get_step_from_sigma(8.0, sigmas), 1) + # Inside second step + self.assertEqual(hook.get_step_from_sigma(7.0, sigmas), 1) + # Exact match end of last step: 0.0 is the start of the "final" interval + # (Though samplers usually stop at 0, this returns len-2 as per hooks.py) + self.assertEqual(hook.get_step_from_sigma(0.0, sigmas), 2) + + def test_step_index_scaling(self): + # sampler has 10 steps, we have 20 embeddings + sigmas = torch.linspace(10, 0, 11) # length 11 -> 10 steps + embeddings = [i for i in range(21)] # length 21 -> step 0 to 20 + hook = A1111StepConditioningHook(step_embeddings=embeddings) + + # Mock the wrapper logic part that calculates step_idx + def get_scaled_step(sigma_val): + raw_idx = hook.get_step_from_sigma(sigma_val, sigmas) + actual_steps = len(sigmas) - 1 + emb_steps = len(embeddings) - 1 + return round(raw_idx * emb_steps / actual_steps) + + # Start (step 0 of 10) -> step 0 of 20 + self.assertEqual(get_scaled_step(10.0), 0) + # Middle (step 5 of 10) -> step 10 of 20 + self.assertEqual(get_scaled_step(5.0), 10) + # End (step 9 of 10) -> step 18 of 20 (it's N-1 for last index usually) + # actually last sigma is step 9. + self.assertEqual(get_scaled_step(sigmas[9].item()), 18) + + def test_model_function_wrapper_swaps(self): + # Prepare 3 embeddings (0, 1, 2) + emb0 = torch.ones((1, 1, 4)) * 0.0 + emb1 = torch.ones((1, 1, 4)) * 1.0 + emb2 = torch.ones((1, 1, 4)) * 2.0 + hook = A1111StepConditioningHook( + step_embeddings=[(emb0, None), (emb1, None), (emb2, None)] + ) + hook.debug = True + + # Mock arguments + apply_model_func = MagicMock(return_value="model_output") + input_x = torch.zeros((1, 4)) + timestep = torch.tensor([10.0]) + sigmas = torch.tensor([10.0, 5.0, 0.0]) # 2 steps: [10->5, 5->0] + + c = { + "c_crossattn": torch.ones((2, 1, 4)) * -1.0, # original (batch=2) + "transformer_options": {"sample_sigmas": sigmas}, + } + cond_or_uncond = [0, 1] # 0 = positive (swap), 1 = negative (keep) + + args = { + "input": input_x, + "timestep": timestep, + "c": c, + "cond_or_uncond": cond_or_uncond, + } + + # Step 0 (sigma 10.0) -> raw_idx 0 + # Scaling: round(0 * 2 / 2) = 0 -> emb0 + hook.model_function_wrapper(apply_model_func, args) + + call_args = apply_model_func.call_args[1] + new_cond = call_args["c_crossattn"] + self.assertTrue(torch.allclose(new_cond[0], emb0[0])) + + # Step 1 (sigma 5.0) -> raw_idx 1 + # Scaling: round(1 * 2 / 2) = 1 -> emb1 + timestep = torch.tensor([5.0]) + args["timestep"] = timestep + hook.model_function_wrapper(apply_model_func, args) + + call_args = apply_model_func.call_args[1] + new_cond = call_args["c_crossattn"] + self.assertTrue(torch.allclose(new_cond[0], emb1[0])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_nodes_logic.py b/tests/test_nodes_logic.py new file mode 100644 index 0000000..79ec080 --- /dev/null +++ b/tests/test_nodes_logic.py @@ -0,0 +1,91 @@ +import unittest +import torch + +# We extract the pure logic from prompt_node.py to test it in isolation +# without dealing with ComfyUI's complex package/import system. + + +def extract_weights(token_batches): + all_weights = [] + for batch in token_batches: + if not batch: + all_weights.append([]) + continue + batch_weights = [ + t[1] if hasattr(t, "__getitem__") and len(t) > 1 else 1.0 for t in batch + ] + all_weights.append(batch_weights) + return all_weights + + +def apply_direct_scaling(cond, weights, normalization): + cond = cond.clone() + flat_weights = [] + for chunk_weights in weights: + flat_weights.extend(chunk_weights) + + seq_len = cond.shape[1] + if len(flat_weights) < seq_len: + flat_weights = flat_weights + [1.0] * (seq_len - len(flat_weights)) + elif len(flat_weights) > seq_len: + flat_weights = flat_weights[:seq_len] + + batch_size = cond.shape[0] + multipliers = torch.tensor( + [flat_weights] * batch_size, device=cond.device, dtype=cond.dtype + ) + + if normalization: + original_mean = cond.mean() + + multipliers = multipliers.unsqueeze(-1).expand_as(cond) + cond = cond * multipliers + + if normalization: + new_mean = cond.mean() + if new_mean.abs() > 1e-8: + scale_factor = original_mean / new_mean + cond = cond * scale_factor + return cond + + +class TestNodesLogic(unittest.TestCase): + def test_extract_weights(self): + tokens = [[(101, 1.0), (102, 1.2), (103, 0.8)], [(201, 1.5), (202, 1.0)]] + weights = extract_weights(tokens) + self.assertEqual(len(weights), 2) + self.assertEqual(weights[0], [1.0, 1.2, 0.8]) + self.assertEqual(weights[1], [1.5, 1.0]) + + def test_apply_direct_scaling_basic(self): + cond = torch.ones((1, 3, 4)) + weights = [[1.0, 2.0, 0.5]] + scaled = apply_direct_scaling(cond, weights, normalization=False) + self.assertTrue(torch.allclose(scaled[0, 0], torch.ones(4) * 1.0)) + self.assertTrue(torch.allclose(scaled[0, 1], torch.ones(4) * 2.0)) + self.assertTrue(torch.allclose(scaled[0, 2], torch.ones(4) * 0.5)) + + def test_apply_direct_scaling_with_break(self): + cond = torch.ones((1, 4, 4)) + weights = [[1.2, 1.0], [0.8, 1.0]] + scaled = apply_direct_scaling(cond, weights, normalization=False) + expected_weights = [1.2, 1.0, 0.8, 1.0] + for i, w in enumerate(expected_weights): + self.assertTrue(torch.allclose(scaled[0, i], torch.ones(4) * w)) + + def test_normalization_mean_preservation(self): + cond = torch.randn((1, 10, 64)) + original_mean = cond.mean().item() + weights = [[2.0] * 10] + + scaled_no_norm = apply_direct_scaling(cond, weights, normalization=False) + self.assertAlmostEqual( + scaled_no_norm.mean().item(), original_mean * 2, places=4 + ) + + scaled_norm = apply_direct_scaling(cond, weights, normalization=True) + self.assertAlmostEqual(scaled_norm.mean().item(), original_mean, places=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..bf361c2 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,94 @@ +import unittest +import sys +import os +from pathlib import Path + +# Add the project root to sys.path so we can import modules +sys.path.append(str(Path(__file__).parent.parent)) + +from parser.scheduler import get_prompt_schedule + + +class TestScheduler(unittest.TestCase): + def test_basic_scheduling(self): + # [before:after:when] + # At 10 steps, 0.5 is step 5 + schedule = get_prompt_schedule("[cat:dog:0.5]", 10) + # Expect (5, "cat"), (10, "dog") + self.assertEqual(len(schedule), 2) + self.assertEqual(schedule[0], (5, "cat")) + self.assertEqual(schedule[1], (10, "dog")) + + def test_step_based_scheduling(self): + # [before:after:step] + schedule = get_prompt_schedule("[cat:dog:7]", 10) + self.assertEqual(len(schedule), 2) + self.assertEqual(schedule[0], (7, "cat")) + self.assertEqual(schedule[1], (10, "dog")) + + def test_alternation(self): + # [A|B] + schedule = get_prompt_schedule("[white|black]", 4) + # Alternation changes every step + self.assertEqual(len(schedule), 4) + self.assertEqual(schedule[0][1], "white") + self.assertEqual(schedule[1][1], "black") + self.assertEqual(schedule[2][1], "white") + self.assertEqual(schedule[3][1], "black") + + def test_scheduled_alternation_add(self): + # [A|B:0.5] - Yield A until 50%, then alternate + schedule = get_prompt_schedule("[A|B:0.5]", 4) + # steps=4, 0.5=2 + # Step 1: A + # Step 2: A + # Step 3: A (cycle (3-1)%2=0) - wait, check implementation + # Actually _at_step calculates idx = (step - 1) % len(options) + # Step 1: yield options[0] -> A + # Step 2: yield options[0] -> A + # Step 3: (3-1)%2 = 0 -> A <-- Is this right? A1111 behavior check. + # Let's check my implementation in scheduler.py: + # idx = (step - 1) % len(options) + # So Step 3 (step > 2) yields options[0] = A + # Step 4: (4-1)%2 = 1 -> B + + self.assertEqual(schedule[0][1], "A") + self.assertEqual(schedule[1][1], "A") + self.assertEqual(schedule[2][1], "A") + self.assertEqual(schedule[3][1], "B") + + def test_scheduled_alternation_remove(self): + # [A|B::0.5] - Alternate until 50%, then first tag + schedule = get_prompt_schedule("[A|B::0.5]", 4) + # Step 1: A + # Step 2: B + # Step 3: A (first tag) + # Step 4: A (first tag) + self.assertEqual(schedule[0][1], "A") + self.assertEqual(schedule[1][1], "B") + self.assertEqual(schedule[2][1], "A") + self.assertEqual(schedule[3][1], "A") + + def test_nested_logic(self): + # [(red:1.2):[blue:green:0.5]:0.3] + # steps=10, 0.3=3, 0.5=5 + schedule = get_prompt_schedule("[(red:1.2):[blue:green:0.5]:0.3]", 10) + # Sequence expected: + # Step 1-3: (red:1.2) + # Step 4-5: blue + # Step 6-10: green + self.assertEqual(len(schedule), 3) + self.assertEqual(schedule[0], (3, "(red:1.2)")) + self.assertEqual(schedule[1], (5, "blue")) + self.assertEqual(schedule[2], (10, "green")) + + def test_break_isolation(self): + # a BREAK [b:c:0.5] + schedule = get_prompt_schedule("a BREAK [b:c:0.5]", 10) + self.assertEqual(len(schedule), 2) + self.assertEqual(schedule[0], (5, "a BREAK b")) + self.assertEqual(schedule[1], (10, "a BREAK c")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py new file mode 100644 index 0000000..70e7474 --- /dev/null +++ b/tests/test_tokenizer.py @@ -0,0 +1,70 @@ +import unittest +import sys +from pathlib import Path +from unittest.mock import MagicMock + +# Add project root to sys.path +sys.path.append(str(Path(__file__).parent.parent)) + +# Mock ComfyUI dependencies before importing logic +sys.modules["aiohttp"] = MagicMock() +sys.modules["server"] = MagicMock() +sys.modules["comfy"] = MagicMock() +sys.modules["comfy.sd1_clip"] = MagicMock() + +from api.tokenize import strip_a1111_syntax + + +class TestTokenizer(unittest.TestCase): + def test_emphasis_stripping(self): + # (word:1.2) -> word + self.assertEqual(strip_a1111_syntax("(dog:1.3)"), "dog") + # (word) -> word + self.assertEqual(strip_a1111_syntax("(cat)"), "cat") + # multiple emphasis + self.assertEqual( + strip_a1111_syntax("a (beautiful:1.2) (cat:1.5)"), "a beautiful cat" + ) + + def test_bracket_emphasis(self): + # [word] -> word + self.assertEqual(strip_a1111_syntax("[bird]"), "bird") + # [word:0.5] -> word + self.assertEqual(strip_a1111_syntax("[fish:0.5]"), "fish") + + def test_scheduling_stripping(self): + # [from:to:when] -> keeps longest + self.assertEqual(strip_a1111_syntax("[mountain:ocean:0.5]"), "mountain") + self.assertEqual( + strip_a1111_syntax("[forest:very_large_city:0.5]"), "very_large_city" + ) + # [add:when] -> add + self.assertEqual(strip_a1111_syntax("[glasses:10]"), "glasses") + # [remove::when] -> remove + self.assertEqual(strip_a1111_syntax("[hat::15]"), "hat") + + def test_alternation_stripping(self): + # [A|B|C] -> keeps longest + self.assertEqual(strip_a1111_syntax("[red|blue|green]"), "green") + self.assertEqual(strip_a1111_syntax("[a|bc|d]"), "bc") + + def test_scheduled_alternation_stripping(self): + # [A|B:0.5] -> keeps longest + self.assertEqual(strip_a1111_syntax("[apple|banana:0.5]"), "banana") + self.assertEqual(strip_a1111_syntax("[apple|banana::0.5]"), "banana") + + def test_escaped_stripping(self): + # \( \) -> ( ) + self.assertEqual(strip_a1111_syntax(r"\(escaped\)"), "(escaped)") + self.assertEqual(strip_a1111_syntax(r"\[bracket\]"), "[bracket]") + + def test_nested_stripping(self): + # [(red:1.2):[blue|green]:0.5] + # Inner [blue|green] -> green (len 5 vs blue len 4) + # Then [(red:1.2):green:0.5] -> red (len 3+tags removal) vs green (len 5) + # Result should be "green" because "red" is len 3 and "green" is len 5 + self.assertEqual(strip_a1111_syntax("[(red:1.2):green:0.5]"), "green") + + +if __name__ == "__main__": + unittest.main()