tests: add tests for python parser

tester.py

@@ -1,204 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import difflib
-import json
-import sys
-from dataclasses import asdict, dataclass, field
-from pathlib import Path
-from typing import Iterator, Optional
-
-from midas.ast.json_serializer import AstJsonSerializer
-from midas.ast.midas import Stmt
-from midas.lexer.base import MidasSyntaxError
-from midas.lexer.midas import MidasLexer
-from midas.lexer.token import Token
-from midas.parser.midas import MidasParser
-
-DEFAULT_BASE_DIR: Path = Path() / "tests"
-
-
-@dataclass
-class CaseResult:
-    tokens: Optional[list[dict]] = None
-    stmts: Optional[list[dict]] = None
-    errors: list[dict] = field(default_factory=list)
-
-    def dumps(self) -> str:
-        return json.dumps(asdict(self), indent=2)
-
-
-class Tester:
-    """A test runner to check for regressions in the lexer and parser"""
-
-    def __init__(self, base_dir: Path):
-        self.base_dir: Path = base_dir
-
-    def _list_tests(self) -> list[Path]:
-        return list(self.base_dir.rglob("*.midas"))
-
-    def run_all_tests(self) -> bool:
-        paths: list[Path] = self._list_tests()
-        return self.run_tests(paths)
-
-    def run_tests(self, tests: list[Path]) -> bool:
-        rule: str = "-" * 80
-        n: int = len(tests)
-        successes: int = 0
-        failures: int = 0
-
-        print(rule)
-        for i, test in enumerate(tests):
-            print(f"Case {i+1}/{n}: {test}")
-            success: bool = self._run_test(test)
-            if success:
-                successes += 1
-            else:
-                failures += 1
-
-        print(rule)
-        print(f"Success: {successes}/{n}")
-        print(f"Failed: {failures}/{n}")
-        print(rule)
-        return failures == 0
-
-    def _run_test(self, path: Path) -> bool:
-        result: CaseResult = self._exec_case(path)
-        result_path: Path = self._result_path(path)
-        expected: str = result_path.read_text()
-        actual: str = result.dumps()
-
-        if expected == actual:
-            return True
-
-        diff = difflib.unified_diff(
-            expected.splitlines(keepends=True),
-            actual.splitlines(keepends=True),
-            fromfile="Snapshot",
-            tofile="Result",
-        )
-        self._print_diff(diff)
-        return False
-
-    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()
-        content: str = path.read_text()
-        lexer: MidasLexer = MidasLexer(content)
-        tokens: list[Token] = []
-        try:
-            tokens = lexer.process()
-            result.tokens = [
-                {
-                    "type": token.type.name,
-                    "lexeme": token.lexeme,
-                    "line": token.position.line,
-                    "column": token.position.column,
-                }
-                for token in tokens
-            ]
-        except MidasSyntaxError as e:
-            result.errors.append(
-                {
-                    "type": "SyntaxError",
-                    "line": e.pos.line,
-                    "column": e.pos.column,
-                    "message": e.message,
-                }
-            )
-            return result
-
-        parser: MidasParser = MidasParser(tokens)
-        stmts: list[Stmt] = parser.parse()
-        result.stmts = AstJsonSerializer().serialize(stmts)
-        result.errors.extend(
-            [
-                {
-                    "line": e.token.position.line,
-                    "column": e.token.position.column,
-                    "message": e.message,
-                }
-                for e in parser.errors
-            ]
-        )
-        return result
-
-    def update_all_tests(self):
-        paths: list[Path] = self._list_tests()
-        return self.update_tests(paths)
-
-    def update_tests(self, tests: list[Path]):
-        updated: int = 0
-        for test in tests:
-            if self._update_test(test):
-                updated += 1
-        print(f"Updated {updated}/{len(tests)} tests")
-
-    def _update_test(self, path: Path) -> bool:
-        result: CaseResult = self._exec_case(path)
-        result_path: Path = self._result_path(path)
-        current: str = result_path.read_text()
-        new: str = result.dumps()
-        if current == new:
-            return False
-        result_path.write_text(new)
-        return True
-
-    def _result_path(self, test_path: Path) -> Path:
-        return test_path.parent / (test_path.name + ".ref.json")
-
-    def _print_diff(self, diff: Iterator[str]):
-        for line in diff:
-            if line.startswith("+") and not line.startswith("+++"):
-                print(f"\033[92m{line}\033[0m", end="")
-            elif line.startswith("-") and not line.startswith("---"):
-                print(f"\033[91m{line}\033[0m", end="")
-            else:
-                print(line, end="")
-        print()
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-D",
-        "--base-dir",
-        help="Base directory containing test files",
-        type=Path,
-        default=DEFAULT_BASE_DIR,
-    )
-    subparsers = parser.add_subparsers(dest="subcommand")
-
-    update = subparsers.add_parser("update")
-    update.add_argument("-a", "--all", action="store_true")
-    update.add_argument("FILE", type=Path, nargs="*")
-
-    run = subparsers.add_parser("run")
-    run.add_argument("-a", "--all", action="store_true")
-    run.add_argument("FILE", type=Path, nargs="*")
-    args = parser.parse_args()
-
-    tester: Tester = Tester(args.base_dir)
-
-    match args.subcommand:
-        case "update":
-            if args.all:
-                tester.update_all_tests()
-            else:
-                tester.update_tests(args.FILE)
-        case "run":
-            success: bool
-            if args.all:
-                success = tester.run_all_tests()
-            else:
-                success = tester.run_tests(args.FILE)
-            if not success:
-                sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()

