docs: add docstrings to types

midas/checker/frames/column_groupby_methods.py

@@ -22,6 +22,8 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
+    """A column group-by method call, implements :class:`utils.MethodCall`"""
+
     location: Location
     call_expr: p.Expr
     groupby: ColumnGroupBy
@@ -35,6 +37,8 @@ class Call:
 
 
 class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
+    """The method registry for column group-by types"""
+
     NAMED_ARGS: dict[str, str] = {
         "numeric_only": "bool",
         "skipna": "bool",
@@ -49,6 +53,21 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
         *,
         preserve_inner_type: bool = False,
     ) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            params (list[str | tuple[str, str, bool], optional): a list of extra
+                mixed parameters. The list can contain strings to include
+                parameters predefined in `NAMED_ARGS`, or tuples containing the
+                parameter's name, type and required flag. Defaults to [].
+            preserve_inner_type (bool, optional): If `True`, the result type
+                will preserve the column's inner type (e.g. for `min`/`max`),
+                otherwise the inner type is widened to `TopType`. Defaults to False.
+
+        Returns:
+            Type: the result type
+        """
         real_params: list[Function.Parameter] = []
         for i, param in enumerate(params):
             match param:

midas/checker/frames/column_manager.py

@@ -15,6 +15,8 @@ if TYPE_CHECKING:
 
 
 class ColumnManager:
+    """Helper class to handle methods and subscripts on column types"""
+
     def __init__(self, typer: PythonTyper) -> None:
         self.typer: PythonTyper = typer
         self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer)
@@ -32,6 +34,20 @@ class ColumnManager:
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Compute the result type of a column's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            column (ColumnType): the column type
+            column_expr (p.Expr): the column expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
         call: Call = Call(
             location=location,
             call_expr=call_expr,
@@ -52,6 +68,20 @@ class ColumnManager:
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Compute the result type of a column group-by's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            groupby (ColumnGroupBy): the column group-by object
+            groupby_expr (p.Expr): the column group-by expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
         call: GroupByCall = GroupByCall(
             location=location,
             call_expr=call_expr,
@@ -63,6 +93,15 @@ class ColumnManager:
         return self.groupby_method_resolver.call(method, call)
 
     def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]:
+        """Get the type of a column's attribute
+
+        Args:
+            column (ColumnType): the column type
+            name (str): the attribute's name
+
+        Returns:
+            Optional[Type]: the attribute's type, or `None` if it doesn't exist
+        """
         types: TypesRegistry = self.typer.types
         match name:
             case "ndim" | "size":

midas/checker/frames/column_methods.py

@@ -27,6 +27,8 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
+    """A column method call, implements :class:`utils.MethodCall`"""
+
     location: Location
     call_expr: p.Expr
     column: ColumnType
@@ -40,6 +42,8 @@ class Call:
 
 
 class ColumnMethodRegistry(MethodRegistry[Call]):
+    """The method registry for column types"""
+
     def _element_binary_op(self, call: Call, method: str) -> ColumnType:
         """Compute the result of an element-wise binary operation
 
@@ -75,6 +79,18 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         return new_column
 
     def _element_wise(self, call: Call, method: str) -> Type:
+        """Compute the result of an element-wise method call
+
+        If the call is valid, this method also generates an assertion to check
+        that both operands have the same length at runtime
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
         # TODO: support add with scalar
 
         # Build signature with new column type and generic operand
@@ -170,6 +186,19 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         *,
         preserve_inner_type: bool = False,
     ) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            kwargs (list[Function.Parameter], optional): a list of extra
+                keyword-only parameters. Defaults to [].
+            preserve_inner_type (bool, optional): If `True`, the result type
+                will preserve the column's inner type (e.g. for `min`/`max`),
+                otherwise the inner type is widened to `TopType`. Defaults to False.
+
+        Returns:
+            Type: the result type
+        """
         signature = Function(
             params=ParamSpec(
                 kw=[
@@ -344,6 +373,14 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         return result.result
 
     def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr):
+        """Generate an assertion to check that two columns have the same length
+
+        Args:
+            call_expr (p.Expr): the call expression, to insert the assertion
+                at the right place
+            column1 (p.Expr): the first column expression
+            column2 (p.Expr): the second column expression
+        """
         func_name: str = "__midas_column_same_length__"
 
         # Efficiently compute length

midas/checker/frames/frame_groupby_methods.py

@@ -21,6 +21,8 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
+    """A frame group-by method call, implements :class:`utils.MethodCall`"""
+
     location: Location
     call_expr: p.Expr
     groupby: FrameGroupBy
@@ -34,14 +36,18 @@ class Call:
 
 
 class FrameGroupByMethodRegistry(MethodRegistry[Call]):
-    NAMED_ARGS: dict[str, str] = {
+    """The method registry for frame group-by types"""
-        "numeric_only": "bool",
-        "skipna": "bool",
-        "engine": "str",
-        "engine_kwargs": "dict",
-    }
 
     def _aggregate(self, call: Call, method: str) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
         new_columns: list[DataFrameType.Column] = []
 
         for column in call.groupby.frame.columns:

midas/checker/frames/frame_manager.py

@@ -24,10 +24,20 @@ if TYPE_CHECKING:
 
 
 def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]:
+    """Check whether the given list only contains literal expressions
+
+    Args:
+        exprs (list[p.Expr]): the list to check
+
+    Returns:
+        TypeGuard[list[p.LiteralExpr]]: whether `exprs` only contains literal expressions
+    """
     return all(isinstance(expr, p.LiteralExpr) for expr in exprs)
 
 
 class FrameManager:
+    """Helper class to handle methods and subscripts on frame types"""
+
     def __init__(self, typer: PythonTyper) -> None:
         self.typer: PythonTyper = typer
         self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer)
@@ -43,6 +53,18 @@ class FrameManager:
         index: p.Expr,
         value_type: Type,
     ) -> Type:
+        """Compute the new frame type after assigning a value to an index
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the assignment's location
+            frame (DataFrameType): the frame type
+            index (p.Expr): the index expression
+            value_type (Type): the assigned value
+
+        Returns:
+            Type: the resulting frame type
+        """
         match index:
             case p.LiteralExpr(value=str() as name):
                 return self.assign_column(reporter, location, frame, name, value_type)
@@ -93,6 +115,18 @@ class FrameManager:
         name: str,
         type: Type,
     ) -> Type:
+        """Compute the new frame type after assigning a single value to a column
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the assignment's location
+            frame (DataFrameType): the frame type
+            name (str): the column name
+            type (Type): the assigned value type
+
+        Returns:
+            Type: the resulting frame type
+        """
         if not isinstance(type, ColumnType):
             reporter.error(
                 location,
@@ -108,6 +142,17 @@ class FrameManager:
         frame: DataFrameType,
         index: p.Expr,
     ) -> Type:
+        """Compute the type of a subscript access
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the subscript's location
+            frame (DataFrameType): the frame type
+            index (p.Expr): the index expression
+
+        Returns:
+            Type: the resulting type
+        """
         match index:
             case p.LiteralExpr(value=str() as name):
                 column: Optional[ColumnType] = FrameManager._get_column(frame, name)
@@ -142,6 +187,17 @@ class FrameManager:
         groupby: FrameGroupBy,
         index: p.Expr,
     ) -> Type:
+        """Compute the type of a subscript access on a frame group-by object
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the subscript's location
+            groupby (FrameGroupBy): the group-by object
+            index (p.Expr): the index expression
+
+        Returns:
+            Type: the resulting type
+        """
         result: Type = self.get(reporter, location, groupby.frame, index)
         match result:
             case ColumnType():
@@ -159,6 +215,16 @@ class FrameManager:
     def _set_column(
         cls, frame: DataFrameType, name: str, column: ColumnType
     ) -> DataFrameType:
+        """Set a frame's column to the given type
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the column's name
+            column (ColumnType): the new column's type
+
+        Returns:
+            DataFrameType: the new frame type
+        """
         new_columns: list[DataFrameType.Column] = []
         index: int = len(frame.columns)
         replace: bool = False
