fix(checker): delegate frame aggregate methods to columns

midas/checker/frames/column_groupby_methods.py

@@ -7,7 +7,7 @@ import midas.ast.python as p
 from midas.ast.location import Location
 from midas.checker.dispatcher import CallResult
 from midas.checker.frames.utils import MethodRegistry, method
-from midas.checker.types import ColumnGroupBy, Function, Type
+from midas.checker.types import ColumnGroupBy, ColumnType, Function, TopType, Type
 
 if TYPE_CHECKING:
     from midas.checker.python import TypedExpr
@@ -28,37 +28,46 @@ class Call:
 
 
 class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
-    @method()
+    NAMED_ARGS: dict[str, str] = {
-    def mean(self, call: Call) -> Type:
+        "numeric_only": "bool",
-        bool_ = self.types.get_type("bool")
+        "skipna": "bool",
+        "engine": "str",
+        "engine_kwargs": "dict",
+    }
+
+    def _aggregate(
+        self,
+        call: Call,
+        args: list[str | tuple[str, str, bool]] = [],
+        *,
+        preserve_inner_type: bool = False,
+    ) -> Type:
+        real_args: list[Function.Argument] = []
+        for i, arg in enumerate(args):
+            match arg:
+                case str() as name:
+                    arg = Function.Argument(
+                        pos=i,
+                        name=name,
+                        type=self.types.get_type(self.NAMED_ARGS[name]),
+                        required=False,
+                    )
+                case (name, type, required):
+                    arg = Function.Argument(
+                        pos=i,
+                        name=name,
+                        type=self.types.get_type(type),
+                        required=required,
+                    )
+            real_args.append(arg)
+
         signature = Function(
-            args=[
+            args=real_args,
-                Function.Argument(
+            returns=(
-                    pos=0,
+                call.groupby.column
-                    name="numeric_only",
+                if preserve_inner_type
-                    type=bool_,
+                else ColumnType(type=TopType())
-                    required=False,
+            ),
-                ),
-                Function.Argument(
-                    pos=1,
-                    name="skipna",
-                    type=bool_,
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=2,
-                    name="engine",
-                    type=self.types.get_type("str"),
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=3,
-                    name="engine_kwargs",
-                    type=self.types.get_type("dict"),
-                    required=False,
-                ),
-            ],
-            returns=call.groupby.column,
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -68,3 +77,127 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
             keywords=call.keywords,
         )
         return result.result
+
+    @method()
+    def kurt(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            ["skipna", "numeric_only"],
+        )
+
+    @method()
+    def max(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+            preserve_inner_type=True,
+        )
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            ["numeric_only", "skipna", "engine", "engine_kwargs"],
+        )
+
+    @method()
+    def median(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            ["numeric_only", "skipna"],
+            preserve_inner_type=True,
+        )
+
+    @method()
+    def min(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+            preserve_inner_type=True,
+        )
+
+    @method()
+    def prod(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+            ],
+        )
+
+    @method()
+    def std(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                (
+                    "ddof",
+                    "int",
+                    False,
+                ),
+                "engine",
+                "engine_kwargs",
+                "numeric_only",
+                "skipna",
+            ],
+        )
+
+    @method()
+    def sum(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+        )
+
+    @method()
+    def var(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                (
+                    "var",
+                    "int",
+                    False,
+                ),
+                "engine",
+                "engine_kwargs",
+                "numeric_only",
+                "skipna",
+            ],
+        )

midas/checker/frames/column_methods.py

@@ -160,7 +160,13 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
     def eq(self, call: Call) -> Type:
         return self._element_wise(call, "__eq__")
 
