HEL/midas
1
0
Fork 0
Code Issues Pull Requests Releases Activity

docs: add docstrings to frame classes

Browse Source
This commit is contained in:
Louis Heredero
2026-07-06 11:03:31 +02:00
parent bac0e334d5
commit b11a9bb8c6
7 changed files with 301 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
midas/checker/frames/column_groupby_methods.py
View File

@@ -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:

39
midas/checker/frames/column_manager.py
View File

@@ -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":

37
midas/checker/frames/column_methods.py
View File

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

18
midas/checker/frames/frame_groupby_methods.py
View File

@@ -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:

131
midas/checker/frames/frame_manager.py
View File

@@ -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":

34
midas/checker/frames/frame_methods.py
View File

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

29
midas/checker/frames/utils.py
View File

@@ -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: