From ec80b1e92ee5c5e69a8882a770faaae9066c7502 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 12:13:30 +0200 Subject: [PATCH] feat(checker): add head/tail methods --- midas/checker/frames/column_methods.py | 44 ++++ midas/checker/frames/frame_methods.py | 44 ++++ tests/cases/checker/09_frame_ops.py | 7 + tests/cases/checker/09_frame_ops.py.ref.json | 224 +++++++++++++++++++ 4 files changed, 319 insertions(+) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index d2559b8..02de5a3 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -248,6 +248,50 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ], ) + @method() + def head(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.column, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + @method() + def tail(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.column, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 2bd8c8b..bd9e3ee 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -327,6 +327,50 @@ class FrameMethodRegistry(MethodRegistry[Call]): ], ) + @method() + def head(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.frame, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + @method() + def tail(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.frame, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 6183698..6d478f6 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -108,3 +108,10 @@ _ = col1.ndim # int _ = col1.size # int _ = col1.shape # (int) _ = col1.T # Column[int] + + +# Misc +_ = df1.head() +_ = df1.tail() +_ = col1.head() +_ = col1.tail() diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 79d7266..6c77757 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -4695,6 +4695,230 @@ "name": "int" } } + }, + { + "location": { + "from": "L114:4", + "to": "L114:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L114:4", + "to": "L114:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "head" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L115:4", + "to": "L115:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L115:4", + "to": "L115:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "tail" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L116:4", + "to": "L116:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L116:4", + "to": "L116:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "head" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L117:4", + "to": "L117:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L117:4", + "to": "L117:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "tail" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } } ] } \ No newline at end of file