Compare commits

...

4 Commits

8 changed files with 676 additions and 561 deletions

View File

@@ -10,6 +10,8 @@ from midas.checker.types import (
ColumnType,
DataFrameType,
Function,
OverloadedFunction,
TopType,
Type,
UnknownType,
unfold_type,
@@ -85,6 +87,9 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
self,
call: Call,
) -> Type:
# TODO: support add with scalar, sequence, Series, dict
# TODO: check operation exists on inner column types
new_columns: list[DataFrameType.Column] = []
by_name: dict[str, DataFrameType.Column] = {}
@@ -151,3 +156,43 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
)
or UnknownType()
)
@frame_method()
def mean(self, call: Call) -> Type:
with_axis = Function(
kw_args=[
Function.Argument(
pos=0,
name="axis",
type=self.types.get_type("int"),
required=False,
)
],
returns=ColumnType(type=TopType()),
)
without_axis = Function(
kw_args=[
Function.Argument(
pos=0,
name="axis",
type=self.types.get_type("None"),
required=True,
)
],
returns=TopType(),
)
overload = OverloadedFunction(
overloads=[
with_axis,
without_axis,
]
)
return (
self.typer._get_call_result(
location=call.location,
callee=overload,
positional=call.positional,
keywords=call.keywords,
)
or UnknownType()
)

View File

@@ -47,6 +47,10 @@ class ReturnException(Exception):
pass
class UndefinedMethodException(Exception):
pass
@dataclass(frozen=True, kw_only=True)
class MappedArgument:
expr: p.Expr
@@ -195,6 +199,36 @@ class PythonTyper(
return self.env.get_at(distance, name)
return self.global_env.get(name)
def call_method(
self,
location: Location,
obj: Type,
method_name: str,
positional: list[TypedExpr],
keywords: dict[str, TypedExpr],
) -> Optional[Type]:
unfolded: Type = unfold_type(obj)
match unfolded:
case DataFrameType():
return self.frame_mgr.call(
method=method_name,
location=location,
frame=unfolded,
positional=positional,
keywords=keywords,
)
method: Optional[Type] = self.types.lookup_member(obj, method_name)
if method is None:
raise UndefinedMethodException
return self._get_call_result(
location,
method,
positional,
keywords,
)
def is_subtype(self, type1: Type, type2: Type) -> bool:
return self.types.is_subtype(type1, type2)
@@ -468,20 +502,16 @@ class PythonTyper(
left: Type = self.type_of(left_expr)
right: Type = self.type_of(right_expr)
operation: Optional[Type] = self.types.lookup_member(left, method)
if operation is None:
result: Optional[Type]
try:
result = self.call_method(location, left, method, [(right_expr, right)], {})
except UndefinedMethodException:
self.reporter.error(
location,
f"Undefined operation {method} between {left} and {right}",
)
return UnknownType()
result: Optional[Type] = self._get_call_result(
location,
operation,
[(right_expr, right)],
{},
)
return result or UnknownType()
def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
@@ -494,20 +524,17 @@ class PythonTyper(
return UnknownType()
operand: Type = self.type_of(expr.right)
operation: Optional[Type] = self.types.lookup_member(operand, method)
if operation is None:
result: Optional[Type]
try:
result = self.call_method(expr.location, operand, method, [], {})
except UndefinedMethodException:
self.reporter.error(
expr.location,
f"Undefined operation {method} for {operand}",
)
return UnknownType()
result: Optional[Type] = self._get_call_result(
expr.location,
operation,
[],
{},
)
return result or UnknownType()
def visit_call_expr(self, expr: p.CallExpr) -> Type:
@@ -567,6 +594,8 @@ class PythonTyper(
return self.types.get_type("float")
case str():
return self.types.get_type("str")
case None:
return self.types.get_type("None")
case _:
self.reporter.warning(expr.location, f"Unknown literal {expr}")
return UnknownType()

View File

@@ -4,7 +4,35 @@
"type": "Warning",
"location": {
"start": [
6,
8,
12
],
"end": [
8,
43
]
},
"message": "ConstraintType not yet supported"
},
{
"type": "Warning",
"location": {
"start": [
10,
10
],
"end": [
10,
18
]
},
"message": "Unknown type 'datetime'"
},
{
"type": "Warning",
"location": {
"start": [
13,
4
],
"end": [
@@ -12,7 +40,7 @@
5
]
},
"message": "FrameType not yet supported"
"message": "Unknown type '_'"
}
],
"judgments": []

View File

@@ -328,6 +328,19 @@
},
"type": {}
},
{
"location": {
"from": "L6:9",
"to": "L6:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L6:5",
@@ -373,19 +386,6 @@
}
}
},
{
"location": {
"from": "L6:9",
"to": "L6:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L6:5",
@@ -407,6 +407,32 @@
},
"type": {}
},
{
"location": {
"from": "L7:9",
"to": "L7:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L7:12",
"to": "L7:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L7:5",
@@ -452,32 +478,6 @@
}
}
},
{
"location": {
"from": "L7:9",
"to": "L7:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L7:12",
"to": "L7:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L7:5",
@@ -503,6 +503,32 @@
},
"type": {}
},
{
"location": {
"from": "L8:9",
"to": "L8:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L8:14",
"to": "L8:17"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L8:5",
@@ -548,32 +574,6 @@
}
}
},
{
"location": {
"from": "L8:9",
"to": "L8:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L8:14",
"to": "L8:17"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L8:5",
@@ -600,6 +600,45 @@
},
"type": {}
},
{
"location": {
"from": "L9:9",
"to": "L9:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L9:12",
"to": "L9:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L9:17",
"to": "L9:23"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L9:5",
@@ -645,45 +684,6 @@
}
}
},
{
"location": {
"from": "L9:9",
"to": "L9:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L9:12",
"to": "L9:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L9:17",
"to": "L9:23"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L9:5",
@@ -713,6 +713,45 @@
},
"type": {}
},
{
"location": {
"from": "L10:9",
"to": "L10:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L10:12",
"to": "L10:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L10:19",
"to": "L10:22"
},
"expr": {
"_type": "LiteralExpr",
"value": 3.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L10:5",
@@ -758,45 +797,6 @@
}
}
},
{
"location": {
"from": "L10:9",
"to": "L10:10"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L10:12",
"to": "L10:15"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L10:19",
"to": "L10:22"
},
"expr": {
"_type": "LiteralExpr",
"value": 3.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L10:5",
@@ -827,6 +827,19 @@
},
"type": {}
},
{
"location": {
"from": "L11:11",
"to": "L11:12"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L11:5",
@@ -872,19 +885,6 @@
}
}
},
{
"location": {
"from": "L11:11",
"to": "L11:12"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L11:5",
@@ -906,6 +906,19 @@
},
"type": {}
},
{
"location": {
"from": "L12:11",
"to": "L12:17"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L12:5",
@@ -951,19 +964,6 @@
}
}
},
{
"location": {
"from": "L12:11",
"to": "L12:17"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L12:5",
@@ -985,6 +985,45 @@
},
"type": {}
},
{
"location": {
"from": "L14:10",
"to": "L14:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L14:13",
"to": "L14:16"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L14:20",
"to": "L14:26"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L14:6",
@@ -1030,45 +1069,6 @@
}
}
},
{
"location": {
"from": "L14:10",
"to": "L14:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L14:13",
"to": "L14:16"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L14:20",
"to": "L14:26"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L14:6",
@@ -1101,6 +1101,45 @@
"name": "bool"
}
},
{
"location": {
"from": "L15:10",
"to": "L15:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L15:15",
"to": "L15:18"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L15:22",
"to": "L15:28"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L15:6",
@@ -1146,45 +1185,6 @@
}
}
},
{
"location": {
"from": "L15:10",
"to": "L15:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L15:15",
"to": "L15:18"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L15:22",
"to": "L15:28"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L15:6",
@@ -1217,6 +1217,45 @@
"name": "bool"
}
},
{
"location": {
"from": "L16:10",
"to": "L16:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L16:15",
"to": "L16:21"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L16:25",
"to": "L16:28"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L16:6",
@@ -1262,45 +1301,6 @@
}
}
},
{
"location": {
"from": "L16:10",
"to": "L16:11"
},
"expr": {
"_type": "LiteralExpr",
"value": 1
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L16:15",
"to": "L16:21"
},
"expr": {
"_type": "LiteralExpr",
"value": "test"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L16:25",
"to": "L16:28"
},
"expr": {
"_type": "LiteralExpr",
"value": 2.0
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L16:6",
@@ -1333,6 +1333,45 @@
"name": "bool"
}
},
{
"location": {
"from": "L18:10",
"to": "L18:13"
},
"expr": {
"_type": "LiteralExpr",
"value": "a"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L18:15",
"to": "L18:16"
},
"expr": {
"_type": "LiteralExpr",
"value": 3
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L18:20",
"to": "L18:25"
},
"expr": {
"_type": "LiteralExpr",
"value": false
},
"type": {
"name": "bool"
}
},
{
"location": {
"from": "L18:6",
@@ -1378,45 +1417,6 @@
}
}
},
{
"location": {
"from": "L18:10",
"to": "L18:13"
},
"expr": {
"_type": "LiteralExpr",
"value": "a"
},
"type": {
"name": "str"
}
},
{
"location": {
"from": "L18:15",
"to": "L18:16"
},
"expr": {
"_type": "LiteralExpr",
"value": 3
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L18:20",
"to": "L18:25"
},
"expr": {
"_type": "LiteralExpr",
"value": false
},
"type": {
"name": "bool"
}
},
{
"location": {
"from": "L18:6",

View File

@@ -100,6 +100,32 @@
"name": "float"
}
},
{
"location": {
"from": "L11:13",
"to": "L11:15"
},
"expr": {
"_type": "VariableExpr",
"name": "v1"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L11:17",
"to": "L11:19"
},
"expr": {
"_type": "VariableExpr",
"name": "v2"
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L11:5",
@@ -135,32 +161,6 @@
}
}
},
{
"location": {
"from": "L11:13",
"to": "L11:15"
},
"expr": {
"_type": "VariableExpr",
"name": "v1"
},
"type": {
"name": "int"
}
},
{
"location": {
"from": "L11:17",
"to": "L11:19"
},
"expr": {
"_type": "VariableExpr",
"name": "v2"
},
"type": {
"name": "float"
}
},
{
"location": {
"from": "L11:5",

View File

@@ -72,29 +72,6 @@
}
],
"judgments": [
{
"location": {
"from": "L26:0",
"to": "L26:5"
},
"expr": {
"_type": "VariableExpr",
"name": "print"
},
"type": {
"pos_args": [
{
"pos": 0,
"name": "object",
"type": {},
"required": true
}
],
"args": [],
"kw_args": [],
"returns": {}
}
},
{
"location": {
"from": "L27:4",
@@ -325,6 +302,29 @@
}
}
},
{
"location": {
"from": "L26:0",
"to": "L26:5"
},
"expr": {
"_type": "VariableExpr",
"name": "print"
},
"type": {
"pos_args": [
{
"pos": 0,
"name": "object",
"type": {},
"required": true
}
],
"args": [],
"kw_args": [],
"returns": {}
}
},
{
"location": {
"from": "L26:0",

View File

@@ -63,31 +63,6 @@
"name": "float"
}
},
{
"location": {
"from": "L6:11",
"to": "L6:15"
},
"expr": {
"_type": "VariableExpr",
"name": "bool"
},
"type": {
"pos_args": [
{
"pos": 0,
"name": "object",
"type": {},
"required": false
}
],
"args": [],
"kw_args": [],
"returns": {
"name": "bool"
}
}
},
{
"location": {
"from": "L6:16",
@@ -135,6 +110,31 @@
"name": "int"
}
},
{
"location": {
"from": "L6:11",
"to": "L6:15"
},
"expr": {
"_type": "VariableExpr",
"name": "bool"
},
"type": {
"pos_args": [
{
"pos": 0,
"name": "object",
"type": {},
"required": false
}
],
"args": [],
"kw_args": [],
"returns": {
"name": "bool"
}
}
},
{
"location": {
"from": "L6:11",
@@ -367,6 +367,54 @@
}
}
},
{
"location": {
"from": "L12:21",
"to": "L12:27"
},
"expr": {
"_type": "VariableExpr",
"name": "double"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "float"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "float"
}
}
},
{
"location": {
"from": "L12:29",
"to": "L12:35"
},
"expr": {
"_type": "VariableExpr",
"name": "floats"
},
"type": {
"name": "list",
"args": [
{
"name": "float"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L12:17",
@@ -455,54 +503,6 @@
}
}
},
{
"location": {
"from": "L12:21",
"to": "L12:27"
},
"expr": {
"_type": "VariableExpr",
"name": "double"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "float"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "float"
}
}
},
{
"location": {
"from": "L12:29",
"to": "L12:35"
},
"expr": {
"_type": "VariableExpr",
"name": "floats"
},
"type": {
"name": "list",
"args": [
{
"name": "float"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L12:17",
@@ -538,6 +538,54 @@
}
}
},
{
"location": {
"from": "L13:19",
"to": "L13:25"
},
"expr": {
"_type": "VariableExpr",
"name": "double"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "float"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "float"
}
}
},
{
"location": {
"from": "L13:27",
"to": "L13:31"
},
"expr": {
"_type": "VariableExpr",
"name": "ints"
},
"type": {
"name": "list",
"args": [
{
"name": "int"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L13:15",
@@ -626,54 +674,6 @@
}
}
},
{
"location": {
"from": "L13:19",
"to": "L13:25"
},
"expr": {
"_type": "VariableExpr",
"name": "double"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "float"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "float"
}
}
},
{
"location": {
"from": "L13:27",
"to": "L13:31"
},
"expr": {
"_type": "VariableExpr",
"name": "ints"
},
"type": {
"name": "list",
"args": [
{
"name": "int"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L13:15",
@@ -699,6 +699,54 @@
},
"type": {}
},
{
"location": {
"from": "L14:15",
"to": "L14:21"
},
"expr": {
"_type": "VariableExpr",
"name": "is_odd"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "int"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "bool"
}
}
},
{
"location": {
"from": "L14:23",
"to": "L14:27"
},
"expr": {
"_type": "VariableExpr",
"name": "ints"
},
"type": {
"name": "list",
"args": [
{
"name": "int"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L14:11",
@@ -787,54 +835,6 @@
}
}
},
{
"location": {
"from": "L14:15",
"to": "L14:21"
},
"expr": {
"_type": "VariableExpr",
"name": "is_odd"
},
"type": {
"pos_args": [],
"args": [
{
"pos": 0,
"name": "value",
"type": {
"name": "int"
},
"required": true
}
],
"kw_args": [],
"returns": {
"name": "bool"
}
}
},
{
"location": {
"from": "L14:23",
"to": "L14:27"
},
"expr": {
"_type": "VariableExpr",
"name": "ints"
},
"type": {
"name": "list",
"args": [
{
"name": "int"
}
],
"body": {
"name": "list"
}
}
},
{
"location": {
"from": "L14:11",

View File

@@ -8,6 +8,7 @@ from midas.ast.midas import (
Expr,
ExtendStmt,
ExtensionType,
FrameType,
FunctionType,
GenericType,
GetExpr,
@@ -197,3 +198,15 @@ class MidasAstJsonSerializer(
"base": type.base.accept(self),
"extension": type.extension.accept(self),
}
def visit_frame_type(self, type: FrameType) -> dict:
return {
"_type": "FrameType",
"columns": [self._serialize_column(col) for col in type.columns],
}
def _serialize_column(self, column: FrameType.Column):
return {
"name": column.name.lexeme,
"type": column.type.accept(self),
}