Fix #4 (FTL-Extract can't see nested kwargs) (#5)

andrew000 · Aug 18, 2024 · 7c50f7e · 7c50f7e
1 parent a4f50a3
commit 7c50f7e
Show file tree

Hide file tree

Showing 2 changed files with 168 additions and 9 deletions.
diff --git a/src/ftl_extract/process/kwargs_extractor.py b/src/ftl_extract/process/kwargs_extractor.py
@@ -3,6 +3,36 @@
 from ftl_extract.matcher import FluentKey
 
 
+def _extract_kwargs_from_variable_reference(
+    variable_reference: ast.VariableReference,
+    kwargs: set[str],
+) -> None:
+    kwargs.add(variable_reference.id.name)
+
+
+def _extract_kwargs_from_selector_expression(
+    selector_expression: ast.SelectExpression,
+    kwargs: set[str],
+) -> None:
+    if isinstance(selector_expression.selector, ast.VariableReference):
+        _extract_kwargs_from_variable_reference(selector_expression.selector, kwargs)
+
+    for variant in selector_expression.variants:
+        for placeable in variant.value.elements:
+            if isinstance(placeable, ast.Placeable):
+                _extract_kwargs_from_placeable(placeable, kwargs)
+
+
+def _extract_kwargs_from_placeable(placeable: ast.Placeable, kwargs: set[str]) -> None:
+    expression = placeable.expression
+
+    if isinstance(expression, ast.VariableReference):
+        _extract_kwargs_from_variable_reference(expression, kwargs)
+
+    elif isinstance(expression, ast.SelectExpression):
+        _extract_kwargs_from_selector_expression(expression, kwargs)
+
+
 def extract_kwargs(key: FluentKey) -> set[str]:
     kwargs: set[str] = set()
 
@@ -14,14 +44,6 @@ def extract_kwargs(key: FluentKey) -> set[str]:
 
     for placeable in key.translation.value.elements:
         if isinstance(placeable, ast.Placeable):
-            expression = placeable.expression
-
-            if isinstance(expression, ast.VariableReference):
-                kwargs.add(expression.id.name)
-
-            elif isinstance(expression, ast.SelectExpression) and isinstance(
-                expression.selector, ast.VariableReference
-            ):
-                kwargs.add(expression.selector.id.name)
+            _extract_kwargs_from_placeable(placeable, kwargs)
 
     return kwargs
diff --git a/tests/test_kwargs_extractor.py b/tests/test_kwargs_extractor.py
@@ -161,3 +161,140 @@ def test_extracts_selector_variable_name_from_select_expression() -> None:
         )
     )
     assert kwargs == {"user_status"}
+
+
+def test_nested_extraction() -> None:
+    kwargs = extract_kwargs(
+        FluentKey(
+            code_path=Path("test.py"),
+            key="trade-waiting_for_answer",
+            translation=ast.Message(
+                id=ast.Identifier(name="nested-key"),
+                value=ast.Pattern(
+                    elements=[
+                        ast.TextElement(value="nested-key\n"),
+                        ast.Placeable(
+                            expression=ast.SelectExpression(
+                                selector=ast.VariableReference(
+                                    id=ast.Identifier(name="first_level_key")
+                                ),
+                                variants=[
+                                    ast.Variant(
+                                        key=ast.NumberLiteral(value="1"),
+                                        value=ast.Pattern(
+                                            elements=[
+                                                ast.TextElement(value="✅ "),
+                                                ast.Placeable(
+                                                    expression=ast.SelectExpression(
+                                                        selector=ast.VariableReference(
+                                                            id=ast.Identifier(
+                                                                name="second_level_key"
+                                                            )
+                                                        ),
+                                                        variants=[
+                                                            ast.Variant(
+                                                                key=ast.NumberLiteral(value="1"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(value="OK")
+                                                                    ]
+                                                                ),
+                                                            ),
+                                                            ast.Variant(
+                                                                key=ast.NumberLiteral(value="2"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(value="NO")
+                                                                    ]
+                                                                ),
+                                                            ),
+                                                            ast.Variant(
+                                                                key=ast.Identifier(name="other"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(
+                                                                            value="ANOTHER"
+                                                                        )
+                                                                    ]
+                                                                ),
+                                                                default=True,
+                                                            ),
+                                                        ],
+                                                    )
+                                                ),
+                                            ]
+                                        ),
+                                    ),
+                                    ast.Variant(
+                                        key=ast.NumberLiteral(value="0"),
+                                        value=ast.Pattern(
+                                            elements=[
+                                                ast.TextElement(value="❌ "),
+                                                ast.Placeable(
+                                                    expression=ast.SelectExpression(
+                                                        selector=ast.VariableReference(
+                                                            id=ast.Identifier(
+                                                                name="second_level_key"
+                                                            )
+                                                        ),
+                                                        variants=[
+                                                            ast.Variant(
+                                                                key=ast.NumberLiteral(value="1"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(value="OK")
+                                                                    ]
+                                                                ),
+                                                            ),
+                                                            ast.Variant(
+                                                                key=ast.NumberLiteral(value="2"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(value="NO")
+                                                                    ]
+                                                                ),
+                                                            ),
+                                                            ast.Variant(
+                                                                key=ast.Identifier(name="other"),
+                                                                value=ast.Pattern(
+                                                                    elements=[
+                                                                        ast.TextElement(
+                                                                            value="ANOTHER"
+                                                                        )
+                                                                    ]
+                                                                ),
+                                                                default=True,
+                                                            ),
+                                                        ],
+                                                    )
+                                                ),
+                                            ]
+                                        ),
+                                    ),
+                                    ast.Variant(
+                                        key=ast.Identifier(name="other"),
+                                        value=ast.Pattern(
+                                            elements=[
+                                                ast.TextElement(value="⏳ "),
+                                                ast.Placeable(
+                                                    expression=ast.VariableReference(
+                                                        id=ast.Identifier(name="second_level_key")
+                                                    )
+                                                ),
+                                                ast.TextElement(value=" ANOTHER"),
+                                            ]
+                                        ),
+                                        default=True,
+                                    ),
+                                ],
+                            )
+                        ),
+                    ]
+                ),
+            ),
+            locale="en",
+            position=0,
+        )
+    )
+
+    assert kwargs == {"first_level_key", "second_level_key"}