From b11a9bb8c6cb886eb6caf46ba6fe6e624d23ed82 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 6 Jul 2026 11:03:31 +0200 Subject: [PATCH] docs: add docstrings to frame classes --- .../checker/frames/column_groupby_methods.py | 19 +++ midas/checker/frames/column_manager.py | 39 ++++++ midas/checker/frames/column_methods.py | 37 +++++ midas/checker/frames/frame_groupby_methods.py | 18 ++- midas/checker/frames/frame_manager.py | 131 ++++++++++++++++++ midas/checker/frames/frame_methods.py | 34 +++++ midas/checker/frames/utils.py | 29 ++++ 7 files changed, 301 insertions(+), 6 deletions(-) diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py index e633be7..f84d8df 100644 --- a/midas/checker/frames/column_groupby_methods.py +++ b/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: diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index f62d025..8adeb1b 100644 --- a/midas/checker/frames/column_manager.py +++ b/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": diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 9396d0b..5b65314 100644 --- a/midas/checker/frames/column_methods.py +++ b/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 diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index 755cab7..cabf88a 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/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] = { - "numeric_only": "bool", - "skipna": "bool", - "engine": "str", - "engine_kwargs": "dict", - } + """The method registry for frame group-by types""" 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: diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 8b99f14..ca75187 100644 --- a/midas/checker/frames/frame_manager.py +++ b/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": diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 7866784..3832865 100644 --- a/midas/checker/frames/frame_methods.py +++ b/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 diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index dee4696..044ca8e 100644 --- a/midas/checker/frames/utils.py +++ b/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: