diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py new file mode 100644 index 0000000..1ff9799 --- /dev/null +++ b/midas/checker/frames/column_manager.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.frames.column_methods import Call, ColumnMethodRegistry +from midas.checker.types import ColumnType, Type + +if TYPE_CHECKING: + from midas.checker.python import PythonTyper, TypedExpr + + +class ColumnManager: + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer) + + def call( + self, + method: str, + location: Location, + call_expr: p.Expr, + column: ColumnType, + column_expr: p.Expr, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + call: Call = Call( + location=location, + call_expr=call_expr, + column=column, + column_expr=column_expr, + positional=positional, + keywords=keywords, + ) + return self.method_resolver.call(method, call) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py new file mode 100644 index 0000000..da22a43 --- /dev/null +++ b/midas/checker/frames/column_methods.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.frames.utils import MethodRegistry +from midas.checker.types import ( + ColumnType, +) + +if TYPE_CHECKING: + from midas.checker.python import TypedExpr + + +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + call_expr: p.Expr + column: ColumnType + column_expr: p.Expr + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +class ColumnMethodRegistry(MethodRegistry[Call]): ... diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index d6ac0d8..95fa40a 100644 --- a/midas/checker/frames/utils.py +++ b/midas/checker/frames/utils.py @@ -23,18 +23,6 @@ if TYPE_CHECKING: from midas.checker.python import PythonTyper -@staticmethod -def method(*names: str): - def wrapper(func): - names_: tuple[str, ...] = names - if len(names_) == 0: - names_ = (func.__name__,) - setattr(func, "__method_names__", names_) - return func - - return wrapper - - class _MethodRegistryMeta(type): _methods: dict[str, Callable[..., Type]] = {} @@ -87,3 +75,18 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): self.reporter.warning(call.location, f"Unknown method {method}") return UnknownType() return func(self, call) + + +_Self = TypeVar("_Self", bound=MethodRegistry[Any]) +Method = Callable[[_Self, T], Type] + + +def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]: + def wrapper(func: Method[_Self, T]) -> Method[_Self, T]: + names_: tuple[str, ...] = names + if len(names_) == 0: + names_ = (func.__name__,) + setattr(func, "__method_names__", names_) + return func + + return wrapper diff --git a/midas/checker/python.py b/midas/checker/python.py index 2031542..e014681 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -9,6 +9,7 @@ from midas.ast.printer import MidasPrinter from midas.checker.dispatcher import CallDispatcher, CallResult from midas.checker.environment import Environment from midas.checker.evaluator import Evaluator +from midas.checker.frames.column_manager import ColumnManager from midas.checker.frames.frame_manager import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, @@ -82,6 +83,7 @@ class PythonTyper( self.reporter: FileReporter = reporter.for_file(None) self.types: TypesRegistry = types self.frame_mgr: FrameManager = FrameManager(self) + self.column_mgr: ColumnManager = ColumnManager(self) self.global_env: Environment = Preamble(self.types) self.env: Environment = self.global_env self.locals: dict[p.Expr, int] = {} @@ -233,6 +235,17 @@ class PythonTyper( keywords=keywords, ) + case ColumnType(): + return self.column_mgr.call( + method=method_name, + location=location, + call_expr=call_expr, + column=unfolded, + column_expr=obj[0], + positional=positional, + keywords=keywords, + ) + method: Optional[Type] = self.types.lookup_member(obj[1], method_name) if method is None: raise UndefinedMethodException