tests/base.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import argparse
+import difflib
+import sys
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Iterator, Protocol
+
+
+class CaseResult(Protocol):
+    def dumps(self) -> str: ...
+
+
+class Tester(ABC):
+    """A test runner to check for regressions in the lexer and parser"""
+
+    CASES_DIR: Path = Path(__file__).parent / "cases"
+
+    @property
+    @abstractmethod
+    def namespace(self) -> str: ...
+
+    @property
+    def base_dir(self) -> Path:
+        return self.CASES_DIR / self.namespace
+
+    @abstractmethod
+    def _list_tests(self) -> list[Path]: ...
+
+    def run_all_tests(self) -> bool:
+        paths: list[Path] = self._list_tests()
+        return self.run_tests(paths)
+
+    def run_tests(self, tests: list[Path]) -> bool:
+        rule: str = "-" * 80
+        n: int = len(tests)
+        successes: int = 0
+        failures: int = 0
+
+        print(rule)
+        for i, test in enumerate(tests):
+            print(f"Case {i+1}/{n}: {test.relative_to(self.CASES_DIR)}")
+            success: bool = self._run_test(test)
+            if success:
+                successes += 1
+            else:
+                failures += 1
+
+        print(rule)
+        print(f"Success: {successes}/{n}")
+        print(f"Failed: {failures}/{n}")
+        print(rule)
+        return failures == 0
+
+    def _run_test(self, path: Path) -> bool:
+        result_path: Path = self._result_path(path)
+        if not result_path.exists():
+            print("Missing snapshot. Please run the update command first")
+            return False
+        result: CaseResult = self._exec_case(path)
+        expected: str = result_path.read_text()
+        actual: str = result.dumps()
+
+        if expected == actual:
+            return True
+
+        diff = difflib.unified_diff(
+            expected.splitlines(keepends=True),
+            actual.splitlines(keepends=True),
+            fromfile="Snapshot",
+            tofile="Result",
+        )
+        self._print_diff(diff)
+        return False
+
+    @abstractmethod
+    def _exec_case(self, path: Path) -> CaseResult: ...
+
+    def update_all_tests(self):
+        paths: list[Path] = self._list_tests()
+        return self.update_tests(paths)
+
+    def update_tests(self, tests: list[Path]):
+        updated: int = 0
+        for test in tests:
+            if self._update_test(test):
+                updated += 1
+        print(f"Updated {updated}/{len(tests)} tests")
+
+    def _update_test(self, path: Path) -> bool:
+        result: CaseResult = self._exec_case(path)
+        result_path: Path = self._result_path(path)
+        current: str = result_path.read_text() if result_path.exists() else ""
+        new: str = result.dumps()
+        if current == new:
+            return False
+        result_path.write_text(new)
+        return True
+
+    def _result_path(self, test_path: Path) -> Path:
+        return test_path.parent / (test_path.name + ".ref.json")
+
+    def _print_diff(self, diff: Iterator[str]):
+        for line in diff:
+            if line.startswith("+") and not line.startswith("+++"):
+                print(f"\033[92m{line}\033[0m", end="")
+            elif line.startswith("-") and not line.startswith("---"):
+                print(f"\033[91m{line}\033[0m", end="")
+            else:
+                print(line, end="")
+        print()
+
+    @classmethod
+    def main(cls):
+        parser = argparse.ArgumentParser()
+        subparsers = parser.add_subparsers(dest="subcommand")
+
+        update = subparsers.add_parser("update")
+        update.add_argument("-a", "--all", action="store_true")
+        update.add_argument("FILE", type=Path, nargs="*")
+
+        run = subparsers.add_parser("run")
+        run.add_argument("-a", "--all", action="store_true")
+        run.add_argument("FILE", type=Path, nargs="*")
+        args = parser.parse_args()
+
+        tester: Tester = cls()
+
+        match args.subcommand:
+            case "update":
+                if args.all:
+                    tester.update_all_tests()
+                else:
+                    tester.update_tests(args.FILE)
+            case "run":
+                success: bool
+                if args.all:
+                    success = tester.run_all_tests()
+                else:
+                    success = tester.run_tests(args.FILE)
+                if not success:
+                    sys.exit(1)

