feat(checker): add head/tail methods

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

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

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

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

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

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

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),
+            ],
+        )

tests/cases/checker/09_frame_ops.py

@@ -99,3 +99,19 @@ _ = 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]
+
+
+# Misc
+_ = df1.head()
+_ = df1.tail()
+_ = col1.head()
+_ = col1.tail()

tests/cases/checker/09_frame_ops.py.ref.json

@@ -4407,6 +4407,518 @@
       "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"
+        }
+      }
+    },
+    {
+      "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"
+        }
+      }
     }
   ]
 }