feat: adds type extraction tests

Author: chalmerlowe

scripts/microgenerator/generate.py

@@ -27,9 +27,8 @@
 import argparse
 import glob
 import logging
-import re
 from collections import defaultdict
-from typing import List, Dict, Any, Iterator
+from typing import List, Dict, Any
 
 from . import name_utils
 from . import utils
@@ -83,6 +82,11 @@ def _get_type_str(self, node: ast.AST | None) -> str | None:
         # Handles forward references as strings, e.g., '"Dataset"'
         if isinstance(node, ast.Constant):
             return repr(node.value)
+        # Handles | union types, e.g., int | float
+        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
+            left_str = self._get_type_str(node.left)
+            right_str = self._get_type_str(node.right)
+            return f"{left_str} | {right_str}"
         return None  # Fallback for unhandled types
 
     def _collect_types_from_node(self, node: ast.AST | None) -> None:

scripts/microgenerator/tests/unit/test_generate_analyzer.py

@@ -15,9 +15,24 @@
 #
 
 import ast
-from scripts.microgenerator.generate import CodeAnalyzer
+import pytest
+from scripts.microgenerator.generate import parse_code, CodeAnalyzer
 
 
+# --- Mock Types ---
+class MyClass:
+    pass
+
+
+class AnotherClass:
+    pass
+
+
+class YetAnotherClass:
+    pass
+
+
+# --- Existing Tests ---
 def test_codeanalyzer_finds_class():
     code = """
 class MyClass:
@@ -100,3 +115,106 @@ class MyClass:
     assert len(analyzer.structure) == 1
     assert analyzer.structure[0]["class_name"] == "MyClass"
     assert len(analyzer.structure[0]["methods"]) == 0
+
+
+# --- Test Data for Parameterization ---
+TYPE_TEST_CASES = [
+    pytest.param(
+        """class TestClass:
+    def func(self, a: int, b: str) -> bool: return True""",
+        [("a", "int"), ("b", "str")],
+        "bool",
+        id="simple_types",
+    ),
+    pytest.param(
+        """from typing import Optional
+class TestClass:
+    def func(self, a: Optional[int]) -> str | None: return 'hello'""",
+        [("a", "Optional[int]")],
+        "str | None",
+        id="optional_union_none",
+    ),
+    pytest.param(
+        """from typing import Union
+class TestClass:
+    def func(self, a: int | float, b: Union[str, bytes]) -> None: pass""",
+        [("a", "int | float"), ("b", "Union[str, bytes]")],
+        "None",
+        id="union_types",
+    ),
+    pytest.param(
+        """from typing import List, Dict, Tuple
+class TestClass:
+    def func(self, a: List[int], b: Dict[str, float]) -> Tuple[int, str]: return (1, 'a')""",
+        [("a", "List[int]"), ("b", "Dict[str, float]")],
+        "Tuple[int, str]",
+        id="generic_types",
+    ),
+    pytest.param(
+        """import datetime
+from scripts.microgenerator.tests.unit.test_generate_analyzer import MyClass
+class TestClass:
+    def func(self, a: datetime.date, b: MyClass) -> MyClass: return b""",
+        [("a", "datetime.date"), ("b", "MyClass")],
+        "MyClass",
+        id="imported_types",
+    ),
+    pytest.param(
+        """from scripts.microgenerator.tests.unit.test_generate_analyzer import AnotherClass, YetAnotherClass
+class TestClass:
+    def func(self, a: 'AnotherClass') -> 'YetAnotherClass': return AnotherClass()""",
+        [("a", "'AnotherClass'")],
+        "'YetAnotherClass'",
+        id="forward_refs",
+    ),
+    pytest.param(
+        """class TestClass:
+    def func(self, a, b): return a + b""",
+        [("a", None), ("b", None)],  # No annotations means type is None
+        None,
+        id="no_annotations",
+    ),
+    pytest.param(
+        """from typing import List, Optional, Dict, Union, Any
+class TestClass:
+    def func(self, a: List[Optional[Dict[str, Union[int, str]]]]) -> Dict[str, Any]: return {}""",
+        [("a", "List[Optional[Dict[str, Union[int, str]]]]")],
+        "Dict[str, Any]",
+        id="complex_nested",
+    ),
+    pytest.param(
+        """from typing import Literal
+class TestClass:
+    def func(self, a: Literal['one', 'two']) -> Literal[True]: return True""",
+        [("a", "Literal['one', 'two']")],
+        "Literal[True]",
+        id="literal_type",
+    ),
+]
+
+
+# --- Tests ---
+class TestCodeAnalyzerArgsReturns:
+    @pytest.mark.parametrize(
+        "code_snippet, expected_args, expected_return", TYPE_TEST_CASES
+    )
+    def test_type_extraction(self, code_snippet, expected_args, expected_return):
+        structure, imports, types = parse_code(code_snippet)
+
+        assert len(structure) == 1, "Should parse one class"
+        class_info = structure[0]
+        assert class_info["class_name"] == "TestClass"
+
+        assert len(class_info["methods"]) == 1, "Should find one method"
+        method_info = class_info["methods"][0]
+        assert method_info["method_name"] == "func"
+
+        # Extract args, skipping 'self'
+        extracted_args = []
+        for arg in method_info.get("args", []):
+            if arg["name"] == "self":
+                continue
+            extracted_args.append((arg["name"], arg["type"]))
+
+        assert extracted_args == expected_args
+        assert method_info.get("return_type") == expected_return