From 4ea15519f38f2b6df557573d027a5c7962d9f8b5 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 12:07:36 +0200 Subject: [PATCH] feate(checker): add some frame/column attributes --- midas/checker/frames/column_manager.py | 18 +- midas/checker/frames/frame_manager.py | 13 + midas/checker/preamble.py | 4 +- midas/checker/python.py | 8 + midas/checker/registry.py | 26 ++ tests/cases/checker/09_frame_ops.py | 9 + tests/cases/checker/09_frame_ops.py.ref.json | 288 +++++++++++++++++++ 7 files changed, 363 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index 0793034..f62d025 100644 --- a/midas/checker/frames/column_manager.py +++ b/midas/checker/frames/column_manager.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import midas.ast.python as p 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 ColumnGroupByMethodRegistry from midas.checker.frames.column_methods import Call, ColumnMethodRegistry +from midas.checker.registry import TypesRegistry from midas.checker.types import ColumnGroupBy, ColumnType, Type if TYPE_CHECKING: @@ -60,3 +61,18 @@ class ColumnManager: keywords=keywords, ) 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 diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 32bc782..8b99f14 100644 --- a/midas/checker/frames/frame_manager.py +++ b/midas/checker/frames/frame_manager.py @@ -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 FrameGroupByMethodRegistry from midas.checker.frames.frame_methods import Call, FrameMethodRegistry +from midas.checker.registry import TypesRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ( ColumnGroupBy, @@ -240,3 +241,15 @@ class FrameManager: keywords=keywords, ) 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 diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index b35d4dc..a60dabb 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -108,8 +108,8 @@ class Preamble(Environment): ], ) - def _list_of(self, item_type: Type) -> Type: - return self._types.apply_generic(self._types.get_type("list"), [item_type]) + def _list_of(self, item_type: str | Type) -> Type: + return self._types.list_of(item_type) def _def_type_constructor( self, name: str, py_function: Optional[Callable[..., Any]] = None diff --git a/midas/checker/python.py b/midas/checker/python.py index 99766d8..9103cf4 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -659,6 +659,14 @@ class PythonTyper( def visit_get_expr(self, expr: p.GetExpr) -> Type: object: Type = self.type_of(expr.object) 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: self.reporter.warning( expr.location, f"Unknown member '{expr.name}' of {object}" diff --git a/midas/checker/registry.py b/midas/checker/registry.py index d1074a7..92b54a6 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -452,3 +452,29 @@ class TypesRegistry: def lookup_predicate(self, name: str) -> Optional[Predicate]: 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), + ], + ) diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 18951f5..6183698 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -99,3 +99,12 @@ _ = col_gb.prod() _ = col_gb.std() _ = col_gb.sum() _ = 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] diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 83540db..79d7266 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -4407,6 +4407,294 @@ "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" + } + } } ] } \ No newline at end of file