From a640b8b3ddac56a13b26642b04e1ba0495d2d59e Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 6 Jul 2026 15:56:57 +0200 Subject: [PATCH] feat(checker): add copy and info methods --- midas/checker/frames/column_methods.py | 91 ++++++++++++++++++++++++++ midas/checker/frames/frame_methods.py | 90 +++++++++++++++++++++++++ 2 files changed, 181 insertions(+) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 5b65314..345a884 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -13,10 +13,12 @@ from midas.checker.types import ( ColumnType, Function, GenericType, + OverloadedFunction, ParamSpec, TopType, Type, TypeVar, + UnitType, UnknownType, unfold_type, ) @@ -44,6 +46,25 @@ class Call: class ColumnMethodRegistry(MethodRegistry[Call]): """The method registry for column types""" + def _simple_call(self, call: Call, function: Type) -> Type: + """Get the result of calling a simple method + + This function is a simple wrapper around :func:`dispatcher.CallDispatcher.get_result` + Args: + call (Call): the call that triggered this resolution + function (Type): the function type + + Returns: + Type: the return type + """ + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=function, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + def _element_binary_op(self, call: Call, method: str) -> ColumnType: """Compute the result of an element-wise binary operation @@ -127,6 +148,76 @@ class ColumnMethodRegistry(MethodRegistry[Call]): return result.result + @method() + def copy(self, call: Call) -> Type: + return self._simple_call( + call, + Function( + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="deep", + type=self.types.get_type("bool"), + required=False, + ) + ] + ), + returns=call.column, + ), + ) + + @method() + def info(self, call: Call) -> Type: + def make_overload(memory_usage: Type, required: bool = False) -> Type: + return Function( + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="verbose", + type=self.types.get_type("bool"), + required=False, + ), + Function.Parameter( + pos=1, + name="buf", + type=TopType(), + required=False, + ), + Function.Parameter( + pos=2, + name="max_cols", + type=self.types.get_type("int"), + required=False, + ), + Function.Parameter( + pos=3, + name="memory_usage", + type=memory_usage, + required=required, + ), + Function.Parameter( + pos=4, + name="show_counts", + type=self.types.get_type("bool"), + required=False, + ), + ] + ), + returns=UnitType(), + ) + + return self._simple_call( + call, + OverloadedFunction( + overloads=[ + make_overload(self.types.get_type("bool"), False), + make_overload(self.types.get_type("str"), True), + ], + ), + ) + @method("add", "__add__") def add(self, call: Call) -> Type: return self._element_wise(call, "__add__") diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 3832865..c5be389 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -17,6 +17,7 @@ from midas.checker.types import ( ParamSpec, TopType, Type, + UnitType, UnknownType, unfold_type, ) @@ -44,6 +45,25 @@ class Call: class FrameMethodRegistry(MethodRegistry[Call]): """The method registry for frame types""" + def _simple_call(self, call: Call, function: Type) -> Type: + """Get the result of calling a simple method + + This function is a simple wrapper around :func:`dispatcher.CallDispatcher.get_result` + Args: + call (Call): the call that triggered this resolution + function (Type): the function type + + Returns: + Type: the return type + """ + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=function, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + def _get_method_result( self, call: Call, @@ -194,6 +214,76 @@ class FrameMethodRegistry(MethodRegistry[Call]): return result.result + @method() + def copy(self, call: Call) -> Type: + return self._simple_call( + call, + Function( + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="deep", + type=self.types.get_type("bool"), + required=False, + ) + ] + ), + returns=call.frame, + ), + ) + + @method() + def info(self, call: Call) -> Type: + def make_overload(memory_usage: Type, required: bool = False) -> Type: + return Function( + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="verbose", + type=self.types.get_type("bool"), + required=False, + ), + Function.Parameter( + pos=1, + name="buf", + type=TopType(), + required=False, + ), + Function.Parameter( + pos=2, + name="max_cols", + type=self.types.get_type("int"), + required=False, + ), + Function.Parameter( + pos=3, + name="memory_usage", + type=memory_usage, + required=required, + ), + Function.Parameter( + pos=4, + name="show_counts", + type=self.types.get_type("bool"), + required=False, + ), + ] + ), + returns=UnitType(), + ) + + return self._simple_call( + call, + OverloadedFunction( + overloads=[ + make_overload(self.types.get_type("bool"), False), + make_overload(self.types.get_type("str"), True), + ], + ), + ) + @method("add", "__add__") def add(self, call: Call) -> Type: return self._element_wise(call, "__add__")