-    def _statistical(self, call: Call, kwargs: list[Function.Argument] = []) -> Type:
+    def _aggregate(
+        self,
+        call: Call,
+        kwargs: list[Function.Argument] = [],
+        *,
+        preserve_inner_type: bool = False,
+    ) -> Type:
         signature = Function(
             kw_args=[
                 Function.Argument(
@@ -171,7 +177,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
                 ),
                 *kwargs,
             ],
-            returns=ColumnType(type=TopType()),
+            returns=call.column if preserve_inner_type else ColumnType(type=TopType()),
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -184,35 +190,35 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method("kurtosis", "kurt")
     def kurtosis(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call)
 
     @method()
     def max(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call, preserve_inner_type=True)
 
     @method()
     def mean(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call)
 
     @method()
     def median(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call, preserve_inner_type=True)
 
     @method()
     def min(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call, preserve_inner_type=True)
 
     @method()
     def mode(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call, preserve_inner_type=True)
 
     @method("product", "prod")
     def product(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call)
 
     @method()
     def std(self, call: Call) -> Type:
-        return self._statistical(
+        return self._aggregate(
             call,
             [
                 Function.Argument(
@@ -226,11 +232,11 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method()
     def sum(self, call: Call) -> Type:
-        return self._statistical(call)
+        return self._aggregate(call)
 
     @method()
     def var(self, call: Call) -> Type:
-        return self._statistical(
+        return self._aggregate(
             call,
             [
                 Function.Argument(

midas/checker/frames/frame_groupby_methods.py

@@ -5,9 +5,15 @@ from typing import TYPE_CHECKING
 
 import midas.ast.python as p
 from midas.ast.location import Location
-from midas.checker.dispatcher import CallResult
 from midas.checker.frames.utils import MethodRegistry, method
-from midas.checker.types import FrameGroupBy, Function, Type
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    DataFrameType,
+    FrameGroupBy,
+    Type,
+    UnknownType,
+)
 
 if TYPE_CHECKING:
     from midas.checker.python import TypedExpr
@@ -35,161 +41,63 @@ class FrameGroupByMethodRegistry(MethodRegistry[Call]):
         "engine_kwargs": "dict",
     }
 
-    def _aggregate(
+    def _aggregate(self, call: Call, method: str) -> Type:
-        self, call: Call, args: list[str | tuple[str, str, bool]] = []
+        new_columns: list[DataFrameType.Column] = []
-    ) -> Type:
-        real_args: list[Function.Argument] = []
-        for i, arg in enumerate(args):
-            match arg:
-                case str() as name:
-                    arg = Function.Argument(
-                        pos=i,
-                        name=name,
-                        type=self.types.get_type(self.NAMED_ARGS[name]),
-                        required=False,
-                    )
-                case (name, type, required):
-                    arg = Function.Argument(
-                        pos=i,
-                        name=name,
-                        type=self.types.get_type(type),
-                        required=required,
-                    )
-            real_args.append(arg)
 
-        signature = Function(
+        for column in call.groupby.frame.columns:
-            args=real_args,
+            column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type)
-            returns=call.groupby.frame,
+            result_type: Type = self.typer.call_method(
-        )
+                location=call.location,
+                call_expr=call.call_expr,
+                obj=(call.groupby_expr, column_groupby),
+                method_name=method,
+                positional=call.positional,
+                keywords=call.keywords,
+            )
+            if not isinstance(result_type, ColumnType):
+                result_type = ColumnType(type=UnknownType())
+            new_columns.append(
+                DataFrameType.Column(
+                    index=column.index,
+                    name=column.name,
+                    type=result_type,
+                )
+            )
 
-        result: CallResult = self.dispatcher.get_result(
+        return DataFrameType(columns=new_columns)
-            location=call.location,
-            callee=signature,
-            positional=call.positional,
-            keywords=call.keywords,
-        )
-        return result.result
 
     @method()
     def kurt(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "kurt")
-            call,
-            [
-                "skipna",
-                "numeric_only",
-            ],
-        )
 
     @method()
     def max(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "max")
-            call,
-            [
-                "numeric_only",
-                (
-                    "min_count",
-                    "int",
-                    False,
-                ),
-                "skipna",
-                "engine",
-                "engine_kwargs",
-            ],
-        )
 
     @method()
     def mean(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "mean")
-            call,
-            ["numeric_only", "skipna", "engine", "engine_kwargs"],
-        )
 
     @method()
     def median(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "median")
-            call,
-            ["numeric_only", "skipna"],
-        )
 
     @method()
     def min(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "min")
-            call,
-            [
-                "numeric_only",
-                (
-                    "min_count",
-                    "int",
-                    False,
-                ),
-                "skipna",
-                "engine",
-                "engine_kwargs",
-            ],
-        )
 
     @method()
     def prod(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "prod")
-            call,
-            [
-                "numeric_only",
-                (
-                    "min_count",
-                    "int",
-                    False,
-                ),
-                "skipna",
-            ],
-        )
 
     @method()
     def std(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "std")