tests/cases/midas-parser/01_simple_types.midas.ref.json

@@ -2170,8 +2170,8 @@
   "stmts": [
     {
       "_type": "SimpleTypeStmt",
-      "template": null,
       "name": "Custom",
+      "template": null,
       "base": {
         "_type": "TypeExpr",
         "name": "float",
@@ -2182,8 +2182,8 @@
     },
     {
       "_type": "SimpleTypeStmt",
-      "template": null,
       "name": "Latitude",
+      "template": null,
       "base": {
         "_type": "TypeExpr",
         "name": "float",
@@ -2219,8 +2219,8 @@
     },
     {
       "_type": "SimpleTypeStmt",
-      "template": null,
       "name": "Longitude",
+      "template": null,
       "base": {
         "_type": "TypeExpr",
         "name": "float",
@@ -2256,6 +2256,7 @@
     },
     {
       "_type": "SimpleTypeStmt",
+      "name": "Difference",
       "template": {
         "_type": "TemplateExpr",
         "type": {
@@ -2265,7 +2266,6 @@
           "optional": false
         }
       },
-      "name": "Difference",
       "base": {
         "_type": "TypeExpr",
         "name": "T",

tests/cases/python-parser/01_simple_types.py

@@ -0,0 +1,14 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+df: Frame[
+    verified: bool,
+    birth_year: int,
+    height: float + ( _ > 0 ) + ( _ < 250 ),
+    name: str,
+    date: datetime,
+    float,
+    unknown: _,
+    _
+]

tests/cases/python-parser/01_simple_types.py.ref.json

@@ -0,0 +1,85 @@
+{
+  "stmts": [
+    {
+      "_type": "TypeAssign",
+      "name": "df",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "verified",
+            "type": {
+              "_type": "BaseType",
+              "base": "bool",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "birth_year",
+            "type": {
+              "_type": "BaseType",
+              "base": "int",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "(_ > 0) + (_ < 250)"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "name",
+            "type": {
+              "_type": "BaseType",
+              "base": "str",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "date",
+            "type": {
+              "_type": "BaseType",
+              "base": "datetime",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "float",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "unknown",
+            "type": null
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "_",
+              "param": null
+            }
+          }
+        ]
+      }
+    }
+  ]
+}

tests/cases/python-parser/02_custom_types.py

@@ -0,0 +1,29 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+import midas
+
+midas.using("02_custom_types.midas")
+
+df: Frame[
+    location: GeoLocation
+]
+
+lat: Column[GeoLocation] = df["location"].lat
+lon: Column[GeoLocation] = df["location"].lon
+
+lat + lon
+
+lat1: Latitude = lat[0]
+lat2: Latitude = lat[1]
+lat_diff: Difference[Latitude] = lat2 - lat1
+
+df2: Frame[
+    age: int + (_ >= 0),
+    height: float + (_ >= 0),
+]
+df2_bis: Frame[
+    age: int + Positive,
+    height: float + Positive,
+]

tests/cases/python-parser/02_custom_types.py.ref.json

@@ -0,0 +1,162 @@
+{
+  "stmts": [
+    {
+      "_type": "ExpressionStmt",
+      "expr": {
+        "_type": "CallExpr",
+        "callee": {
+          "_type": "GetExpr",
+          "object": {
+            "_type": "VariableExpr",
+            "name": "midas"
+          },
+          "name": "using"
+        },
+        "arguments": [
+          {
+            "_type": "LiteralExpr",
+            "value": "02_custom_types.midas"
+          }
+        ],
+        "keywords": {}
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "location",
+            "type": {
+              "_type": "BaseType",
+              "base": "GeoLocation",
+              "param": null
+            }
+          }
+        ]
+      }
+    },
+    {
+      "_type": "ExpressionStmt",
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "lat"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "lon"
+        }
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "lat_diff",
+      "type": {
+        "_type": "BaseType",
+        "base": "Difference",
+        "param": {
+          "_type": "BaseType",
+          "base": "Latitude",
+          "param": null
+        }
+      }
+    },
+    {
+      "_type": "AssignStmt",
+      "targets": [
+        {
+          "_type": "VariableExpr",
+          "name": "lat_diff"
+        }
+      ],
+      "value": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "lat2"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "lat1"
+        }
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df2",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "age",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "int",
+                "param": null
+              },
+              "constraint": "_ >= 0"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "_ >= 0"
+            }
+          }
+        ]
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df2_bis",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "age",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "int",
+                "param": null
+              },
+              "constraint": "Positive"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "Positive"
+            }
+          }
+        ]
+      }
+    }
+  ]
+}

