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

feate(checker): add some frame/column attributes

Browse Source
This commit is contained in:
Louis Heredero
2026-07-03 12:07:36 +02:00
parent 7a6e01cff8
commit 4ea15519f3
7 changed files with 363 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
midas/checker/frames/column_manager.py
View File

@@ -1,12 +1,13 @@
from __future__ import annotations from __future__ import annotations
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Optional
import midas.ast.python as p import midas.ast.python as p
from midas.ast.location import Location from midas.ast.location import Location
from midas.checker.frames.column_groupby_methods import Call as GroupByCall from midas.checker.frames.column_groupby_methods import Call as GroupByCall
from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry
from midas.checker.frames.column_methods import Call, ColumnMethodRegistry from midas.checker.frames.column_methods import Call, ColumnMethodRegistry
from midas.checker.registry import TypesRegistry
from midas.checker.types import ColumnGroupBy, ColumnType, Type from midas.checker.types import ColumnGroupBy, ColumnType, Type
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -60,3 +61,18 @@ class ColumnManager:
keywords=keywords, keywords=keywords,
) )
return self.groupby_method_resolver.call(method, call) return self.groupby_method_resolver.call(method, call)
def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]:
types: TypesRegistry = self.typer.types
match name:
case "ndim" | "size":
return types.get_type("int")
case "shape":
return types.tuple_of("int")
case "T":
return column
case _:
return None

13
midas/checker/frames/frame_manager.py
View File

@@ -7,6 +7,7 @@ from midas.ast.location import Location
from midas.checker.frames.frame_groupby_methods import Call as GroupByCall from midas.checker.frames.frame_groupby_methods import Call as GroupByCall
from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry
from midas.checker.frames.frame_methods import Call, FrameMethodRegistry from midas.checker.frames.frame_methods import Call, FrameMethodRegistry
from midas.checker.registry import TypesRegistry
from midas.checker.reporter import FileReporter from midas.checker.reporter import FileReporter
from midas.checker.types import ( from midas.checker.types import (
ColumnGroupBy, ColumnGroupBy,
@@ -240,3 +241,15 @@ class FrameManager:
keywords=keywords, keywords=keywords,
) )
return self.groupby_method_resolver.call(method, call) return self.groupby_method_resolver.call(method, call)
def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]:
types: TypesRegistry = self.typer.types
match name:
case "ndim" | "size":
return types.get_type("int")
case "shape":
return types.tuple_of("int", "int")
case _:
return None

4
midas/checker/preamble.py
View File

@@ -108,8 +108,8 @@ class Preamble(Environment):
], ],
) )
def _list_of(self, item_type: Type) -> Type: def _list_of(self, item_type: str | Type) -> Type:
return self._types.apply_generic(self._types.get_type("list"), [item_type]) return self._types.list_of(item_type)
def _def_type_constructor( def _def_type_constructor(
self, name: str, py_function: Optional[Callable[..., Any]] = None self, name: str, py_function: Optional[Callable[..., Any]] = None

8
midas/checker/python.py
View File

@@ -659,6 +659,14 @@ class PythonTyper(
def visit_get_expr(self, expr: p.GetExpr) -> Type: def visit_get_expr(self, expr: p.GetExpr) -> Type:
object: Type = self.type_of(expr.object) object: Type = self.type_of(expr.object)
member: Optional[Type] = self.types.lookup_member(object, expr.name) member: Optional[Type] = self.types.lookup_member(object, expr.name)
if member is None:
match object:
case DataFrameType():
member = self.frame_mgr.get_attribute(object, expr.name)
case ColumnType():
member = self.column_mgr.get_attribute(object, expr.name)
if member is None: if member is None:
self.reporter.warning( self.reporter.warning(
expr.location, f"Unknown member '{expr.name}' of {object}" expr.location, f"Unknown member '{expr.name}' of {object}"

26
midas/checker/registry.py
View File

@@ -452,3 +452,29 @@ class TypesRegistry:
def lookup_predicate(self, name: str) -> Optional[Predicate]: def lookup_predicate(self, name: str) -> Optional[Predicate]:
return self._predicates.get(name) return self._predicates.get(name)
def _by_name_or_type(self, name_or_type: str | Type) -> Type:
if isinstance(name_or_type, str):
return self.get_type(name_or_type)
return name_or_type
def list_of(self, item_type: str | Type) -> Type:
list_ = self.get_type("list")
return self.apply_generic(list_, [self._by_name_or_type(item_type)])
def tuple_of(self, *item_types: str | Type) -> Type:
tuple_ = self.get_type("tuple")
return self.apply_generic(
tuple_,
[self._by_name_or_type(item_type) for item_type in item_types],
)
def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type:
dict_ = self.get_type("dict")
return self.apply_generic(
dict_,
[
self._by_name_or_type(key_type),
self._by_name_or_type(value_type),
],
)

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

@@ -99,3 +99,12 @@ _ = col_gb.prod()
_ = col_gb.std() _ = col_gb.std()
_ = col_gb.sum() _ = col_gb.sum()
_ = col_gb.var() _ = col_gb.var()
# Attributes
_ = df1.ndim # int
_ = df1.size # int
_ = df1.shape # (int, int)
_ = col1.ndim # int
_ = col1.size # int
_ = col1.shape # (int)
_ = col1.T # Column[int]

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

@@ -4407,6 +4407,294 @@
"type": { "type": {
"type": {} "type": {}
} }
},
{
"location": {
"from": "L104:4",
"to": "L104: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": "L104:4",
"to": "L104:12"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "df1"
},
"name": "ndim"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L105:4",
"to": "L105: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": "L105:4",
"to": "L105:12"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "df1"
},
"name": "size"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L106:4",
"to": "L106: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": "L106:4",
"to": "L106:13"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "df1"
},
"name": "shape"
},
"type": {
"items": [
{
"name": "int"
},
{
"name": "int"
}
]
}
},
{
"location": {
"from": "L107:4",
"to": "L107:8"
},
"expr": {
"_type": "VariableExpr",
"name": "col1"
},
"type": {
"type": {
"name": "int"
}
}
},
{
"location": {
"from": "L107:4",
"to": "L107:13"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "col1"
},
"name": "ndim"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L108:4",
"to": "L108:8"
},
"expr": {
"_type": "VariableExpr",
"name": "col1"
},
"type": {
"type": {
"name": "int"
}
}
},
{
"location": {
"from": "L108:4",
"to": "L108:13"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "col1"
},
"name": "size"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L109:4",
"to": "L109:8"
},
"expr": {
"_type": "VariableExpr",
"name": "col1"
},
"type": {
"type": {
"name": "int"
}
}
},
{
"location": {
"from": "L109:4",
"to": "L109:14"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "col1"
},
"name": "shape"
},
"type": {
"items": [
{
"name": "int"
}
]
}
},
{
"location": {
"from": "L110:4",
"to": "L110:8"
},
"expr": {
"_type": "VariableExpr",
"name": "col1"
},
"type": {
"type": {
"name": "int"
}
}
},
{
"location": {
"from": "L110:4",
"to": "L110:10"
},
"expr": {
"_type": "GetExpr",
"object": {
"_type": "VariableExpr",
"name": "col1"
},
"name": "T"
},
"type": {
"type": {
"name": "int"
}
}
} }
] ]
} }