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

feat(checker): add head/tail methods

Browse Source
This commit is contained in:
Louis Heredero
2026-07-03 12:13:30 +02:00
parent 4ea15519f3
commit ec80b1e92e
4 changed files with 319 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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

7
tests/cases/checker/09_frame_ops.py
View File

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

224
tests/cases/checker/09_frame_ops.py.ref.json
View File

@@ -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"
}
}
}
]
}