docs(checker): add documentation to checker, resolvers, etc.

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
@@ -73,9 +75,23 @@ class Checker(
         )
 
     def evaluate(self, expr: p.Expr) -> Type:
+        """Evaluate the type of an expression
+
+        Args:
+            expr (p.Expr): the expression to evaluate
+
+        Returns:
+            Type: the type of the given expression
+        """
         return expr.accept(self)
 
     def evaluate_block(self, block: list[p.Stmt], env: Environment) -> None:
+        """Evaluate a sequence of statements
+
+        Args:
+            block (list[p.Stmt]): the statements to evaluate
+            env (Environment): the environment in which to evaluate
+        """
         previous_env: Environment = self.env
         self.env = env
         for stmt in block:
@@ -86,6 +102,14 @@ class Checker(
         self.env = previous_env
 
     def check(self, statements: list[p.Stmt]) -> list[Diagnostic]:
+        """Type check a sequence of statements and returns diagnostics
+
+        Args:
+            statements (list[p.Stmt]): the statements to evaluate and check
+
+        Returns:
+            list[Diagnostic]: the list of diagnostics (errors, warning, etc.)
+        """
         self.diagnostics = []
         for stmt in statements:
             stmt.accept(self)
@@ -94,12 +118,31 @@ class Checker(
         return self.diagnostics
 
     def look_up_variable(self, name: str, expr: p.Expr) -> Optional[Type]:
+        """Look up a variable in the environment it was declared
+
+        Args:
+            name (str): the name of the variable
+            expr (p.Expr): the variable expression, used to lookup the scope distance
+
+        Returns:
+            Optional[Type]: the type of the variable, or None if it was not found
+        """
         distance: Optional[int] = self.locals.get(expr)
         if distance is not None:
             return self.env.get_at(distance, name)
         return self.global_env.get(name)
 
     def parse_midas_import(self, expr: p.CallExpr) -> Optional[Path]:
+        """Parse a Midas import statement
+
+        The statement should be written as `midas.using("path/to/types.midas")`
+
+        Args:
+            expr (p.CallExpr): the import call expression
+
+        Returns:
+            Optional[Path]: the path to the imported file, or None if the expression is malformed
+        """
         match expr:
             case p.CallExpr(
                 callee=p.GetExpr(
@@ -114,6 +157,11 @@ class Checker(
         return None
 
     def import_midas(self, path: Path) -> None:
+        """Import Midas definitions from a path
+
+        Args:
+            path (Path): the import path
+        """
         self.logger.debug(f"Importing type definitions from {path}")
         path = (self.file_path.parent / path).resolve()
         lexer: MidasLexer = MidasLexer(path.read_text())
@@ -306,6 +354,20 @@ class Checker(
     def map_call_arguments(
         self, function: Function, call: p.CallExpr
     ) -> list[MappedArgument]:
+        """Map call arguments to function parameters as defined in its signature
+
+        This method maps positional-only, keyword-only and mixed parameter definitions
+        with the arguments passed at the call site
+
+        Any mismatched, missing or unexpected argument is reported as a diagnostic
+
+        Args:
+            function (Function): the function definition
+            call (p.CallExpr): the call expression
+
+        Returns:
+            list[MappedArgument]: the list of mapped arguments
+        """
         positional: list[tuple[p.Expr, Type]] = [
             (arg, self.evaluate(arg)) for arg in call.arguments
         ]

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,71 @@ 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

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)