@@ -185,12 +251,31 @@ class FrameManager:
     def _set_columns(
         cls, frame: DataFrameType, names: list[str], columns: list[ColumnType]
     ) -> DataFrameType:
+        """Set multiple columns of a frame to the given types
+
+        Args:
+            frame (DataFrameType): the frame type
+            names (list[str]): the column names
+            columns (list[ColumnType]): the new column types
+
+        Returns:
+            DataFrameType: the new frame type
+        """
         for name, col in zip(names, columns):
             frame = cls._set_column(frame, name, col)
         return frame
 
     @classmethod
     def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]:
+        """Get a column's type by name
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the column's name
+
+        Returns:
+            Optional[ColumnType]: the column's type, or `None` if it doesn't exist
+        """
         for col in frame.columns:
             if col.name == name:
                 return col.type
@@ -200,6 +285,15 @@ class FrameManager:
     def _get_columns(
         cls, frame: DataFrameType, names: list[str]
     ) -> list[Optional[ColumnType]]:
+        """Get multiple column types by name
+
+        Args:
+            frame (DataFrameType): the frame type
+            names (list[str]): the column names
+
+        Returns:
+            list[Optional[ColumnType]]: the column types (see :func:`_get_column`)
+        """
         return [cls._get_column(frame, name) for name in names]
 
     def call(
@@ -212,6 +306,20 @@ class FrameManager:
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Compute the result type of a frame's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            frame (DataFrameType): the frame type
+            frame_expr (p.Expr): the frame expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
         call: Call = Call(
             location=location,
             call_expr=call_expr,
@@ -232,6 +340,20 @@ class FrameManager:
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Compute the result type of a frame group-by's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            groupby (FrameGroupBy): the frame group-by object
+            groupby_expr (p.Expr): the frame group-by expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
         call: GroupByCall = GroupByCall(
             location=location,
             call_expr=call_expr,
@@ -243,6 +365,15 @@ class FrameManager:
         return self.groupby_method_resolver.call(method, call)
 
     def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]:
+        """Get the type of a frame's attribute
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the attribute's name
+
+        Returns:
+            Optional[Type]: the attribute's type, or `None` if it doesn't exist
+        """
         types: TypesRegistry = self.typer.types
         match name:
             case "ndim" | "size":

midas/checker/frames/frame_methods.py

@@ -27,6 +27,8 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
+    """A frame method call, implements :class:`utils.MethodCall`"""
+
     location: Location
     call_expr: p.Expr
     frame: DataFrameType
@@ -40,6 +42,8 @@ class Call:
 
 
 class FrameMethodRegistry(MethodRegistry[Call]):
+    """The method registry for frame types"""
+
     def _get_method_result(
         self,
         call: Call,
@@ -148,6 +152,18 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         return DataFrameType(columns=new_columns)
 
     def _element_wise(self, call: Call, method: str) -> Type:
+        """Compute the result of an element-wise method call
+
+        If the call is valid, this method also generates an assertion to check
+        that both operands have the same length at runtime
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
         # TODO: support scalar, sequence, Series, dict operand
         # Build signature with new schema and generic operand
         signature = Function(
@@ -231,6 +247,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         return self._element_wise(call, "__eq__")
 
     def _aggregate(self, call: Call, kwargs: list[Function.Parameter] = []) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            kwargs (list[Function.Parameter], optional): a list of extra
+                keyword-only parameters. Defaults to [].
+
+        Returns:
+            Type: the result type
+        """
         with_axis = Function(
             params=ParamSpec(
                 kw=[
@@ -425,6 +451,14 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         return result.result
 
     def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr):
+        """Generate an assertion to check that two frames have the same length
+
+        Args:
+            call_expr (p.Expr): the call expression, to insert the assertion
+                at the right place
+            frame1 (p.Expr): the first frame expression
+            frame2 (p.Expr): the second frame expression
+        """
         func_name: str = "__midas_frame_same_length__"
 
         # Efficiently compute length

midas/checker/frames/utils.py

@@ -24,6 +24,12 @@ if TYPE_CHECKING:
 
 
 class _MethodRegistryMeta(type):
+    """Meta-class for :class:`MethodRegistry`
+
+    Collects methods marked with the :func:`method` decorator into a dictionary
+    named `_methods` on the class itself
+    """
+
     _methods: dict[str, Callable[..., Type]] = {}
 
     def __new__(
@@ -42,6 +48,11 @@ class _MethodRegistryMeta(type):
 
 
 class MethodCall(Protocol):
+    """A method call object
+
+    Must have at least `location`, `call_expr` and `subject` properties
+    """
+
     @property
     def location(self) -> Location: ...
 
@@ -56,6 +67,8 @@ T = TypeVar("T", bound=MethodCall)
 
 
 class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
+    """A registry of methods"""
+
     def __init__(self, typer: PythonTyper) -> None:
         self.typer: PythonTyper = typer
 
@@ -76,6 +89,15 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
         return self.typer.assertions
 
     def call(self, method: str, call: T) -> Type:
+        """Compute the result type of a call to the given method
+
+        Args:
+            method (str): the method's name
+            call (T): the call
+
+        Returns:
+            Type: the result type
+        """
         func: Optional[Callable[[Self, T], Type]] = self._methods.get(method)
         if func is None:
             self.reporter.warning(
@@ -90,6 +112,13 @@ Method = Callable[[_Self, T], Type]
 
 
 def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]:
+    """Simple decorator to mark a method as part of the registry
+
+    Args:
+        names (str): names by which the method can be called. If left empty, the
+            Python method's name will be used
+    """
+
     def wrapper(func: Method[_Self, T]) -> Method[_Self, T]:
         names_: tuple[str, ...] = names
         if len(names_) == 0:

midas/checker/types.py

@@ -10,12 +10,16 @@ from midas.ast.printer import MidasPrinter
 
 @dataclass(frozen=True, kw_only=True)
 class TopType:
+    """The top type (`Any`)"""
+
     def __str__(self) -> str:
         return "Any"
 
 
 @dataclass(frozen=True, kw_only=True)
 class BaseType:
+    """A base / builtin type"""
+
     name: str
 
     def __str__(self) -> str:
@@ -24,6 +28,8 @@ class BaseType:
 
 @dataclass(frozen=True, kw_only=True)
 class DerivedType:
+    """A derived type, i.e. a named subtype of another type"""
+
     name: str
     type: Type
 
@@ -33,18 +39,24 @@ class DerivedType:
 
 @dataclass(frozen=True, kw_only=True)
 class UnknownType:
+    """An unknown type"""
+
     def __str__(self) -> str:
         return "<Unknown>"
 
 
 @dataclass(frozen=True, kw_only=True)
 class UnitType:
+    """The unit type (`None`)"""
+
     def __str__(self) -> str:
         return "None"
 
 
 @dataclass(frozen=True, kw_only=True)
 class Function:
+    """A function type"""
+
     params: ParamSpec
     returns: Type
 
@@ -65,6 +77,8 @@ class Function:
 
 @dataclass(frozen=True, kw_only=True)
 class ParamSpec:
+    """A function's parameter spec"""
+
     pos: list[Function.Parameter] = field(default_factory=list)
     mixed: list[Function.Parameter] = field(default_factory=list)
     kw: list[Function.Parameter] = field(default_factory=list)
@@ -87,6 +101,8 @@ class ParamSpec:
 
 @dataclass(frozen=True, kw_only=True)
 class OverloadedFunction:
+    """A list of method overloads"""
+
     overloads: list[Type]
 
     def __str__(self) -> str:
@@ -95,6 +111,8 @@ class OverloadedFunction:
 
 @dataclass(frozen=True, kw_only=True)
 class ComplexType:
+    """A type with inline members"""
+
     members: dict[str, Type]
 
     def __str__(self) -> str:
@@ -104,6 +122,8 @@ class ComplexType:
 
 @dataclass(frozen=True, kw_only=True)
 class ExtensionType:
+    """An extension of a type, adding members through a `ComplexType`"""
+
     base: Type
     extension: ComplexType
 
@@ -112,6 +132,8 @@ class ExtensionType:
 
 
 class Variance(StrEnum):
+    """The variance of a :class:`TypeVar`"""
+
     INVARIANT = "INVARIANT"
     COVARIANT = "COVARIANT"
     CONTRAVARIANT = "CONTRAVARIANT"
@@ -119,6 +141,8 @@ class Variance(StrEnum):
 
 @dataclass(frozen=True, kw_only=True)
 class TypeVar:
+    """A type variable, often used as type parameters for a generic type"""
+
     name: str
     bound: Optional[Type]
     variance: Variance = Variance.INVARIANT
@@ -136,6 +160,8 @@ class TypeVar:
 
 @dataclass(frozen=True, kw_only=True)
 class GenericType:
+    """A generic type, with type parameters and a generic body type"""
+
     name: str
     params: list[TypeVar]
     body: Type
@@ -146,6 +172,8 @@ class GenericType:
 
 @dataclass(frozen=True, kw_only=True)
 class AppliedType:
+    """An instance of a :class:`GenericType`, with concrete type arguments substituted in its body"""
+
     name: str
     args: list[Type]
     body: Type
@@ -156,6 +184,8 @@ class AppliedType:
 
 @dataclass(frozen=True, kw_only=True)
 class ConstraintType:
+    """A type with a constraint expression"""
+
     type: Type
     constraint: m.Expr
 
@@ -166,6 +196,8 @@ class ConstraintType:
 
 @dataclass(frozen=True, kw_only=True)
 class TupleType:
+    """A tuple type, containing any number of ordered item types"""
+
     items: tuple[Type, ...]
 
     def __str__(self) -> str:
@@ -174,6 +206,8 @@ class TupleType:
 
 @dataclass(frozen=True, kw_only=True)
 class ColumnType:
+    """A column type containing items of a given unique type"""
+
     type: Type
 
     def __str__(self) -> str:
@@ -182,6 +216,8 @@ class ColumnType:
 
 @dataclass(frozen=True, kw_only=True)
 class DataFrameType:
+    """A data-frame type, containing named columns of specific :class:`ColumnType`"""
+
     columns: list[Column]
 
     def __str__(self) -> str:
@@ -197,6 +233,8 @@ class DataFrameType:
 
 @dataclass(frozen=True, kw_only=True)
 class FrameGroupBy:
+    """A frame group-by object"""
+
     frame: DataFrameType
 
     def __str__(self) -> str:
@@ -205,6 +243,8 @@ class FrameGroupBy:
 
 @dataclass(frozen=True, kw_only=True)
 class ColumnGroupBy:
+    """A column group-by object"""
+
     column: ColumnType
 
     def __str__(self) -> str:
@@ -212,6 +252,19 @@ class ColumnGroupBy:
 
 
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
+    """Substitute type variables in the given type
+
+    This function is called recursively on inner type structures
+
+    Args:
+        type (Type): the type in which to substitute type variables
+        substitutions (dict[str, Type]): a mapping of type variable names to
+            concrete types
+
+    Returns:
+        Type: the resulting type with substitutions applied
+    """
+
     def sub_parameter(param: Function.Parameter):
         return Function.Parameter(
             pos=param.pos,
@@ -354,6 +407,14 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
 
 
 def unfold_type(type: Type) -> Type:
+    """Unfold a chain of :class:`DerivedType` to get the root supertype
+
+    Args:
+        type (Type): the type to unfold
+
+    Returns:
+        Type: the root supertype
+    """
     match type:
         case DerivedType(type=ref_type):
             return unfold_type(ref_type)
@@ -362,6 +423,15 @@ def unfold_type(type: Type) -> Type:
 
 
 def to_annotation(type: Type) -> str:
+    """Convert the given type to a Python annotation string
+
+    Args:
+        type (Type): the type to convert
+
+    Returns:
+        str: the annotation string
+    """
+
     def _params_annotation(spec: ParamSpec) -> str:
         if len(spec.kw) != 0:
             return "..."
@@ -430,6 +500,8 @@ def to_annotation(type: Type) -> str:
 
 @dataclass(frozen=True, kw_only=True)
 class Predicate:
+    """A predicate"""
+
     type: Type
     body: m.Expr
     alias: bool