fix: tidy some TODOs

midas/checker/frames/frame_manager.py

@@ -232,7 +232,7 @@ class FrameManager:
             if col.name == name:
                 index = i
                 replace = True
-                # TODO: check column type here to prevent changing it
+                # TODO: might want to check column type here to disallow changing it
             new_columns.append(col)
 
         new_col: DataFrameType.Column = DataFrameType.Column(

midas/checker/midas.py

@@ -24,7 +24,7 @@ from midas.checker.types import (
     TypeVar,
     UnknownType,
 )
-from midas.checker.variance import VarianceInferrer
+from midas.checker.variance import VarianceManager
 from midas.lexer.midas import MidasLexer
 from midas.lexer.token import Token, TokenType
 from midas.parser.midas import MidasParser
@@ -147,10 +147,8 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         for stmt in stmts:
             stmt.accept(self)
 
-        for name, type in self.types._types.items():
+        manager: VarianceManager = VarianceManager(self.types)
-            if isinstance(type, GenericType):
+        manager.infer_all()
-                inferrer = VarianceInferrer(self.types)
-                self.types._types[name] = inferrer.infer(type)
 
     def assert_bool(self, expr: m.Expr):
         """Check that the given expression is a subtype of `bool` or report an error

midas/checker/python.py

@@ -454,7 +454,6 @@ class PythonTyper(
         self.env.define(stmt.name, function)
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> None:
-        # TODO check not yet defined locally
         type: Type = self.resolve_type_expr(stmt.type)
         self.env.define(stmt.name, type)
 
@@ -844,7 +843,7 @@ class PythonTyper(
     def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
         test_type: Type = self.type_of(expr.test)
 
-        # TODO Allow subtypes or any type
+        # Strict: test must be a subtype of bool, or UnknownType
         if (
             not self.is_subtype(test_type, self.types.get_type("bool"))
             and test_type != UnknownType()
@@ -1154,8 +1153,9 @@ class PythonTyper(
                         return False, None
 
                     if key is None:
-                        # TODO: check that value is always a dict
+                        # If literal value is not a dict, invalid Python -> abort
-                        assert isinstance(value_val, dict)
+                        if not isinstance(value_val, dict):
+                            return False, None
                         pairs.extend(value_val.items())
                     else:
                         pairs.append((key_val, value_val))
@@ -1281,9 +1281,7 @@ class PythonTyper(
 
             case BaseType():
                 # TODO: do we want to allow cast(float, int)? would require runtime conversion
-                if not self.types.is_subtype(
+                if not self.types.are_equivalent(subject_type, target_type):
-                    subject_type, target_type
-                ) or not self.types.is_subtype(target_type, subject_type):
                     self.reporter.error(
                         expr.location,
                         f"Value {lit_value!r} of type {subject_type} cannot be cast as {target_type}",

midas/checker/registry.py

@@ -213,7 +213,6 @@ class TypesRegistry:
                 return True
 
             case (ColumnType(type=inner1), ColumnType(type=inner2)):
-                # TODO: invariant, replace ColumnType with simple GenericType
                 if not self.are_equivalent(inner1, inner2):
                     return False
                 return True

midas/checker/variance.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Literal, Optional, cast
 
 from midas.checker.registry import Member, TypesRegistry
@@ -73,10 +75,14 @@ class Tracker:
 class VarianceInferrer:
     """Helper class to compute type parameter variance"""
 
-    def __init__(self, types: TypesRegistry) -> None:
+    def __init__(self, manager: VarianceManager) -> None:
-        self.types: TypesRegistry = types
+        self.manager: VarianceManager = manager
         self.tracker: Tracker = Tracker([])
 
+    @property
+    def types(self) -> TypesRegistry:
+        return self.manager.types
+
     def infer(self, type: GenericType) -> GenericType:
         """Infer the variance of a generic type's parameters
 
@@ -150,11 +156,13 @@ class VarianceInferrer:
             # Get inferred variance of parameters and multiply with current
             # polarity to recurse through arguments
             case AppliedType(name=name, args=args):
-                # TODO: handle mutually recursive types
+                if self.manager.is_in_queue(name):
-                if name == base_name:
                     return
+
                 generic: Type = self.types.get_type(name)
                 assert isinstance(generic, GenericType)
+                generic = self.manager.infer(name, generic)
+
                 params: list[TypeVar] = generic.params
                 polarities: dict[Variance, Polarity] = {
                     Variance.INVARIANT: 0,
@@ -179,3 +187,66 @@ class VarianceInferrer:
             case TypeVar():
                 if type in self.tracker:
                     self.tracker.record(type, polarity)
+
+
+class VarianceManager:
+    """Coordinator for VarianceInferrer to handle recursive types"""
+
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+        self._queue: list[str] = []
+        self._inferred: set[str] = set()
+
+    def infer_all(self):
+        """Infer variance on all generic types defined in the registry"""
+
+        for name, type in self.types._types.items():
+            if isinstance(type, GenericType):
+                self.infer(name, type)
+
+    def infer(self, name: str, type: GenericType) -> GenericType:
+        """Infer variance of parameters of the given type
+
+        Args:
+            name (str): the type's name
+            type (GenericType): the type
+
+        Returns:
+            GenericType: a new generic type with its parameters updated with
+                their inferred variance
+        """
+        if self.is_inferred(name):
+            return type
+
+        self._queue.append(name)
+
+        inferrer: VarianceInferrer = VarianceInferrer(self)
+        inferred: GenericType = inferrer.infer(type)
+        self.types._types[name] = inferred
+
+        self._queue.pop()
+        self._inferred.add(name)
+
+        return inferred
+
+    def is_in_queue(self, name: str) -> bool:
+        """Whether the given type's variance is currently being inferred
+
+        Args:
+            name (str): the type's name
+
+        Returns:
+            bool: whether the type is in the queue
+        """
+        return name in self._queue
+
+    def is_inferred(self, name: str) -> bool:
+        """Whether the given type's variance has already been inferred
+
+        Args:
+            name (str): the type's name
+
+        Returns:
+            bool: whether the type has been processed
+        """
+        return name in self._inferred

midas/generator/generator.py

@@ -269,7 +269,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         alias: ast.expr = self._make_alias(expr.expr, expr2)
 
         type: Type = self._get_expr_type(expr)
-        asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type)
+        asserts: list[ast.stmt] = self._make_cast_asserts(
+            expr.location, alias, type, context=[]
+        )
         for assert_ in asserts:
             self._add_assert(assert_)
 
@@ -352,7 +354,6 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> ast.stmt:
-        # TODO: is that ok?
         return ast.Pass()
 
     def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
@@ -526,7 +527,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         raise RuntimeError(f"Cannot get type judgement for {query}")
 
     def _make_cast_asserts(
-        self, src_location: Location, expr: ast.expr, type: Type
+        self,
+        src_location: Location,
+        expr: ast.expr,
+        type: Type,
+        *,
+        context: list[str],
     ) -> list[ast.stmt]:
         """Generate assertions for the given cast expression
 
@@ -535,6 +541,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 the source file
             expr (ast.expr): the expression being cast
             type (Type): the target type
+            context (list[str]): the current context
 
         Returns:
             list[ast.stmt]: the generated assertion statements
@@ -551,12 +558,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             args=[expr, ast.Name(id=name)],
                             keywords=[],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     )
                 ]
 
             case DerivedType(type=base):
-                return self._make_cast_asserts(src_location, expr, base)
+                return self._make_cast_asserts(
+                    src_location, expr, base, context=context
+                )
 
             case UnitType():
                 return [
@@ -568,19 +579,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 ast.Constant(value=None),
                             ],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     ),
                 ]
 
             case AppliedType(body=body):
-                return self._make_cast_asserts(src_location, expr, body)
+                return self._make_cast_asserts(
+                    src_location, expr, body, context=context
+                )
 
             case ConstraintType(type=base, constraint=constraint):
                 asserts: list[ast.stmt] = self._make_cast_asserts(
-                    src_location, expr, base
+                    src_location, expr, base, context=context
                 )
                 asserts.append(
-                    self._make_constraint_assert(src_location, expr, constraint)
+                    self._make_constraint_assert(
+                        src_location, expr, constraint, context=context
+                    )
                 )
                 return asserts
 
@@ -588,7 +605,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 # TODO: check with type from arguments / use call-site context
                 if bound is None:
                     return []
-                return self._make_cast_asserts(src_location, expr, bound)
+                return self._make_cast_asserts(
+                    src_location, expr, bound, context=context
+                )
 
             case TupleType(items=items):
                 asserts: list[ast.stmt] = [
@@ -598,13 +617,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             args=[expr, ast.Name(id="tuple")],
                             keywords=[],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     ),
                 ]
                 assert isinstance(expr, ast.Tuple)
                 for item, item_type in zip(expr.elts, items):
                     asserts.extend(
-                        self._make_cast_asserts(src_location, item, item_type)
+                        self._make_cast_asserts(
+                            src_location, item, item_type, context=context
+                        )
                     )
                 return asserts
 
@@ -618,7 +641,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             keywords=[],
                         ),
                         self._make_cast_assert_message(
-                            src_location, expr, type, ": Not a dataframe"
+                            src_location,
+                            expr,
+                            type,
+                            context=context,
+                            extra=": Not a dataframe",
                         ),
                     ),
                 ]
@@ -634,7 +661,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 src_location,
                                 expr,
                                 type,
-                                f": Missing column {column.name}",
+                                context=context,
+                                extra=f": Missing column '{column.name}'",
                             ),
                         )
                     )
@@ -645,6 +673,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 value=expr, slice=ast.Constant(value=column.name)
                             ),
                             column.type,
+                            context=context + [f"in column '{column.name}'"],
                         )
                     )
                 return asserts
@@ -659,12 +688,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             keywords=[],
                         ),
                         self._make_cast_assert_message(
-                            src_location, expr, type, ": Not a column"
+                            src_location,
+                            expr,
+                            type,
+                            context=context,
+                            extra=": Not a column",
                         ),
                     ),
                 ]
                 inner_assert: Optional[ast.stmt] = self._make_column_inner_assert(
-                    src_location, expr, type
+                    src_location,
+                    expr,
+                    type,
+                    context,
                 )
                 if inner_assert is not None:
                     asserts.append(inner_assert)
@@ -689,7 +725,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         location: Location,
         expr: ast.expr,
         type: Type,
-        extra: Optional[str] = None,
+        *,
+        context: list[str],
+        extra: str = "",
     ) -> ast.expr:
         """Build an AST node for a cast assertion message
 
@@ -703,12 +741,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 source file
             expr (ast.expr): the expression being cast
             type (Type): the target type
-            extra (Optional[str], optional): extra text to append at the end of
+            context (list[str]): the current context
-                the message. Defaults to None.
+            extra (str, optional): extra text to append at the end of
+                the message. Defaults to "".
 
         Returns:
             ast.expr: the generated message (as an f-string)
         """
+        context_str: str = "".join(map(lambda c: f", {c}", context))
         loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
         # f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
         return ast.JoinedStr(
@@ -725,12 +765,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                     ),
                     conversion=-1,
                 ),
-                ast.Constant(f" to {type}{extra or ''}"),
+                ast.Constant(f" to {type}{context_str}{extra}"),
             ]
         )
 
     def _make_constraint_assert(
-        self, src_location: Location, expr: ast.expr, constraint: m.Expr
+        self,
+        src_location: Location,
+        expr: ast.expr,
+        constraint: m.Expr,
+        *,
+        context: list[str],
     ) -> ast.stmt:
         """Build an assertion for the given constraint on the given expression
 
@@ -739,6 +784,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 source file
             expr (ast.expr): the expression subject to `constraint`
             constraint (m.Expr): the constraint applied on `expr`
+            context (list[str]): the current context
 
         Returns:
             ast.stmt: the assert statement checking the constraint
@@ -750,11 +796,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 args=[expr],
                 keywords=[],
             ),
-            self._make_constraint_assert_message(src_location, constraint),
+            self._make_constraint_assert_message(
+                src_location, constraint, context=context
+            ),
         )
 
     def _make_constraint_assert_message(
-        self, location: Location, constraint: m.Expr
+        self, location: Location, constraint: m.Expr, *, context: list[str]
     ) -> ast.expr:
         """Build an assert message for the given constraint
 
@@ -762,16 +810,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             location (Location): the location of the cast expression in the
                 source file
             constraint (m.Expr): the constraint
+            context (list[str]): the current context
 
         Returns:
             ast.expr: the assert message
         """
         printer = MidasPrinter()
         constraint_str: str = printer.print(constraint)
+        context_str: str = "".join(map(lambda c: f", {c}", context))
         loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
         # f"file.py:L1:1: ConstraintError: Value does not fit constraint 'v > 0'"
         return ast.Constant(
-            f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'"
+            f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'{context_str}"
         )
 
     def _get_constraint(self, expr: m.Expr) -> ast.expr:
@@ -880,7 +930,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     def _make_column_inner_assert(
-        self, src_location: Location, column: ast.expr, type: ColumnType
+        self,
+        src_location: Location,
+        column: ast.expr,
+        type: ColumnType,
+        context: list[str],
     ) -> Optional[ast.stmt]:
         """Build a for-loop checking the type of values inside a column
 
@@ -889,18 +943,21 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 source file
             column (ast.expr): the column being cast
             type (ColumnType): the type of the column
+            context (list[str]): the current context
 
         Returns:
             Optional[ast.stmt]: a for-loop checking the values, or `None` if no
                 assertions are necessary
         """
-        # TODO: improve message, maybe chain contexts
+
-        col: ast.expr = ast.Name(id="col")
+        value: ast.expr = ast.Name(id="value")
-        body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type)
+        body: list[ast.stmt] = self._make_cast_asserts(
+            src_location, value, type.type, context=context
+        )
         if len(body) == 0:
             return None
         return ast.For(
-            target=col,
+            target=value,
             iter=column,
             body=body,
             orelse=[],

midas/generator/stubs.py

@@ -53,7 +53,7 @@ class StubsGenerator:
         self.import_pandas = False
         for name, type in self.types._types.items():
             # Skip builtin types, not just based on name so the user can override
-            # TODO: check if added members on builtin type
+            # TODO: check if added members on builtin type, or prevent it
             match type:
                 case BaseType(name=name_) if name == name_:
                     continue
@@ -105,7 +105,9 @@ class StubsGenerator:
         """
         base_type: Type = type
 
-        # TODO: improve
+        # Generate simple assignment for type aliases
+        # A type alias will have a name that is different from the type represents
+        # or will neither be a `DeriveType` nor a `GenericType`
         match type:
             case DerivedType(name=name_) | GenericType(name=name_) if name_ == name:
                 pass