core: (ContainerOf) improve python typing

xdslproject · Jan 29, 2025 · 258e8af · 258e8af
1 parent fa315a7
commit 258e8af
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/xdsl/dialects/builtin.py b/xdsl/dialects/builtin.py
@@ -1094,13 +1094,21 @@ def get_element_type(self) -> AttributeCovT:
 
 
 @dataclass(frozen=True, init=False)
-class ContainerOf(AttrConstraint):
+class ContainerOf(
+    Generic[AttributeCovT],
+    GenericAttrConstraint[
+        AttributeCovT | VectorType[AttributeCovT] | TensorType[AttributeCovT]
+    ],
+):
     """A type constraint that can be nested once in a vector or a tensor."""
 
-    elem_constr: AttrConstraint
+    elem_constr: GenericAttrConstraint[AttributeCovT]
 
     def __init__(
-        self, elem_constr: Attribute | type[Attribute] | AttrConstraint
+        self,
+        elem_constr: Attribute
+        | type[AttributeCovT]
+        | GenericAttrConstraint[AttributeCovT],
     ) -> None:
         object.__setattr__(self, "elem_constr", attr_constr_coercion(elem_constr))