diff --git a/src/ftl_extract/process/kwargs_extractor.py b/src/ftl_extract/process/kwargs_extractor.py index de14c40..61e1a0f 100644 --- 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 index 2b42be3..ed3a7fa 100644 --- 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"}