tests: add simple gen test

midas/generator/generator.py

@@ -2,6 +2,7 @@ import ast
 import shutil
 from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Optional
 
 from midas.ast.location import Location
 import midas.ast.python as p
@@ -44,21 +45,28 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self._alias_count: int = 0
         self._scopes: list[Scope] = []
 
-    def generate(self, typed_ast: TypedAST, src_path: Path) -> Path:
+    def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST:
         self.rel_src_path = src_path.relative_to(self.workdir)
         self._typed_ast = typed_ast
         body: list[ast.stmt] = self._visit_body(typed_ast.stmts)
         module = ast.Module(body=body, type_ignores=[])
         module = ast.fix_missing_locations(module)
+        return module
+
+    def generate(
+        self, typed_ast: TypedAST, src_path: Path, out_path: Optional[Path] = None
+    ) -> Path:
+        module: ast.AST = self.generate_ast(typed_ast, src_path)
         compiled: str = ast.unparse(module)
-        out_path: Path = (self.build_dir / self.rel_src_path).resolve()
+        if out_path is None:
-        try:
+            out_path = (self.build_dir / self.rel_src_path).resolve()
-            _ = out_path.relative_to(self.build_dir)
+            try:
-        except ValueError:
+                _ = out_path.relative_to(self.build_dir)
-            raise ValueError(
+            except ValueError:
-                f"Directory traversal, {self.rel_src_path} points outside of parent directory"
+                raise ValueError(
-            )
+                    f"Directory traversal, {self.rel_src_path} points outside of parent directory"
-        out_path.parent.mkdir(parents=True, exist_ok=True)
+                )
+            out_path.parent.mkdir(parents=True, exist_ok=True)
         out_path.write_text(compiled)
         return out_path
 

tests/base.py

@@ -21,6 +21,10 @@ class Tester(ABC):
     @abstractmethod
     def namespace(self) -> str: ...
 
+    @property
+    def extension(self) -> str:
+        return "json"
+
     @property
     def base_dir(self) -> Path:
         return self.CASES_DIR / self.namespace
@@ -99,7 +103,7 @@ class Tester(ABC):
         return True
 
     def _result_path(self, test_path: Path) -> Path:
-        return test_path.parent / (test_path.name + ".ref.json")
+        return test_path.parent / (test_path.name + f".ref.{self.extension}")
 
     def _print_diff(self, diff: Iterator[str]):
         for line in diff:

tests/cases/generator/01_simple_types.midas

@@ -0,0 +1,14 @@
+type Meter = float
+type Second = float
+type MeterPerSecond = float
+
+extend Meter {
+    def __add__: fn(Meter, /) -> Meter
+    def __sub__: fn(Meter, /) -> Meter
+    def __truediv__: fn(Second, /) -> MeterPerSecond
+}
+
+extend Second {
+    def __add__: fn(Second, /) -> Second
+    def __sub__: fn(Second, /) -> Second
+}

tests/cases/generator/01_simple_types.py

@@ -0,0 +1,5 @@
+from midas import cast, Meter, Second
+
+distance: Meter = cast(Meter, 123.45)
+time: Second = cast(Second, 6.7)
+speed = distance / time

tests/cases/generator/01_simple_types.py.ref.txt

@@ -0,0 +1,79 @@
+Module(
+  body=[
+    ImportFrom(
+      module='midas',
+      names=[
+        alias(name='cast'),
+        alias(name='Meter'),
+        alias(name='Second')],
+      level=0),
+    Assign(
+      targets=[
+        Name(id='__midas_alias_0__')],
+      value=Constant(value=123.45)),
+    Assert(
+      test=Call(
+        func=Name(id='isinstance'),
+        args=[
+          Name(id='__midas_alias_0__'),
+          Name(id='float')],
+        keywords=[]),
+      msg=JoinedStr(
+        values=[
+          Constant(value='01_simple_types.py:L3:19: CastError: Cannot cast '),
+          FormattedValue(
+            value=Attribute(
+              value=Call(
+                func=Name(id='type'),
+                args=[
+                  Name(id='__midas_alias_0__')],
+                keywords=[]),
+              attr='__name__'),
+            conversion=-1),
+          Constant(value=' to float')])),
+    Assign(
+      targets=[
+        Name(id='distance')],
+      value=Name(id='__midas_alias_0__')),
+    Delete(
+      targets=[
+        Name(id='__midas_alias_0__')]),
+    Assign(
+      targets=[
+        Name(id='__midas_alias_1__')],
+      value=Constant(value=6.7)),
+    Assert(
+      test=Call(
+        func=Name(id='isinstance'),
+        args=[
+          Name(id='__midas_alias_1__'),
+          Name(id='float')],
+        keywords=[]),
+      msg=JoinedStr(
+        values=[
+          Constant(value='01_simple_types.py:L4:16: CastError: Cannot cast '),
+          FormattedValue(
+            value=Attribute(
+              value=Call(
+                func=Name(id='type'),
+                args=[
+                  Name(id='__midas_alias_1__')],
+                keywords=[]),
+              attr='__name__'),
+            conversion=-1),
+          Constant(value=' to float')])),
+    Assign(
+      targets=[
+        Name(id='time')],
+      value=Name(id='__midas_alias_1__')),
+    Delete(
+      targets=[
+        Name(id='__midas_alias_1__')]),
+    Assign(
+      targets=[
+        Name(id='speed')],
+      value=BinOp(
+        left=Name(id='distance'),
+        op=Div(),
+        right=Name(id='time')))],
+  type_ignores=[])

tests/generator.py

@@ -0,0 +1,55 @@
+import ast
+from dataclasses import dataclass
+from pathlib import Path
+
+from midas.checker.checker import TypeChecker
+from midas.checker.diagnostic import DiagnosticType
+from midas.generator.generator import Generator
+from midas.utils import TypedAST
+from tests.base import Tester
+
+
+@dataclass
+class CaseResult:
+    compiled_ast: ast.AST = ast.Module([], [])
+
+    def dumps(self) -> str:
+        return ast.dump(self.compiled_ast, indent=2)
+
+
+class GeneratorTester(Tester):
+    @property
+    def namespace(self) -> str:
+        return "generator"
+
+    @property
+    def extension(self) -> str:
+        return "txt"
+
+    def _list_tests(self) -> list[Path]:
+        return list(self.base_dir.rglob("*.py"))
+
+    def _exec_case(self, path: Path) -> CaseResult:
+        if not path.exists():
+            raise FileNotFoundError(f"Could not find test '{path}'")
+        if not path.is_file():
+            raise TypeError(f"Test '{path}' is not a file")
+
+        result: CaseResult = CaseResult()
+
+        checker = TypeChecker()
+        types_path: Path = path.with_suffix(".midas")
+        if types_path.exists():
+            checker.import_midas(types_path)
+
+        typed_ast: TypedAST = checker.type_check(path)
+
+        if not any(d.type == DiagnosticType.ERROR for d in checker.diagnostics):
+            generator = Generator(workdir=path.parent)
+            result.compiled_ast = generator.generate_ast(typed_ast, path)
+
+        return result
+
+
+if __name__ == "__main__":
+    GeneratorTester.main()