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

feat(checker): add sort_values on frames and columns

Browse Source
This commit is contained in:
Louis Heredero
2026-07-07 11:28:06 +02:00
parent f48ebd49d1
commit c6b8c0a145
2 changed files with 155 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -420,6 +420,79 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
)
return result.result
@method()
def sort_values(self, call: Call) -> Type:
str_ = self.types.get_type("str")
bool_ = self.types.get_type("bool")
def make_overload(ascending: Type) -> Function:
return Function(
params=ParamSpec(
kw=[
Function.Parameter(
pos=0,
name="axis",
type=TopType(),
required=False,
),
Function.Parameter(
pos=1,
name="ascending",
type=ascending,
required=False,
),
Function.Parameter(
pos=2,
name="inplace",
type=bool_,
required=False,
unsupported=True,
),
Function.Parameter(
pos=3,
name="kind",
type=str_,
required=False,
),
Function.Parameter(
pos=4,
name="na_position",
type=str_,
required=False,
),
Function.Parameter(
pos=5,
name="ignore_index",
type=bool_,
required=False,
),
Function.Parameter(
pos=6,
name="key",
type=TopType(),
required=False,
),
],
),
returns=call.column,
)
list_of = self.types.list_of
overloads: list[Type] = [
make_overload(bool_),
make_overload(bool_),
make_overload(list_of(bool_)),
make_overload(list_of(bool_)),
]
result: CallResult = self.dispatcher.get_result(
location=call.location,
callee=OverloadedFunction(overloads=overloads),
positional=call.positional,
keywords=call.keywords,
)
return result.result
@method()
def groupby(self, call: Call) -> Type:
bool_: Type = self.types.get_type("bool")

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

@@ -513,6 +513,88 @@ class FrameMethodRegistry(MethodRegistry[Call]):
)
return result.result
@method()
def sort_values(self, call: Call) -> Type:
str_ = self.types.get_type("str")
bool_ = self.types.get_type("bool")
def make_overload(by: Type, ascending: Type) -> Function:
return Function(
params=ParamSpec(
mixed=[
Function.Parameter(
pos=0,
name="by",
type=by,
required=True,
),
],
kw=[
Function.Parameter(
pos=1,
name="axis",
type=TopType(),
required=False,
),
Function.Parameter(
pos=2,
name="ascending",
type=ascending,
required=False,
),
Function.Parameter(
pos=3,
name="inplace",
type=bool_,
required=False,
unsupported=True,
),
Function.Parameter(
pos=4,
name="kind",
type=str_,
required=False,
),
Function.Parameter(
pos=5,
name="na_position",
type=str_,
required=False,
),
Function.Parameter(
pos=6,
name="ignore_index",
type=bool_,
required=False,
),
Function.Parameter(
pos=7,
name="key",
type=TopType(),
required=False,
),
],
),
returns=call.frame,
)
list_of = self.types.list_of
overloads: list[Type] = [
make_overload(by=str_, ascending=bool_),
make_overload(by=list_of(str_), ascending=bool_),
make_overload(by=str_, ascending=list_of(bool_)),
make_overload(by=list_of(str_), ascending=list_of(bool_)),
]
# TODO: check that literal strings in `by` are valid columns
result: CallResult = self.dispatcher.get_result(
location=call.location,
callee=OverloadedFunction(overloads=overloads),
positional=call.positional,
keywords=call.keywords,
)
return result.result
@method()
def groupby(self, call: Call) -> Type:
bool_: Type = self.types.get_type("bool")