docs: add docstrings for type checkers

midas/checker/diagnostic.py

@@ -18,9 +18,6 @@ class Diagnostic:
 
     Holds a location, a diagnostic type and a message.
     Optionally bound to a file path
-
-    Returns:
-        _type_: _description_
     """
 
     file_path: Optional[str]
@@ -30,7 +27,7 @@ class Diagnostic:
 
     @property
     def location_str(self) -> str:
-        """The diagnostic type and location as a human readable string
+        """Get diagnostic type and location as a human readable string
 
         The location is formatted as "<Type> in <file> from L<start_line>:<start_col> to <end_line>:<end_col>",
         for example: "Error in /home/user/Desktop/script.py from L12:5 to L12:8"
@@ -39,7 +36,7 @@ class Diagnostic:
         If the location's end is not specified, the formulation "at L<start_line>:<start_col>" is used.
 
         Returns:
-            str: _description_
+            str: the formatted type and location string
         """
 
         start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"

midas/checker/dispatcher.py

@@ -26,10 +26,13 @@ class HasLocation(Protocol):
 E = TypeVar("E", bound=HasLocation)
 
 TypedExpr = tuple[E, Type]
+"""An expression and its type"""
 
 
 @dataclass(frozen=True, kw_only=True)
 class MappedArgument(Generic[E]):
+    """An argument passed in a call and the corresponding parameter"""
+
     arg_expr: E
     arg_type: Type
     parameter: Function.Parameter
@@ -37,11 +40,15 @@ class MappedArgument(Generic[E]):
 
 @dataclass(frozen=True, kw_only=True)
 class OverloadCandidate:
+    """An overloaded function call candidate with its mapped arguments"""
+
     function: Function
     mapped: list[MappedArgument]
 
 
 class CallError(StrEnum):
+    """Reason of a call error"""
+
     INVALID_ARGS = "Invalid arguments"
     NO_MATCHING_OVERLOAD = "No matching overload"
     IMPOSSIBLE_UNIFICATION = "Parameters unification failed"
@@ -50,16 +57,28 @@ class CallError(StrEnum):
 
 @dataclass(frozen=True, kw_only=True)
 class CallResult:
+    """The result of a function call
+
+    Holds a return type, an optional error reason and message
+    """
+
     error: Optional[CallError] = None
+    """The reason of the error, if there is one"""
+
     result: Type = UnknownType()
+    """The result type. `UnknownType()` if the call is invalid"""
+
     message: Optional[str] = None
+    """An optional error message"""
 
     @property
     def is_valid(self) -> bool:
+        """Whether the call is valid (i.e. no error)"""
         return self.error is None
 
     @property
     def error_message(self) -> str:
+        """A descriptive message for the error, if there is one"""
         if self.message is not None:
             return self.message
         if self.error is not None:
@@ -68,6 +87,15 @@ class CallResult:
 
 
 class CallDispatcher(Generic[E]):
+    """Helper class to handle dispatching calls and mapping arguments
+
+    This class is responsible for mapping call-site arguments to function
+    parameters, verifying the validity of calls and computing their
+    return types
+
+    :class:`CallDispatcher` is generic to handle AST nodes from both Midas and Python
+    """
+
     def __init__(self, types: TypesRegistry, reporter: FileReporter) -> None:
         self.types: TypesRegistry = types
         self.reporter: FileReporter = reporter
@@ -86,22 +114,21 @@ class CallDispatcher(Generic[E]):
     ) -> CallResult:
         """Get the result type of a function call
 
-        If the function has overloads, the function will try to resolve the
+        If the callee has overloads, this function will try to resolve the
         appropriate signature.
-        Argument types are matched to the defined parameters.
+        Argument types are matched with the defined parameters.
-        The function doesn't take the raw expression as a parameter to accommodate
+        This function doesn't take the raw expression as a parameter to
-        for desugared calls such as for operators.
+        accommodate for desugared calls such as for operators.
 
         Args:
             location (Location): the call location
             callee (Type): the called function
-            positional (list[TypedExpr]): the list positional arguments
+            positional (list[TypedExpr]): the list of positional arguments
             keywords (dict[str, TypedExpr]): the map of keyword arguments
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
-            Type: the return type of the call, or `None` if either
+            CallResult: the call result, either a type or an error
-            the call is invalid or no overload matched the arguments uniquely
         """
         match callee:
             case Function() as function:
@@ -179,6 +206,18 @@ class CallDispatcher(Generic[E]):
         positional: list[TypedExpr[E]],
         keywords: dict[str, TypedExpr[E]],
     ) -> Union[tuple[Function, None], tuple[None, CallError]]:
+        """Unwrap a type to get a callable `Function`
+
+        Args:
+            callee (Type): the called type
+            positional (list[TypedExpr[E]]): the list of positional arguments
+            keywords (dict[str, TypedExpr[E]]): the map of keyword arguments
+
+        Returns:
+            Union[tuple[Function, None], tuple[None, CallError]]: a tuple
+                containing the callable `Function` type, or `None` if it could
+                not be unwrapped, and an error, or `None` if there was none.
+        """
         match callee:
             case DerivedType(type=base):
                 return self._unwrap_function(base, positional, keywords)
@@ -246,8 +285,9 @@ class CallDispatcher(Generic[E]):
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
-            Optional[Function]: the resolved function signature if it can be
+            Union[tuple[Function, None], tuple[None, str]]: a tuple containing
-            determined unambiguously, or `None`.
+                the resolved function signature if it can be determined
+                unambiguously, or `None`, and an error message, or `None`
         """
         candidates: list[OverloadCandidate] = []
         errors: list[CallError] = []
@@ -464,8 +504,8 @@ class CallDispatcher(Generic[E]):
         of `mapped2`. If any of the parameter type in `mapped1` is not a subtype
         of the corresponding parameter in `mapped2`, `False` is returned.
 
-        This is used to check whether a given overload is
+        This is used to check whether a given overload is a more specific
-        a more specific function/ a subtype of another.
+        function / a subtype of another.
 
         Args:
             mapped1 (list[MappedArgument]): the first argument mappings (subtype)

midas/checker/evaluator.py

@@ -11,10 +11,18 @@ from midas.lexer.token import TokenType
 
 @dataclass(frozen=True, kw_only=True)
 class PartialPredicate(Predicate):
+    """A partially applied predicate"""
+
     scope: dict[str, Any]
+    """A dictionary of already applied parameters"""
 
 
 class Evaluator(m.Expr.Visitor[Any]):
+    """Helper class to evaluate an expression
+
+    This class is used to evaluate constraint types on literals at compile-time.
+    """
+
     def __init__(self, types: TypesRegistry, reporter: Optional[FileReporter] = None):
         self.types: TypesRegistry = types
         self.reporter: Optional[FileReporter] = reporter
@@ -22,16 +30,51 @@ class Evaluator(m.Expr.Visitor[Any]):
         self.scopes: list[dict[str, Any]] = [{}]
 
     def evaluate(self, expr: m.Expr) -> Any:
+        """Evaluate the given expression
+
+        Args:
+            expr (m.Expr): the expression to evaluate
+
+        Returns:
+            Any: the value of the expression
+        """
         value: Any = expr.accept(self)
         if self.reporter is not None:
             self.reporter.debug(expr.location, f"Value: {value}")
         return value
 
     def get_value(self, name: str) -> Any:
+        """Get the value of a variable in the current scope
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            KeyError: if the variable is not defined
+
+        Returns:
+            Any: the value of the variable
+        """
         scope: dict[str, Any] = self.scopes[-1]
         return scope[name]
 
     def set_value(self, name: str, value: Any, force_declare: bool = False):
+        """Set the value of a variable
+
+        If `force_declare` is `False`, this function first tries to find the
+        closest scope in which the variable is defined and assign the value in
+        that scope, if it can find one.
+
+        If `force_declare` is `True` or if the variable is not defined in any
+        scope, it is declare and assigned in the current scope
+
+        Args:
+            name (str): the name of the variable
+            value (Any): the value of the variable
+            force_declare (bool, optional): if `False` and the variable is
+                defined in a scope, the value is assigned in that scope (the
+                closest if there are multiple declarations). Defaults to False.
+        """
         if not force_declare:
             for scope in reversed(self.scopes):
                 if name in scope:
@@ -131,8 +174,21 @@ class Evaluator(m.Expr.Visitor[Any]):
         return self.get_value("_")
 
     def _evaluate_predicate(
-        self, predicate: Predicate, args: list[Any], kwargs: dict[str, Any]
+        self,
+        predicate: Predicate,
+        args: list[Any],
+        kwargs: dict[str, Any],
     ) -> Any:
+        """Evaluate a predicate function call
+
+        Args:
+            predicate (Predicate): the predicate to evaluate
+            args (list[Any]): a list of positional arguments
+            kwargs (dict[str, Any]): a map of keyword arguments
+
+        Returns:
+            Any: the value returned by the predicate call
+        """
         res: Any = None
         if isinstance(predicate, PartialPredicate):
             self.scopes.append(predicate.scope)
@@ -158,6 +214,16 @@ class Evaluator(m.Expr.Visitor[Any]):
         return res
 
     def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
+        """Map call arguments to a function's parameters and set their values in context
+
+        Each argument is mapped to a parameter of the function, then its value
+        is set in the context using :func:`set_value` with the parameter's name
+
+        Args:
+            function (Function): the called function
+            args (list[Any]): a list of positional arguments
+            kwargs (dict[str, Any]): a map of keyword arguments
+        """
         positional: list[Function.Parameter] = (
             function.params.pos + function.params.mixed
         )

midas/checker/midas.py

@@ -1,5 +1,4 @@
 import logging
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -33,23 +32,6 @@ from midas.lexer.token import Token
 from midas.parser.midas import MidasParser
 
 
-class ReturnException(Exception):
-    pass
-
-
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: m.Expr
-    type: Type
-    argument: Function.Parameter
-
-
-@dataclass(frozen=True, kw_only=True)
-class OverloadCandidate:
-    function: Function
-    mapped: list[MappedArgument]
-
-
 class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type]):
     """A resolver which evaluates Midas type definitions and build a registry"""
 
@@ -76,10 +58,21 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         self._preamble: Environment = Preamble(self.types)
 
     def set_reporter(self, reporter: FileReporter):
+        """Set the file reporter to use for diagnostics
+
+        Args:
+            reporter (FileReporter): the file reporter
+        """
         self.reporter = reporter
         self.dispatcher.set_reporter(reporter)
 
     def process(self, source: str, path: Optional[str]):
+        """Process some Midas source code
+
+        Args:
+            source (str): the Midas source code
+            path (Optional[str]): the path of the source file, if known
+        """
         reporter: FileReporter = self.reporter.for_file(path)
         self.set_reporter(reporter)
 
@@ -92,6 +85,14 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         self.resolve(stmts)
 
     def type_of(self, expr: m.Expr) -> Type:
+        """Compute the type of the given expression
+
+        Args:
+            expr (m.Expr): the expression to type
+
+        Returns:
+            Type: the type of the expression
+        """
         type: Type = expr.accept(self)
         return type
 
@@ -112,6 +113,21 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         return self.types.get_type(name)
 
     def get_variable(self, name: str) -> Type:
+        """Get the type of a variable
+
+        This function will first look into the current predicate's parameters if
+        we are in a predicate definition.
+        The the variable is looked up in the preamble (i.e. global environment)
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            NameError: if the variable cannot be found
+
+        Returns:
+            Type: the type of the variable
+        """
         if name in self._predicate_params:
             return self._predicate_params[name]
         predicate: Optional[Predicate] = self.types.lookup_predicate(name)
@@ -139,6 +155,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
                 self.types._types[name] = inferrer.infer(type)
 
     def assert_bool(self, expr: m.Expr):
+        """Check that the given expression is a subtype of `bool` or report an error
+
+        Args:
+            expr (m.Expr): the expression to check
+        """
         type: Type = self.type_of(expr)
         if not self.types.is_subtype(type, self._bool):
             self.reporter.error(expr.location, f"Must be a boolean but is {type}")
@@ -215,6 +236,16 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         )
 
     def _is_valid_predicate(self, body: Type) -> bool:
+        """Check whether the given type is valid as a predicate's body
+
+        Accepted types are either subtypes of `bool` or valid predicates
+
+        Args:
+            body (Type): the potential predicate body
+
+        Returns:
+            bool: `True` if `body` can be a predicate body, `False` otherwise
+        """
         match body:
             case Function(returns=returns):
                 return self._is_valid_predicate(returns)
@@ -240,7 +271,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         return self._visit_binary_expr(expr.location, expr.left, expr.right, method)
 
     def _visit_binary_expr(
-        self, location: Location, left_expr: m.Expr, right_expr: m.Expr, method: str
+        self,
+        location: Location,
+        left_expr: m.Expr,
+        right_expr: m.Expr,
+        method: str,
     ) -> Type:
         left: Type = self.type_of(left_expr)
         right: Type = self.type_of(right_expr)

midas/checker/preamble.py

@@ -23,6 +23,8 @@ class Param:
 
 
 class Preamble(Environment):
+    """The initial environment containing some of Python's builtin functions"""
+
     def __init__(self, types: TypesRegistry) -> None:
         super().__init__()
         self._types: TypesRegistry = types

midas/checker/python.py

@@ -1,6 +1,5 @@
 import ast
 import logging
-from dataclasses import dataclass
 from typing import Any, Optional
 
 import midas.ast.python as p
@@ -56,19 +55,6 @@ class UndefinedMethodException(Exception):
     pass
 
 
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: p.Expr
-    type: Type
-    argument: Function.Parameter
-
-
-@dataclass(frozen=True, kw_only=True)
-class OverloadCandidate:
-    function: Function
-    mapped: list[MappedArgument]
-
-
 class PythonTyper(
     p.Stmt.Visitor[None],
     p.Expr.Visitor[Type],
@@ -97,10 +83,24 @@ class PythonTyper(
         self.assertions: AssertionCollector = AssertionCollector()
 
     def set_reporter(self, reporter: FileReporter):
+        """Set the file reporter to use for diagnostics
+
+        Args:
+            reporter (FileReporter): the file reporter
+        """
         self.reporter = reporter
         self.dispatcher.set_reporter(self.reporter)
 
     def process(self, source: str, path: Optional[str]) -> TypedAST:
+        """Process some Python source code
+
+        Args:
+            source (str): the Python source code
+            path (Optional[str]): the path of the source file, if known
+
+        Returns:
+            TypedAST: all generated typechecking information
+        """
         reporter: FileReporter = self.reporter.for_file(path)
         self.set_reporter(reporter)
 
@@ -125,7 +125,7 @@ class PythonTyper(
         )
 
     def judge(self, expr: p.Expr, type: Type):
-        """Record a typing judgement
+        """Record a typing judgement for the given expression
 
         Args:
             expr (p.Expr): the judged expression
@@ -134,7 +134,7 @@ class PythonTyper(
         self.judgements.append((expr, type))
 
     def compute_type(self, expr: p.Expr) -> Type:
-        """Evaluate the type of an expression
+        """Evaluate the type of the given expression
 
         Args:
             expr (p.Expr): the expression to type
@@ -145,7 +145,7 @@ class PythonTyper(
         return expr.accept(self)
 
     def type_of(self, expr: p.Expr) -> Type:
-        """Evaluate the type of an expression and record the judgement
+        """Evaluate the type of the given expression and record the judgement
 
         Args:
             expr (p.Expr): the expression to evaluate
@@ -158,9 +158,22 @@ class PythonTyper(
         return type
 
     def resolve_type_expr(self, expr: p.MidasType) -> Type:
+        """Resolve the type of a type expression (annotation)
+
+        Args:
+            expr (p.MidasType): the type expression
+
+        Returns:
+            Type: the resolved type
+        """
         return expr.accept(self)
 
     def process_stmt(self, stmt: p.Stmt) -> None:
+        """Type check the given statement
+
+        Args:
+            stmt (p.Stmt): the statement to type-check
+        """
         stmt.accept(self)
 
     def process_block(self, block: list[p.Stmt], env: Environment) -> bool:
@@ -224,6 +237,24 @@ class PythonTyper(
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Evaluate a method call on an object
+
+        Calls to dataframes and columns types are delegated to the appropriate manager
+
+        Args:
+            location (Location): the location of the call
+            call_expr (p.Expr): the call expression
+            obj (TypedExpr): the object on which the method is called
+            method_name (str): the method name
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Raises:
+            UndefinedMethodException: if the method is not defined
+
+        Returns:
+            Type: the return type of the call
+        """
         unfolded: Type = unfold_type(obj[1])
         match unfolded:
             case DataFrameType():
@@ -283,6 +314,15 @@ class PythonTyper(
         return result.result
 
     def is_subtype(self, type1: Type, type2: Type) -> bool:
+        """Check whether `type1` is a subtype of `type2`
+
+        Args:
+            type1 (Type): the potential "subtype"
+            type2 (Type): the potential "supertype"
+
+        Returns:
+            bool: whether `type1` is a subtype of `type2`
+        """
         return self.types.is_subtype(type1, type2)
 
     def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
@@ -408,6 +448,15 @@ class PythonTyper(
             self._assign(stmt.location, target, value_type)
 
     def _assign(self, location: Location, target: p.Expr, value_type: Type):
+        """Handle an assignment to the given target
+
+        Delegate to the appropriate method according to the target type
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.Expr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         match target:
             case p.VariableExpr():
                 self._assign_var(location, target, value_type)
@@ -429,6 +478,13 @@ class PythonTyper(
                     )
 
     def _assign_var(self, location: Location, target: p.VariableExpr, value_type: Type):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         name: str = target.name
         var_type: Optional[Type] = self.look_up_variable(name, target)
 
@@ -447,6 +503,13 @@ class PythonTyper(
     def _assign_attr(
         self, location: Location, object: p.Expr, name: str, value_type: Type
     ):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         object_type: Type = self.type_of(object)
         member: Optional[Type] = self.types.lookup_member(object_type, name)
         if member is None:
@@ -466,6 +529,13 @@ class PythonTyper(
         index: p.Expr,
         value_type: Type,
     ):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         var_type: Type = self.type_of(var)
         unfolded_type: Type = unfold_type(var_type)
         # TODO: what happens if type is an alias of a dataframe type
@@ -887,6 +957,15 @@ class PythonTyper(
         )
 
     def _get_iterator_type(self, expr: p.Expr, type: Type) -> Optional[Type]:
+        """Get the item type of an iterator type
+
+        Args:
+            expr (p.Expr): the iterator expression
+            type (Type): the iterator type
+
+        Returns:
+            Optional[Type]: the item type, or `None` if it cannot be determined
+        """
         # TODO: lookup __iter__
         getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__")
         if getitem is None:
@@ -906,6 +985,16 @@ class PythonTyper(
         return result.result
 
     def define_typevar(self, call: p.CallExpr) -> Optional[TypeVar]:
+        """Define a type variable from a call to `typing.TypeVar`
+
+        Args:
+            call (p.CallExpr): the call to `typing.TypeVar`
+
+        Returns:
+            Optional[TypeVar]: the define type variable, or `None` if the call
+                is invalid
+        """
+
         def is_kw_true(name: str) -> bool:
             match call.keywords.get(name):
                 case p.LiteralExpr(value=True):
@@ -948,6 +1037,19 @@ class PythonTyper(
                 return None
 
     def _parse_type_from_expr(self, expr: p.Expr) -> p.MidasType:
+        """Parse a type expression from a raw expression
+
+        This is useful for expressions inside a `TypeVar`'s `bound` parameter
+
+        Args:
+            expr (p.Expr): the expression to parse
+
+        Raises:
+            NotImplementedError: if the expression is not supported
+
+        Returns:
+            p.MidasType: the parsed type node
+        """
         location: Location = expr.location
         parser = PythonParser()
         match expr:
@@ -960,6 +1062,16 @@ class PythonTyper(
                 raise NotImplementedError
 
     def _get_literal(self, expr: p.Expr) -> tuple[bool, Any]:
+        """Get the literal value of a literal-like expression
+
+        Args:
+            expr (p.Expr): the expression
+
+        Returns:
+            tuple[bool, Any]: a tuple containing a boolean indicating whether
+                the given expression is literal-like, and the literal value (or
+                `None` if the first value is `False`)
+        """
         match expr:
             case p.LiteralExpr(value=value):
                 return True, value
@@ -1016,6 +1128,17 @@ class PythonTyper(
     def _evaluate_cast_statically(
         self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any
     ) -> bool:
+        """Evaluate the given cast expression statically
+
+        Args:
+            expr (p.CastExpr): the cast expression
+            subject_type (Type): the subject type being casted
+            target_type (Type): the target type to which the expression is casted
+            lit_value (Any): the literal value of the expression
+
+        Returns:
+            bool: whether the cast expression could be evaluated successfully
+        """
         match target_type:
             case TopType():
                 return True

midas/checker/registry.py

@@ -29,11 +29,15 @@ from midas.checker.types import (
 
 @dataclass
 class Member:
+    """A member of a type (property or method)"""
+
     kind: MemberKind
     type: Type
 
 
 class TypesRegistry:
+    """A registry of types, type members and predicates"""
+
     def __init__(self) -> None:
         self.logger: logging.Logger = logging.getLogger("TypesRegistry")
         self._types: dict[str, Type] = {}
@@ -81,6 +85,25 @@ class TypesRegistry:
         member_type: Type,
         kind: MemberKind,
     ):
+        """Define a member on a type
+
+        If the member is a method and a member with the same name is already
+        defined on the given type, the two are combined into an :class:`OverloadedFunction`.
+
+        If the member is a property and a member with the same name is already
+        defined on the given type, the new definition is dropped and an error
+        is reported.
+
+        In any case, if a member with the same name but a different kind is
+        already defined on the given type, the new definition is dropped and
+        an error is reported.
+
+        Args:
+            type_name (str): the name of the type on which the member is defined
+            member_name (str): the name of the new member
+            member_type (Type): the type of the new member
+            kind (MemberKind): the kind of member to define (property or method)
+        """
         members: dict[str, Member] = self._members.setdefault(type_name, {})
         if member_name in members:
             current: Member = members[member_name]
@@ -109,11 +132,29 @@ class TypesRegistry:
             members[member_name] = Member(kind=kind, type=member_type)
 
     def define_predicate(self, name: str, predicate: Predicate):
+        """Define a predicate
+
+        Args:
+            name (str): the name of the new predicate
+            predicate (Predicate): the predicate to define
+
+        Raises:
+            ValueError: if a predicate with the same name is already defined
+        """
         if name in self._predicates:
             raise ValueError(f"Predicate {name} already defined")
         self._predicates[name] = predicate
 
     def is_builtin_subtype(self, name1: str, name2: str) -> bool:
+        """Check whether a type is a subtype of another base on builtin subtype rules
+
+        Args:
+            name1 (str): the name of the potential subtype
+            name2 (str): the name of the potential supertype
+
+        Returns:
+            bool: _description_
+        """
         subtypes: set[str] = BUILTIN_SUBTYPES.get(name2, set())
         if name1 in subtypes:
             return True
@@ -218,6 +259,15 @@ class TypesRegistry:
         return False
 
     def are_equivalent(self, type1: Type, type2: Type) -> bool:
+        """Check whether two types are equivalent (T <: S and S <: T)
+
+        Args:
+            type1 (Type): the first type
+            type2 (Type): the second type
+
+        Returns:
+            bool: whether `type1` is a subtype and a supertype of `type2`
+        """
         return self.is_subtype(type1, type2) and self.is_subtype(type2, type1)
 
     # TODO: verify the logic in here
@@ -334,6 +384,18 @@ class TypesRegistry:
         return True
 
     def apply_generic(self, type: Type, args: list[Type]) -> Type:
+        """Instantiate a generic type with the given type arguments
+
+        Args:
+            type (Type): the generic
+            args (list[Type]): the type arguments
+
+        Raises:
+            ValueError: if the arguments are invalid (wrong number, bound violation)
+
+        Returns:
+            Type: the applied generic type
+        """
         match type:
             case DerivedType(name=name, type=base):
                 return DerivedType(name=name, type=self.apply_generic(base, args))
@@ -399,6 +461,19 @@ class TypesRegistry:
         return [types[i] for i in keep]
 
     def lookup_member(self, type: Type, member_name: str) -> Optional[Type]:
+        """Lookup a member by name on a given type
+
+        This function first looks up directly on the specified type, then
+        recurse through supertypes until it finds the member or reaches
+        the root type
+
+        Args:
+            type (Type): the type on which to lookup the member
+            member_name (str): the member's name
+
+        Returns:
+            Optional[Type]: the member's type, or `None` if it is not defined
+        """
         match type:
             case BaseType(name=name):
                 if name in self._members:
@@ -459,18 +534,54 @@ class TypesRegistry:
                 return None
 
     def lookup_predicate(self, name: str) -> Optional[Predicate]:
+        """Lookup a predicate by name
+
+        Args:
+            name (str): the name of the predicate
+
+        Returns:
+            Optional[Predicate]: the predicate, or `None` if is not defined
+        """
         return self._predicates.get(name)
 
     def _by_name_or_type(self, name_or_type: str | Type) -> Type:
+        """Get a type by name or return it as is
+
+        If `name_or_type` is a string, the associated type is looked up and returned.
+        Otherwise, the type is returned as is.
+
+        Args:
+            name_or_type (str | Type): the type or type's name
+
+        Returns:
+            Type: the type
+        """
         if isinstance(name_or_type, str):
             return self.get_type(name_or_type)
         return name_or_type
 
     def list_of(self, item_type: str | Type) -> Type:
+        """Helper method to type a list of a given item type
+
+        Args:
+            item_type (str | Type): the item type
+
+        Returns:
+            Type: the list type
+        """
         list_ = self.get_type("list")
         return self.apply_generic(list_, [self._by_name_or_type(item_type)])
 
     def tuple_of(self, *item_types: str | Type) -> Type:
+        """Helper method to type a tuple of given item types
+
+        Args:
+            item_type (str | Type): the item types
+
+        Returns:
+            Type: the tuple type
+        """
+
         tuple_ = self.get_type("tuple")
         return self.apply_generic(
             tuple_,
@@ -478,6 +589,15 @@ class TypesRegistry:
         )
 
     def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type:
+        """Helper method to type a dict of given key and value types
+
+        Args:
+            key_type (str | Type): the key type
+            value_type (str | Type): the value type
+
+        Returns:
+            Type: the dict type
+        """
         dict_ = self.get_type("dict")
         return self.apply_generic(
             dict_,

midas/checker/reporter.py

@@ -7,6 +7,8 @@ from midas.checker.diagnostic import Diagnostic, DiagnosticType
 
 
 class Reporter:
+    """Helper class to store diagnostics"""
+
     def __init__(self):
         self.diagnostics: list[Diagnostic] = []
 
@@ -17,6 +19,14 @@ class Reporter:
         location: Location,
         message: str,
     ):
+        """Create and record a diagnostic
+
+        Args:
+            path (Optional[str]): the path linked to this diagnostic
+            type (DiagnosticType): the type of diagnostic
+            location (Location): the location if the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.diagnostics.append(
             Diagnostic(
                 file_path=path,
@@ -27,21 +37,52 @@ class Reporter:
         )
 
     def for_file(self, path: Optional[str]) -> FileReporter:
+        """Create a new file reporter for the given path using this reporter
+
+        Args:
+            path (Optional[str]): the path for the new file reporter
+
+        Returns:
+            FileReporter: the new file reporter, linked to this reporter
+        """
         return FileReporter(self, path)
 
 
 class FileReporter:
+    """Helper class to manage diagnostics for a file"""
+
     def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None:
         self.base_reporter: Reporter = base_reporter
         self.path: Optional[str] = path
 
     def for_file(self, path: Optional[str]) -> FileReporter:
+        """Create a new file reporter for the given path with the same base reporter
+
+        Args:
+            path (Optional[str]): the path for the new file reporter
+
+        Returns:
+            FileReporter: the file reporter
+        """
         return FileReporter(self.base_reporter, path)
 
     def report(self, type: DiagnosticType, location: Location, message: str):
+        """Report a diagnostic to the base reporter
+
+        Args:
+            type (DiagnosticType): the type of diagnostic
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.base_reporter.report(self.path, type, location, message)
 
     def error(self, location: Location, message: str):
+        """Report an error diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.report(
             type=DiagnosticType.ERROR,
             location=location,
@@ -49,6 +90,12 @@ class FileReporter:
         )
 
     def warning(self, location: Location, message: str):
+        """Report a warning diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.report(
             type=DiagnosticType.WARNING,
             location=location,
@@ -56,6 +103,12 @@ class FileReporter:
         )
 
     def info(self, location: Location, message: str):
+        """Report an info diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.report(
             type=DiagnosticType.INFO,
             location=location,
@@ -63,6 +116,12 @@ class FileReporter:
         )
 
     def debug(self, location: Location, message: str):
+        """Report a debug diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
         self.report(
             type=DiagnosticType.DEBUG,
             location=location,

midas/checker/resolver.py

@@ -78,6 +78,14 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 return
 
     def is_defined(self, name: str) -> bool:
+        """Check whether the given variable is defined in any scope
+
+        Args:
+            name (str): the name of the variable
+
+        Returns:
+            bool: `True` if the variable is defined in a scope, `False` otherwise
+        """
         for scope in self.scopes:
             if name in scope:
                 return True

midas/checker/unifier.py

@@ -19,6 +19,14 @@ class UnificationError(Exception): ...
 
 
 class Unifier:
+    """
+    Helper class to unify generic types in concrete usages
+
+    This can be used for example when a generic function is called with concrete
+    arguments, at which point the type parameters of the function signature
+    should be resolvable
+    """
+
     def __init__(self, types: TypesRegistry) -> None:
         self.types: TypesRegistry = types
         self.logger: logging.Logger = logging.getLogger("Unifier")
@@ -29,6 +37,16 @@ class Unifier:
         positional: list[Type],
         keywords: dict[str, Type],
     ) -> Optional[Type]:
+        """Try and unify a generic function call given concrete arguments
+
+        Args:
+            type (GenericType): the generic function type
+            positional (list[Type]): the list of positional arguments
+            keywords (dict[str, Type]): the map of keyword arguments
+
+        Returns:
+            Optional[Type]: the concrete function type if unifiable, or `None`
+        """
         concrete_func: Function = Function(
             params=ParamSpec(
                 pos=[
@@ -60,6 +78,18 @@ class Unifier:
         concrete: Type,
         match_return: bool = True,
     ) -> Optional[Type]:
+        """Unify a generic type's parameters given a concrete usage
+
+        Args:
+            template (GenericType): the generic type
+            concrete (Type): a concrete usage
+            match_return (bool, optional): if `template` is a function type,
+                whether its return type must be matched (see :func:`match`).
+                Defaults to True.
+
+        Returns:
+            Optional[Type]: the concrete type if unifiable, or `None`
+        """
         substitutions: dict[str, Type]
         try:
             substitutions = self.match(template.body, concrete, match_return)
@@ -81,6 +111,22 @@ class Unifier:
         concrete: Type,
         match_return: bool = True,
     ) -> dict[str, Type]:
+        """Match a generic type with a concrete usage, recording parameter substitutions
+
+        Args:
+            template (Type): the generic type
+            concrete (Type): a concrete usage
+            match_return (bool, optional): if `template` and `concrete` are both
+                :class:`Function`, whether their return types are also matched.
+                Defaults to True.
+
+        Raises:
+            UnificationError: if there is a conflict in parameter substitutions
+
+        Returns:
+            dict[str, Type]: the parameter substitutions which,
+                applied to `template`, yield `concrete`
+        """
         # TODO: if concrete is Generic, record bound TypeVar. Then when merging
         # substitutions, check that the constraint is respected
         match (template, concrete):
@@ -150,6 +196,18 @@ class Unifier:
                 return {}
 
     def merge(self, subs1: dict[str, Type], subs2: dict[str, Type]) -> dict[str, Type]:
+        """Merge two maps of substitutions and raise an error if incompatible
+
+        Args:
+            subs1 (dict[str, Type]): the first substitutions
+            subs2 (dict[str, Type]): the second substitutions
+
+        Raises:
+            UnificationError: if there is a conflict between the two maps
+
+        Returns:
+            dict[str, Type]: the merged map of substitutions
+        """
         merged: dict[str, Type] = subs1.copy()
 
         for k, v in subs2.items():
@@ -164,6 +222,15 @@ class Unifier:
     def map_params(
         self, func1: Function, func2: Function
     ) -> list[tuple[Function.Parameter, Function.Parameter]]:
+        """Map parameters of two functions
+
+        Args:
+            func1 (Function): the first function
+            func2 (Function): the second function
+
+        Returns:
+            list[tuple[Function.Parameter, Function.Parameter]]: the list of parameter pairs
+        """
         pos1: list[Function.Parameter] = func1.params.pos
         mixed1: list[Function.Parameter] = func1.params.mixed
         kw1: list[Function.Parameter] = func1.params.kw

midas/checker/variance.py

@@ -16,14 +16,27 @@ Polarity = Literal[-1, 0, 1]
 
 
 class Tracker:
+    """Helper class to track the polarity of type parameter references and computer their variance"""
+
     def __init__(self, vars: list[TypeVar]) -> None:
         self.vars: list[TypeVar] = vars
         self.refs: dict[str, set[Polarity]] = {var.name: set() for var in self.vars}
 
     def record(self, var: TypeVar, polarity: Polarity):
+        """Record a polarity of the given type parameter
+
+        Args:
+            var (TypeVar): the type parameter
+            polarity (Polarity): the polarity
+        """
         self.refs[var.name].add(polarity)
 
     def get_updated_vars(self) -> list[TypeVar]:
+        """Get a list of the tracked type variables with their recorded variance
+
+        Returns:
+            list[TypeVar]: the list of update type parameters
+        """
         return [
             TypeVar(
                 name=var.name, bound=var.bound, variance=self.get_variance(var.name)
@@ -32,6 +45,18 @@ class Tracker:
         ]
 
     def get_variance(self, name: str) -> Variance:
+        """Get the variance of a type parameter
+
+        If the type parameter is only referenced in positive positions, it is
+        covariant. If it is only referenced in negative positions, it is
+        contravariant. Otherwise, it is invariant
+
+        Args:
+            name (str): the name of the type parameter
+
+        Returns:
+            Variance: the variance of the type parameter
+        """
         refs: set[Polarity] = self.refs[name]
         if refs == {-1}:
             return Variance.CONTRAVARIANT
@@ -46,11 +71,22 @@ class Tracker:
 
 
 class VarianceInferrer:
+    """Helper class to compute type parameter variance"""
+
     def __init__(self, types: TypesRegistry) -> None:
         self.types: TypesRegistry = types
         self.tracker: Tracker = Tracker([])
 
     def infer(self, type: GenericType) -> GenericType:
+        """Infer the variance of a generic type's parameters
+
+        Args:
+            type (GenericType): the generic type
+
+        Returns:
+            GenericType: a new generic type with its parameters updated with
+                their inferred variance
+        """
         self.tracker = Tracker(type.params)
 
         self.walk(type.body, 1, type.name)
@@ -71,6 +107,22 @@ class VarianceInferrer:
         base_name: str,
         path: Optional[list[str]] = None,
     ):
+        """Walk the type nodes and record variance
+
+        This function recurses into type substructures (e.g. function parameters,
+        overloads, constraint type bases, etc.)
+
+        When recursing, the polarity is flipped for consumer positions (e.g. function
+        parameters) or kept the same for producer positions (e.g. return type)
+
+        Args:
+            type (Type): the type to visit
+            polarity (Polarity): the current polarity
+            base_name (str): the root generic type name (used to detect and
+                handle cyclic references)
+            path (Optional[list[str]], optional): the path to reach the current
+                type from the root generic type (used for debugging). Defaults to None.
+        """
         if path is None:
             path = []
 
@@ -109,10 +161,10 @@ class VarianceInferrer:
                     Variance.COVARIANT: 1,
                     Variance.CONTRAVARIANT: -1,
                 }
-                for param, param in zip(args, params):
+                for arg, param in zip(args, params):
                     param_polarity: Polarity = polarities[param.variance]
                     self.walk(
-                        param,
+                        arg,
                         cast(Polarity, polarity * param_polarity),
                         base_name,
                         path + [f"applied:'{name}'"],