tests/cases/python-parser/03_functions.py

@@ -0,0 +1,15 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+
+def func(
+    col1: Column[float + (0 <= _ <= 1)],
+    col2: Column[float + (0 <= _ <= 1)],
+) -> Column[float + (0 <= _ <= 2)]:
+    result: Column[float + (0 <= _ <= 2)] = col1 + col2
+    return result
+
+
+def func2(a: int, /, b: float, *, c: str):
+    pass

tests/cases/python-parser/03_functions.py.ref.json

@@ -0,0 +1,149 @@
+{
+  "stmts": [
+    {
+      "_type": "Function",
+      "name": "func",
+      "posonlyargs": [],
+      "args": [
+        {
+          "name": "col1",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 1"
+            }
+          },
+          "default": null
+        },
+        {
+          "name": "col2",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 1"
+            }
+          },
+          "default": null
+        }
+      ],
+      "sink": null,
+      "kwonlyargs": [],
+      "kw_sink": null,
+      "returns": {
+        "_type": "BaseType",
+        "base": "Column",
+        "param": {
+          "_type": "ConstraintType",
+          "type": {
+            "_type": "BaseType",
+            "base": "float",
+            "param": null
+          },
+          "constraint": "0 <= _ <= 2"
+        }
+      },
+      "body": [
+        {
+          "_type": "TypeAssign",
+          "name": "result",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 2"
+            }
+          }
+        },
+        {
+          "_type": "AssignStmt",
+          "targets": [
+            {
+              "_type": "VariableExpr",
+              "name": "result"
+            }
+          ],
+          "value": {
+            "_type": "BinaryExpr",
+            "left": {
+              "_type": "VariableExpr",
+              "name": "col1"
+            },
+            "operator": "+",
+            "right": {
+              "_type": "VariableExpr",
+              "name": "col2"
+            }
+          }
+        },
+        {
+          "_type": "ReturnStmt",
+          "value": {
+            "_type": "VariableExpr",
+            "name": "result"
+          }
+        }
+      ]
+    },
+    {
+      "_type": "Function",
+      "name": "func2",
+      "posonlyargs": [
+        {
+          "name": "a",
+          "type": {
+            "_type": "BaseType",
+            "base": "int",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "args": [
+        {
+          "name": "b",
+          "type": {
+            "_type": "BaseType",
+            "base": "float",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "sink": null,
+      "kwonlyargs": [
+        {
+          "name": "c",
+          "type": {
+            "_type": "BaseType",
+            "base": "str",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "kw_sink": null,
+      "returns": null,
+      "body": []
+    }
+  ]
+}

tests/midas.py

@@ -0,0 +1,82 @@
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from midas.ast.midas import Stmt
+from midas.lexer.base import MidasSyntaxError
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token
+from midas.parser.midas import MidasParser
+from tests.base import Tester
+from tests.serializer.midas import MidasAstJsonSerializer
+
+
+@dataclass
+class CaseResult:
+    tokens: Optional[list[dict]] = None
+    stmts: Optional[list[dict]] = None
+    errors: list[dict] = field(default_factory=list)
+
+    def dumps(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+
+
+class MidasTester(Tester):
+    @property
+    def namespace(self) -> str:
+        return "midas-parser"
+
+    def _list_tests(self) -> list[Path]:
+        return list(self.base_dir.rglob("*.midas"))
+
+    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()
+        content: str = path.read_text()
+        lexer: MidasLexer = MidasLexer(content)
+        tokens: list[Token] = []
+        try:
+            tokens = lexer.process()
+            result.tokens = [
+                {
+                    "type": token.type.name,
+                    "lexeme": token.lexeme,
+                    "line": token.position.line,
+                    "column": token.position.column,
+                }
+                for token in tokens
+            ]
+        except MidasSyntaxError as e:
+            result.errors.append(
+                {
+                    "type": "SyntaxError",
+                    "line": e.pos.line,
+                    "column": e.pos.column,
+                    "message": e.message,
+                }
+            )
+            return result
+
+        parser: MidasParser = MidasParser(tokens)
+        stmts: list[Stmt] = parser.parse()
+        result.stmts = MidasAstJsonSerializer().serialize(stmts)
+        result.errors.extend(
+            [
+                {
+                    "line": e.token.position.line,
+                    "column": e.token.position.column,
+                    "message": e.message,
+                }
+                for e in parser.errors
+            ]
+        )
+        return result
+
+
+if __name__ == "__main__":
+    MidasTester.main()

tests/python.py

@@ -0,0 +1,46 @@
+import ast
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Optional
+
+from midas.ast.python import Stmt
+from midas.parser.python import PythonParser
+from tests.base import Tester
+from tests.serializer.python import PythonAstJsonSerializer
+
+
+@dataclass
+class CaseResult:
+    stmts: Optional[list[dict]] = None
+
+    def dumps(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+
+
+class PythonTester(Tester):
+    @property
+    def namespace(self) -> str:
+        return "python-parser"
+
+    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()
+        content: str = path.read_text()
+        tree: ast.Module = ast.parse(content)
+
+        parser: PythonParser = PythonParser()
+        stmts: list[Stmt] = parser.parse_module(tree)
+        result.stmts = PythonAstJsonSerializer().serialize(stmts)
+        return result
+
+
+if __name__ == "__main__":
+    PythonTester.main()

tests/serializer/midas.py

@@ -23,7 +23,7 @@ from midas.ast.midas import (
 )
 
 
-class AstJsonSerializer(Stmt.Visitor[dict], Expr.Visitor[dict]):
+class MidasAstJsonSerializer(Stmt.Visitor[dict], Expr.Visitor[dict]):
     """An AST serializer which produces a JSON-compatible structure"""
 
     def serialize(self, stmts: list[Stmt]) -> list[dict]:
@@ -40,8 +40,8 @@ class AstJsonSerializer(Stmt.Visitor[dict], Expr.Visitor[dict]):
     def visit_simple_type_stmt(self, stmt: SimpleTypeStmt) -> dict:
         return {
             "_type": "SimpleTypeStmt",
-            "template": self._serialize_optional(stmt.template),
             "name": stmt.name.lexeme,
+            "template": self._serialize_optional(stmt.template),
             "base": stmt.base.accept(self),
             "constraint": self._serialize_optional(stmt.constraint),
         }

tests/serializer/python.py

@@ -0,0 +1,230 @@
+import ast
+from typing import Optional, Sequence, Type
+
+from midas.ast.python import (
+    AssignStmt,
+    BaseType,
+    BinaryExpr,
+    CallExpr,
+    CompareExpr,
+    ConstraintType,
+    Expr,
+    ExpressionStmt,
+    FrameColumn,
+    FrameType,
+    Function,
+    GetExpr,
+    LiteralExpr,
+    LogicalExpr,
+    MidasType,
+    ReturnStmt,
+    SetExpr,
+    Stmt,
+    TypeAssign,
+    UnaryExpr,
+    VariableExpr,
+)
+
+unary_ops: dict[Type[ast.unaryop], str] = {
+    ast.Invert: "~",
+    ast.Not: "not",
+    ast.UAdd: "+",
+    ast.USub: "-",
+}
+binary_ops: dict[Type[ast.operator], str] = {
+    ast.Add: "+",
+    ast.Sub: "-",
+    ast.Mult: "*",
+    ast.MatMult: "@",
+    ast.Div: "/",
+    ast.Mod: "%",
+    ast.LShift: "<<",
+    ast.RShift: ">>",
+    ast.BitOr: "|",
+    ast.BitXor: "^",
+    ast.BitAnd: "&",
+    ast.FloorDiv: "//",
+    ast.Pow: "**",
+}
+compare_ops: dict[Type[ast.cmpop], str] = {
+    ast.Eq: "==",
+    ast.NotEq: "!=",
+    ast.Lt: "<",
+    ast.LtE: "<=",
+    ast.Gt: ">",
+    ast.GtE: ">=",
+    ast.Is: "is",
+    ast.IsNot: "is not",
+    ast.In: "in",
+    ast.NotIn: "not in",
+}
+boolean_ops: dict[Type[ast.boolop], str] = {
+    ast.And: "and",
+    ast.Or: "or",
+}
+
+
+class PythonAstJsonSerializer(
+    Stmt.Visitor[dict], Expr.Visitor[dict], MidasType.Visitor[dict]
+):
+    """An AST serializer which produces a JSON-compatible structure"""
+
+    def serialize(self, stmts: list[Stmt]) -> list[dict]:
+        return [stmt.accept(self) for stmt in stmts]
+
+    def _serialize_optional(
+        self, element: Optional[Stmt | Expr | MidasType]
+    ) -> Optional[dict]:
+        if element is None:
+            return None
+        return element.accept(self)
+
+    def _serialize_list(
+        self, elements: Sequence[Stmt | Expr | MidasType]
+    ) -> list[dict]:
+        return [element.accept(self) for element in elements]
+
+    def visit_base_type(self, node: BaseType) -> dict:
+        return {
+            "_type": "BaseType",
+            "base": node.base,
+            "param": self._serialize_optional(node.param),
+        }
+
+    def visit_constraint_type(self, node: ConstraintType) -> dict:
+        return {
+            "_type": "ConstraintType",
+            "type": node.type.accept(self),
+            "constraint": ast.unparse(node.constraint),
+        }
+
+    def visit_frame_column(self, node: FrameColumn) -> dict:
+        return {
+            "_type": "FrameColumn",
+            "name": node.name,
+            "type": self._serialize_optional(node.type),
+        }
+
+    def visit_frame_type(self, node: FrameType) -> dict:
+        return {
+            "_type": "FrameType",
+            "columns": self._serialize_list(node.columns),
+        }
+
+    def visit_expression_stmt(self, stmt: ExpressionStmt) -> dict:
+        return {
+            "_type": "ExpressionStmt",
+            "expr": stmt.expr.accept(self),
+        }
+
+    def _serialize_argument(self, arg: Function.Argument) -> dict:
+        return {
+            "name": arg.name,
+            "type": self._serialize_optional(arg.type),
+            "default": self._serialize_optional(arg.default),
+        }
+
+    def visit_function(self, stmt: Function) -> dict:
+        return {
+            "_type": "Function",
+            "name": stmt.name,
+            "posonlyargs": [self._serialize_argument(arg) for arg in stmt.posonlyargs],
+            "args": [self._serialize_argument(arg) for arg in stmt.args],
+            "sink": (
+                self._serialize_argument(stmt.sink) if stmt.sink is not None else None
+            ),
+            "kwonlyargs": [self._serialize_argument(arg) for arg in stmt.kwonlyargs],
+            "kw_sink": (
+                self._serialize_argument(stmt.kw_sink)
+                if stmt.kw_sink is not None
+                else None
+            ),
+            "returns": self._serialize_optional(stmt.returns),
+            "body": self._serialize_list(stmt.body),
+        }
+
+    def visit_type_assign(self, stmt: TypeAssign) -> dict:
+        return {
+            "_type": "TypeAssign",
+            "name": stmt.name,
+            "type": stmt.type.accept(self),
+        }
+
+    def visit_assign_stmt(self, stmt: AssignStmt) -> dict:
+        return {
+            "_type": "AssignStmt",
+            "targets": self._serialize_list(stmt.targets),
+            "value": stmt.value.accept(self),
+        }
+
+    def visit_return_stmt(self, stmt: ReturnStmt) -> dict:
+        return {
+            "_type": "ReturnStmt",
+            "value": self._serialize_optional(stmt.value),
+        }
+
+    def visit_binary_expr(self, expr: BinaryExpr) -> dict:
+        return {
+            "_type": "BinaryExpr",
+            "left": expr.left.accept(self),
+            "operator": binary_ops[expr.operator.__class__],
+            "right": expr.right.accept(self),
+        }
+
+    def visit_compare_expr(self, expr: CompareExpr) -> dict:
+        return {
+            "_type": "CompareExpr",
+            "left": expr.left.accept(self),
+            "operator": compare_ops[expr.operator.__class__],
+            "right": expr.right.accept(self),
+        }
+
+    def visit_unary_expr(self, expr: UnaryExpr) -> dict:
+        return {
+            "_type": "UnaryExpr",
+            "operator": unary_ops[expr.operator.__class__],
+            "right": expr.right.accept(self),
+        }
+
+    def visit_call_expr(self, expr: CallExpr) -> dict:
+        return {
+            "_type": "CallExpr",
+            "callee": expr.callee.accept(self),
+            "arguments": self._serialize_list(expr.arguments),
+            "keywords": {name: arg.accept(self) for name, arg in expr.keywords.items()},
+        }
+
+    def visit_get_expr(self, expr: GetExpr) -> dict:
+        return {
+            "_type": "GetExpr",
+            "object": expr.object.accept(self),
+            "name": expr.name,
+        }
+
+    def visit_literal_expr(self, expr: LiteralExpr) -> dict:
+        return {
+            "_type": "LiteralExpr",
+            "value": expr.value,
+        }
+
+    def visit_variable_expr(self, expr: VariableExpr) -> dict:
+        return {
+            "_type": "VariableExpr",
+            "name": expr.name,
+        }
+
+    def visit_logical_expr(self, expr: LogicalExpr) -> dict:
+        return {
+            "_type": "LogicalExpr",
+            "left": expr.left.accept(self),
+            "operator": boolean_ops[expr.operator.__class__],
+            "right": expr.right.accept(self),
+        }
+
+    def visit_set_expr(self, expr: SetExpr) -> dict:
+        return {
+            "_type": "SetExpr",
+            "object": expr.object.accept(self),
+            "name": expr.name,
+            "value": expr.value.accept(self),
+        }