-            call,
-            [
-                (
-                    "ddof",
-                    "int",
-                    False,
-                ),
-                "engine",
-                "engine_kwargs",
-                "numeric_only",
-                "skipna",
-            ],
-        )
 
     @method()
     def sum(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "sum")
-            call,
-            [
-                "numeric_only",
-                (
-                    "min_count",
-                    "int",
-                    False,
-                ),
-                "skipna",
-                "engine",
-                "engine_kwargs",
-            ],
-        )
 
     @method()
     def var(self, call: Call) -> Type:
-        return self._aggregate(
+        return self._aggregate(call, "var")
-            call,
-            [
-                (
-                    "var",
-                    "int",
-                    False,
-                ),
-                "engine",
-                "engine_kwargs",
-                "numeric_only",
-                "skipna",
-            ],
-        )

midas/checker/python.py

@@ -222,7 +222,7 @@ class PythonTyper(
         method_name: str,
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
-    ) -> Optional[Type]:
+    ) -> Type:
         unfolded: Type = unfold_type(obj[1])
         match unfolded:
             case DataFrameType():
@@ -580,9 +580,8 @@ class PythonTyper(
         right: TypedExpr,
         method: str,
     ) -> Type:
-        result: Optional[Type]
         try:
-            result = self.call_method(
+            return self.call_method(
                 location=location,
                 call_expr=expr,
                 obj=left,
@@ -597,8 +596,6 @@ class PythonTyper(
             )
             return UnknownType()
 
-        return result or UnknownType()
-
     def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
         method: Optional[str] = PY_UNARY_METHODS.get(expr.operator.__class__)
         if method is None:
@@ -610,9 +607,8 @@ class PythonTyper(
 
         operand: Type = self.type_of(expr.right)
 
-        result: Optional[Type]
         try:
-            result = self.call_method(
+            return self.call_method(
                 location=expr.location,
                 call_expr=expr,
                 obj=(expr.right, operand),
@@ -627,8 +623,6 @@ class PythonTyper(
             )
             return UnknownType()
 
-        return result or UnknownType()
-
     def visit_call_expr(self, expr: p.CallExpr) -> Type:
         match expr.callee:
             case p.VariableExpr(name="TypeVar"):
@@ -644,16 +638,13 @@ class PythonTyper(
         match expr.callee:
             case p.GetExpr(object=obj, name=method):
                 obj_type: Type = self.type_of(obj)
-                return (
+                return self.call_method(
-                    self.call_method(
+                    location=expr.location,
-                        location=expr.location,
+                    call_expr=expr,
-                        call_expr=expr,
+                    obj=(obj, obj_type),
-                        obj=(obj, obj_type),
+                    method_name=method,
-                        method_name=method,
+                    positional=positional,
-                        positional=positional,
+                    keywords=keywords,
-                        keywords=keywords,
-                    )
-                    or UnknownType()
                 )
 
         callee: Type = self.type_of(expr.callee)

tests/__main__.py

@@ -0,0 +1,43 @@
+from typing import Type
+
+from midas.cli.ansi import Ansi
+from tests.base import Tester
+from tests.checker import CheckerTester
+from tests.generator import GeneratorTester
+from tests.midas import MidasTester
+from tests.python import PythonTester
+
+
+def print_banner(name: str):
+    horizontal: str = "+" + "-" * (len(name) + 2) + "+"
+    print(horizontal)
+    print(f"| {name} |")
+    print(horizontal)
+
+
+def run_tests(tester_cls: Type[Tester]) -> bool:
+    print_banner(tester_cls.__name__)
+    tester: Tester = tester_cls()
+    success: bool = tester.run_all_tests()
+    print()
+    return success
+
+
+def main():
+    testers: list[Type[Tester]] = [
+        PythonTester,
+        MidasTester,
+        CheckerTester,
+        GeneratorTester,
+    ]
+
+    success: bool = all(map(run_tests, testers))
+
+    if success:
+        print(Ansi.FG(Ansi.BRIGHT_GREEN) + "All tests passed!" + Ansi.RESET)
+    else:
+        print(Ansi.FG(Ansi.BRIGHT_RED) + "Some tests failed!" + Ansi.RESET)
+
+
+if __name__ == "__main__":
+    main()

tests/base.py

@@ -7,6 +7,8 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Iterator, Protocol
 
+from midas.cli.ansi import Ansi
+
 
 class CaseResult(Protocol):
     def dumps(self) -> str: ...
@@ -44,8 +46,11 @@ class Tester(ABC):
 
         print(rule)
         for i, test in enumerate(tests):
-            print(f"Case {i+1}/{n}: {test.resolve().relative_to(self.CASES_DIR)}")
+            path: Path = test.resolve().relative_to(self.CASES_DIR)
+            print(f"{Ansi.FG(Ansi.BRIGHT_CYAN)}Case {i+1}/{n}: {path}{Ansi.RESET}")
+            print(Ansi.DIM, end="")
             success: bool = self._run_test(test)
+            print(Ansi.RESET, end="")
             if success:
                 successes += 1
             else:
@@ -146,8 +151,9 @@ class Tester(ABC):
                 if not success:
                     sys.exit(1)
             case None:
-                print("No subcommand provided. Available subcommands: run, update")
+                success: bool = tester.run_all_tests()
-                sys.exit(1)
+                if not success:
+                    sys.exit(1)
             case _:
                 print(f"Unknown subcommand '{args.subcommand}'")
                 sys.exit(1)

tests/cases/checker/09_frame_ops.py

@@ -38,14 +38,64 @@ _ = df1.sum()
 _ = df1.var()
 
 # Groupby
-gb = df1.groupby(by="a")
+df_gb = df1.groupby(by="a")
 
-_ = gb.kurt()
+_ = df_gb.kurt()
-_ = gb.max()
+_ = df_gb.max()
-_ = gb.mean()
+_ = df_gb.mean()
-_ = gb.median()
+_ = df_gb.median()
-_ = gb.min()
+_ = df_gb.min()
-_ = gb.prod()
+_ = df_gb.prod()
-_ = gb.std()
+_ = df_gb.std()
-_ = gb.sum()
+_ = df_gb.sum()
-_ = gb.var()
+_ = df_gb.var()
+
+
+# Columns
+
+col1 = df1["a"]
+col2 = df1["a"]
+
+# Arithmetic
+_ = col1 + col2
+_ = col1 - col2
+_ = col1 * col2
+_ = col1 / col2
+_ = col1 // col2
+_ = col1 % col2
+_ = col1**col2
+
+# Comparisons
+_ = col1 < col2
+_ = col1 > col2
+_ = col1 <= col2
+_ = col1 >= col2
+_ = col1 != col2
+_ = col1 == col2
+
+# Aggregate
+_ = col1.kurt()
+_ = col1.kurtosis()
+_ = col1.max()
+_ = col1.mean()
+_ = col1.median()
+_ = col1.min()
+_ = col1.mode()
+_ = col1.prod()
+_ = col1.product()
+_ = col1.std()
+_ = col1.sum()
+_ = col1.var()
+
+# Groupby
+col_gb = col1.groupby(level=0)
+
+_ = col_gb.kurt()
+_ = col_gb.max()
+_ = col_gb.mean()
+_ = col_gb.median()
+_ = col_gb.min()
+_ = col_gb.prod()
+_ = col_gb.std()
+_ = col_gb.sum()
+_ = col_gb.var()