docs: add docstrings for most files in checker module

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] = []
@@ -345,7 +385,7 @@ class CallDispatcher(Generic[E]):
 
         Returns:
             tuple[bool, list[MappedArgument]]: a boolean reporting whether
-            the call is valid and the list of mapped arguments
+                the call is valid and the list of mapped arguments
         """
         set_params: set[str] = set()
 
@@ -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"""
 

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/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}'"],