fix: use generic Difference type in example

examples/00_syntax_prototype/02_custom_types.py

@@ -21,7 +21,7 @@ lat + lon  # Invalid operation
 # Registered operations are permitted
 lat1: Latitude = lat[0]
 lat2: Latitude = lat[1]
-lat_diff: LatitudeDiff = lat2 - lat1  # Valid operation
+lat_diff: Difference[Latitude] = lat2 - lat1  # Valid operation
 
 # In addition to the type, a column can have one or more constraints, either defined inline or in a separate file
 df2: Frame[

midas/ast/printer.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import ast
 import io
 from contextlib import contextmanager
 from enum import Enum, auto
 from typing import Generator, Generic, Optional, Protocol, TypeVar
 
 import midas.ast.midas as m
+import midas.ast.python as p
 
 
 class _Level(Enum):
@@ -346,3 +348,32 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str]):
     def visit_type_expr(self, expr: m.TypeExpr):
         template: str = expr.template.accept(self) if expr.template is not None else ""
         return f"{expr.name.lexeme}{template}{'?' if expr.optional else ''}"
+
+
+class PythonAstPrinter(AstPrinter, p.MidasType.Visitor[None]):
+    def visit_base_type(self, node: p.BaseType) -> None:
+        self._write_line("BaseType")
+        with self._child_level():
+            self._write_line(f"base: {node.base}")
+            self._write_optional_child("param", node.param)
+            constraint_str: str = "None"
+            if node.constraint is not None:
+                constraint_str = ast.unparse(node.constraint)
+            self._write_line(f"constraint: {constraint_str}", last=True)
+
+    def visit_frame_column(self, node: p.FrameColumn) -> None:
+        self._write_line("FrameColumn")
+        with self._child_level():
+            self._write_line(f"name: {node.name}")
+            self._write_optional_child("type", node.type, last=True)
+
+    def visit_frame_type(self, node: p.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level():
+            self._write_line("columns", last=True)
+            with self._child_level():
+                for i, col in enumerate(node.columns):
+                    self._idx = i
+                    if i == len(node.columns) - 1:
+                        self._mark_last()
+                    col.accept(self)

midas/ast/python.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import ast
+from dataclasses import dataclass
+from typing import Generic, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class MidasType(ABC):
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_base_type(self, node: BaseType) -> T: ...
+
+        @abstractmethod
+        def visit_frame_column(self, node: FrameColumn) -> T: ...
+
+        @abstractmethod
+        def visit_frame_type(self, node: FrameType) -> T: ...
+
+
+@dataclass(frozen=True)
+class BaseType(MidasType):
+    base: str
+    param: Optional[MidasType]
+    constraint: Optional[ast.expr] = None
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_base_type(self)
+
+
+@dataclass(frozen=True)
+class FrameColumn(MidasType):
+    name: Optional[str]
+    type: Optional[MidasType]
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_frame_column(self)
+
+
+@dataclass(frozen=True)
+class FrameType(MidasType):
+    columns: list[FrameColumn]
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_frame_type(self)

midas/cli/main.py

@@ -1,6 +1,53 @@
+import ast
+from typing import Optional, TextIO
+
 import click
 
+from midas.ast.printer import PythonAstPrinter
+from midas.parser.python import PythonParser
 
-@click.command()
+
+@click.group()
 def midas():
     click.echo("Welcome to Midas!")
+
+
+@midas.command()
+@click.argument("file", type=click.File("r"))
+def compile(file: TextIO):
+    raise NotImplementedError
+
+
+@midas.group()
+def utils():
+    pass
+
+
+@utils.command()
+@click.option("-o", "--output", type=click.File("w"))
+@click.option("-p", "--parse", is_flag=True)
+@click.argument("file", type=click.File("r"))
+def dump_ast(output: Optional[TextIO], parse: bool, file: TextIO):
+    source: str = file.read()
+    tree: ast.Module = ast.parse(source, filename=file.name)
+    dump: str
+
+    if parse:
+        parser = PythonParser()
+        parser.visit(tree)
+        printer = PythonAstPrinter()
+        dump = ""
+        for name, annotation in parser.annotations:
+            dump += f"{name} = "
+            if annotation is None:
+                dump += "None"
+            else:
+                dump += printer.print(annotation)
+            dump += "\n"
+    else:
+        dump = ast.dump(tree, indent=4)
+
+    if output is None:
+        click.echo(dump)
+    else:
+        output.write(dump)

midas/parser/python.py

@@ -0,0 +1,92 @@
+import ast
+from typing import Any, Optional
+
+from midas.ast.python import BaseType, FrameColumn, FrameType, MidasType
+
+
+class InvalidSyntaxError(Exception):
+    pass
+
+
+class UnsupportedSyntaxError(Exception):
+    def __init__(self, expr: ast.expr) -> None:
+        super().__init__(f"Unsupported syntax: {ast.unparse(expr)}")
+
+
+class PythonParser(ast.NodeVisitor):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.annotations: list[tuple[str, Optional[MidasType]]] = []
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> Any:
+        match node:
+            case ast.AnnAssign(
+                target=ast.Name(id=target), annotation=annotation, simple=1
+            ):
+                self.annotations.append(
+                    (target, self._parse_type(annotation, root=True))
+                )
+
+            case _:
+                print(f"Unsupported annotation: {ast.unparse(node)}")
+
+    def _parse_type(
+        self, type_expr: ast.expr, root: bool = False
+    ) -> Optional[MidasType]:
+        match type_expr:
+            case ast.Subscript(value=ast.Name(id="Frame"), slice=schema):
+                return self._parse_frame_type(schema)
+
+            case ast.Subscript(value=ast.Name(id=name), slice=param):
+                return BaseType(
+                    base=name, param=self._parse_type(param), constraint=None
+                )
+
+            case ast.Name(id=name):
+                return BaseType(base=name, param=None, constraint=None)
+
+            case ast.BinOp(left=left_expr, op=ast.Add(), right=right_expr):
+                print("Constraints not implemented yet")
+                return None
+
+            case _:
+                if root:
+                    return None
+                raise UnsupportedSyntaxError(type_expr)
+
+    def _parse_frame_type(self, schema: ast.expr) -> FrameType:
+        columns: list[FrameColumn] = []
+
+        match schema:
+            case ast.Tuple(elts=cols):
+                for col in cols:
+                    columns.append(self._parse_frame_column(col))
+            case ast.Slice() | ast.Name():
+                columns.append(self._parse_frame_column(schema))
+            case _:
+                raise UnsupportedSyntaxError(schema)
+
+        return FrameType(columns=columns)
+
+    def _parse_frame_column(self, column: ast.expr) -> FrameColumn:
+        match column:
+            case ast.Name():
+                return FrameColumn(name=None, type=self._parse_type(column))
+            case ast.Slice(lower=ast.Name(id=name), upper=type_expr):
+                if name == "_":
+                    name = None
+
+                type: Optional[MidasType] = None
+                match type_expr:
+                    case None:
+                        raise InvalidSyntaxError("Missing column type")
+                    case ast.Name(id="_"):
+                        type = None
+                    case ast.expr():
+                        type = self._parse_type(type_expr)
+                    case _:
+                        raise UnsupportedSyntaxError(type_expr)
+                return FrameColumn(name=name, type=type)
+            case _:
+                raise UnsupportedSyntaxError(column)

pyproject.toml

@@ -4,8 +4,16 @@ version = "0.1.0"
 description = "A static-first type checking framework for Python data-frames"
 readme = "README.md"
 requires-python = ">=3.11"
+authors = [
+    { name = "Louis Heredero", email = "louis.heredero@students.hevs.ch" },
+]
+classifiers = ["Programming Language :: Python :: 3"]
 dependencies = ["click>=8.4.1"]
 
+[project.urls]
+Homepage = "https://git.kbk28.ch/HEL/midas"
+Repository = "https://git.kbk28.ch/HEL/midas"
+
 [project.scripts]
 midas = "midas.cli.main:midas"