feat(cli): dump environment after compile

midas/checker/checker.py

@@ -32,6 +32,8 @@ class Checker(
     p.Expr.Visitor[Type],
     p.MidasType.Visitor[Type],
 ):
+    """A type checker which can use custom type definitions"""
+
     def __init__(self, locals: dict[p.Expr, int], file_path: Path):
         self.logger: logging.Logger = logging.getLogger("Checker")
         self.file_path: Path = file_path
@@ -271,7 +273,7 @@ class Checker(
 
     def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
         type: Type = stmt.value.accept(self) if stmt.value is not None else UnitType()
-        self.env.return_types.add(type)
+        self.env.return_types.append(type)
         raise ReturnException()
 
     def visit_binary_expr(self, expr: p.BinaryExpr) -> Type:

midas/checker/environment.py

@@ -6,15 +6,39 @@ from midas.checker.types import Type
 
 
 class Environment:
+    """
+    A scoped environment in which variables are defined
+
+    Each environment can inherit from a parent/enclosing environment.
+    """
+
     def __init__(self, enclosing: Optional[Environment] = None) -> None:
         self.enclosing: Optional[Environment] = enclosing
         self.values: dict[str, Type] = {}
-        self.return_types: set[Type] = set()
+        self.return_types: list[Type] = []
 
-    def define(self, name: str, value: Type):
+        self._children: list[Environment] = []
+        if enclosing is not None:
+            enclosing._children.append(self)
+
+    def define(self, name: str, value: Type) -> None:
+        """Define a variable in this environment
+
+        Args:
+            name (str): the name of the variable
+            value (Type): the value
+        """
         self.values[name] = value
 
     def get(self, name: str) -> Optional[Type]:
+        """Get a variable in the closest environment which has a definition for it
+
+        Args:
+            name (str): the name of the variable
+
+        Returns:
+            Optional[Type]: the value of the variable, or None if it was not found
+        """
         if name in self.values:
             return self.values[name]
         if self.enclosing is not None:
@@ -23,6 +47,15 @@ class Environment:
         return None
 
     def assign(self, name: str, value: Type) -> bool:
+        """Assign a new value to a variable in the environment it was defined in
+
+        Args:
+            name (str): the name of the variable
+            value (Type): the new value
+
+        Returns:
+            bool: True if the variable was assigned in this environment or an ancestor, False otherwise
+        """
         if name not in self.values:
             if self.enclosing is None:
                 return False
@@ -32,22 +65,78 @@ class Environment:
         return True
 
     def clear(self):
+        """Clear all definitions in this environment"""
         self.values = {}
 
     def get_at(self, distance: int, name: str) -> Optional[Type]:
+        """Get the value of a variable at a given distance
+
+        A distance of 0 looks up in this environment, 1 in the parent environment, etc.
+        This methods expects `distance` to be valid. An error will be raised if
+        the stack does not extend far enough to reach `distance`
+
+        Args:
+            distance (int): the scope distance
+            name (str): the name of the variable
+
+        Returns:
+            Optional[Type]: the value at the given distance, or None if it is not defined in that environment
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
         return self.ancestor(distance).values.get(name)
 
-    def assign_at(self, distance: int, name: str, value: Type):
+    def assign_at(self, distance: int, name: str, value: Type) -> None:
+        """Assign a new value to a variable at a given distance
+
+        A distance of 0 assigns in this environment, 1 in the parent environment, etc.
+
+        Args:
+            distance (int): the scope distance
+            name (str): the name of the variable
+            value (Type): the new value
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
         self.ancestor(distance).values[name] = value
 
     def ancestor(self, distance: int) -> Environment:
+        """Get the ancestor at a given distance
+
+        A distance of 0 references this environment, 1 the parent environment, etc.
+
+        Args:
+            distance (int): the scope distance
+
+        Returns:
+            Environment: the environment
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
         env: Environment = self
         for _ in range(distance):
             assert env.enclosing is not None
             env = env.enclosing
         return env
 
-    def flat_dict(self) -> dict:
+    def flat_dict(self) -> dict[str, Type]:
+        """Get the current environment including definitions in its ancestor as a flat dictionary
+
+        This method recursively combines this environment definitions with its ancestor's
+
+        Returns:
+            dict: the combined environment
+        """
         if self.enclosing is None:
             return self.values
         return self.enclosing.flat_dict() | self.values
+
+    def dump(self) -> dict:
+        return {
+            "values": self.values,
+            "return_types": self.return_types,
+            "children": [child.dump() for child in self._children],
+        }

midas/cli/main.py

@@ -1,4 +1,5 @@
 import ast
+import json
 import logging
 from dataclasses import dataclass
 from pathlib import Path
@@ -18,6 +19,7 @@ from midas.lexer.token import Token, TokenType
 from midas.parser.midas import MidasParser
 from midas.parser.python import PythonParser
 from midas.resolver.resolver import Resolver
+from midas.utils import UniversalJSONDumper
 
 
 @click.group()
@@ -40,6 +42,15 @@ def compile(file: TextIO):
     for diagnostic in diagnostics:
         print(diagnostic)
 
+    print(
+        json.dumps(
+            UniversalJSONDumper.dump(
+                checker.global_env, [("Environment", "_children")]
+            ),
+            indent=4,
+        )
+    )
+
 
 @midas.group()
 def utils():

midas/resolver/midas.py

@@ -5,6 +5,8 @@ from midas.checker.types import BaseType, SimpleType, Type
 
 
 class MidasResolver(m.Stmt.Visitor[None], m.Expr.Visitor[Type]):
+    """A resolver which evaluates Midas type definitions and build a registry"""
+
     def __init__(self) -> None:
         self._types: dict[str, Type] = {}
         self._operations: dict[tuple[Type, str, Type], Type] = {}
@@ -12,6 +14,17 @@ class MidasResolver(m.Stmt.Visitor[None], m.Expr.Visitor[Type]):
         self._define_builtin()
 
     def get_type(self, name: str) -> Type:
+        """Get a type from its name
+
+        Args:
+            name (str): the name of the type
+
+        Raises:
+            NameError: if the type is not defined
+
+        Returns:
+            Type: the type
+        """
         type: Optional[Type] = self._types.get(name)
         if type is None:
             raise NameError(f"Undefined type {name}")
@@ -20,11 +33,22 @@ class MidasResolver(m.Stmt.Visitor[None], m.Expr.Visitor[Type]):
     def get_operation_result(
         self, left: Type, operator: str, right: Type
     ) -> Optional[Type]:
+        """Get the resulting type of an operation
+
+        Args:
+            left (Type): the type of the left operand
+            operator (str): the operation name
+            right (Type): the type of the right operand
+
+        Returns:
+            Optional[Type]: the result type, or None if no matching operation was found
+        """
         operation: tuple[Type, str, Type] = (left, operator, right)
         result: Optional[Type] = self._operations.get(operation)
         return result
 
     def _define_builtin(self):
+        """Define builtin types and operations"""
         self.define_type("bool", BaseType(name="bool"))
         self.define_type("int", BaseType(name="int"))
         self.define_type("float", BaseType(name="float"))
@@ -37,12 +61,35 @@ class MidasResolver(m.Stmt.Visitor[None], m.Expr.Visitor[Type]):
         )
 
     def define_type(self, name: str, type: Type) -> Type:
+        """Define a type in the registry
+
+        Args:
+            name (str): the name of the type
+            type (Type): the type to define
+
+        Raises:
+            ValueError: if a type is already defined with that name
+
+        Returns:
+            Type: the defined type
+        """
         if name in self._types:
             raise ValueError(f"Type {name} already defined")
         self._types[name] = type
         return type
 
     def define_operation(self, left: Type, operator: str, right: Type, result: Type):
+        """Define an operation in the registry
+
+        Args:
+            left (Type): the type of the left operand
+            operator (str): the operation name
+            right (Type): the type of the right operand
+            result (Type): the result type
+
+        Raises:
+            ValueError: if an operation is already defined with these operands and name
+        """
         operation: tuple[Type, str, Type] = (left, operator, right)
         if operation in self._operations:
             raise ValueError(
@@ -51,6 +98,11 @@ class MidasResolver(m.Stmt.Visitor[None], m.Expr.Visitor[Type]):
         self._operations[operation] = result
 
     def resolve(self, stmts: list[m.Stmt]):
+        """Process a sequence of statements
+
+        Args:
+            stmts (list[m.Stmt]): the statements
+        """
         for stmt in stmts:
             stmt.accept(self)
 

midas/resolver/resolver.py

@@ -5,21 +5,41 @@ class ResolverError(Exception): ...
 
 
 class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
+    """A variable assignment and reference resolver
+
+    This class keeps track of which scope a variable is defined in and which
+    scope is referred to when a variable is referenced
+    """
+
     def __init__(self):
         self.locals: dict[p.Expr, int] = {}
         self.scopes: list[dict[str, bool]] = []
 
     def resolve(self, *objects: p.Stmt | p.Expr) -> None:
+        """Resolve the given statements or expressions"""
+
         for obj in objects:
             obj.accept(self)
 
     def begin_scope(self):
+        """Begin a new scope inside the current one"""
         self.scopes.append({})
 
     def end_scope(self):
+        """Close the current scope"""
         self.scopes.pop()
 
     def declare(self, name: str) -> None:
+        """Declare a variable in the current scope
+
+        This method must be called *before* evaluating the variable initializer
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            ResolverError: if the variable has already been declared in the current scope
+        """
         if len(self.scopes) == 0:
             return
         scope: dict[str, bool] = self.scopes[-1]
@@ -30,17 +50,42 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         scope[name] = False
 
     def define(self, name: str) -> None:
+        """Define a variable in the current scope
+
+        This method must be called *after* evaluating the variable initializer
+
+        Args:
+            name (str): the name of the variable
+        """
         if len(self.scopes) == 0:
             return
         self.scopes[-1][name] = True
 
     def resolve_local(self, expr: p.Expr, name: str) -> None:
+        """Resolve a variable reference and store the scope distance
+
+        This method associates to the variable expression a number representing
+        the "distance" of the variable declaration, i.e. the number of scope
+        levels to go "up" to find the closest declaration for that variable.
+
+        Args:
+            expr (p.Expr): the variable expression
+            name (str): the name of the variable
+        """
         for i, scope in enumerate(reversed(self.scopes)):
             if name in scope:
                 self.locals[expr] = i
                 return
 
     def resolve_function(self, function: p.Function) -> None:
+        """Resolve a function definition
+
+        This method creates a new scope for the function, resolves all the
+        parameter declarations and then the body.
+
+        Args:
+            function (p.Function): the function to resolve
+        """
         self.begin_scope()
         for param in function.all_args:
             self.declare(param.name)

midas/utils.py

@@ -0,0 +1,40 @@
+from typing import Any, Optional
+
+
+class UniversalJSONDumper:
+    @classmethod
+    def dump(
+        cls, obj: Any, include_keys: Optional[list[str | tuple[str, str]]] = None
+    ) -> Any:
+        if include_keys is None:
+            include_keys = []
+        return cls._dump(obj, include_keys, [])
+
+    @classmethod
+    def _dump(
+        cls, obj: Any, include_keys: list[str | tuple[str, str]], visited: list[Any]
+    ) -> Any:
+        if obj in visited:
+            return None
+        match obj:
+            case str() | int() | float() | None:
+                return obj
+            case list() | set() | tuple():
+                return [cls._dump(child, include_keys, visited) for child in obj]
+            case dict():
+                return {
+                    str(k): cls._dump(v, include_keys, visited) for k, v in obj.items()
+                }
+            case object():
+                visited.append(obj)
+                return {
+                    "_type": obj.__class__.__name__,
+                } | {
+                    k: cls._dump(v, include_keys, visited)
+                    for k, v in obj.__dict__.items()
+                    if not k.startswith("_")
+                    or k in include_keys
+                    or (obj.__class__.__name__, k) in include_keys
+                }
+            case _:
+                raise ValueError(f"Unsupported value: {obj}")