Make mypy happy

Author: zrho

hugr-py/src/hugr/model/__init__.py

@@ -76,8 +76,8 @@ def from_str(s: str) -> "Param":
 @dataclass
 class Symbol:
     name: str
-    params: list[Param] = field(default_factory = list)
-    constraints: list[Term] = field(default_factory = list)
+    params: Sequence[Param] = field(default_factory = list)
+    constraints: Sequence[Term] = field(default_factory = list)
     signature: Term = field(default_factory = Wildcard)
 
     def __str__(self):
@@ -150,10 +150,10 @@ class Import(Op):
 @dataclass
 class Node:
     operation: Op = field(default_factory=lambda: InvalidOp())
-    inputs: list[str] = field(default_factory=list)
-    outputs: list[str] = field(default_factory=list)
-    regions: list["Region"] = field(default_factory=list)
-    meta: list[Term] = field(default_factory=list)
+    inputs: Sequence[str] = field(default_factory=list)
+    outputs: Sequence[str] = field(default_factory=list)
+    regions: Sequence["Region"] = field(default_factory=list)
+    meta: Sequence[Term] = field(default_factory=list)
     signature: Optional[Term] = None
 
     def __str__(self) -> str:
@@ -171,10 +171,10 @@ class ScopeClosure(Enum):
 @dataclass
 class Region:
     kind: RegionKind = RegionKind.DATA_FLOW
-    sources: list[str] = field(default_factory=list)
-    targets: list[str] = field(default_factory=list)
-    children: list[Node] = field(default_factory=list)
-    meta: list[Term] = field(default_factory=list)
+    sources: Sequence[str] = field(default_factory=list)
+    targets: Sequence[str] = field(default_factory=list)
+    children: Sequence[Node] = field(default_factory=list)
+    meta: Sequence[Term] = field(default_factory=list)
     signature: Optional[Term] = None
 
     def __str__(self):

hugr-py/src/hugr/model/export.py

@@ -4,7 +4,7 @@
 from hugr.ops import DFG, Input, Output, Custom, AsExtOp, Conditional, TailLoop, FuncDefn, FuncDecl, Call, CallIndirect, LoadFunc, AliasDecl, AliasDefn, LoadConst, Const, CFG, ExitBlock, DataflowBlock, Tag
 from hugr.hugr.node_port import InPort, OutPort
 from hugr.tys import FunctionKind, ConstKind, TypeParam, Type, TypeTypeParam, TypeBound
-from typing import cast, Sequence
+from typing import cast, Sequence, TypeVar, Generic
 
 class ModelExport:
     def __init__(self, hugr: Hugr):
@@ -127,15 +127,15 @@ def export_node(self, node: Node) -> Optional[model.Node]:
                     signature = model.Apply("core.type")
                 )
 
-                value = cast(model.Term, op.definition.to_model())
+                alias_value = cast(model.Term, op.definition.to_model())
 
                 return model.Node(
-                    operation = model.DefineAlias(symbol, value)
+                    operation = model.DefineAlias(symbol, alias_value)
                 )
 
             case Call() as op:
-                input_types = model.List([type.to_model() for type in op.instantiation.input ])
-                output_types = model.List([type.to_model() for type in op.instantiation.output ])
+                input_types = [type.to_model() for type in op.instantiation.input ]
+                output_types = [type.to_model() for type in op.instantiation.output ]
                 signature = op.instantiation.to_model()
                 func_args = cast(list[model.Term], [type.to_model() for type in op.type_args])
                 func_name = self.find_func_input(node)
@@ -146,15 +146,12 @@ def export_node(self, node: Node) -> Optional[model.Node]:
                 func = model.Apply(func_name, func_args)
 
                 return model.Node(
-                    operation = model.CustomOp(model.Apply(
-                        "core.call",
-                        [
-                            input_types,
-                            output_types,
-                            model.ExtSet(),
-                            func
-                        ]
-                    )),
+                    operation = model.CustomOp(model.Apply("core.call", [
+                        model.List(input_types),
+                        model.List(output_types),
+                        model.ExtSet(),
+                        func
+                    ])),
                     signature = signature,
                     inputs = inputs,
                     outputs = outputs,
@@ -246,22 +243,26 @@ def export_node(self, node: Node) -> Optional[model.Node]:
             case DataflowBlock() as op:
                 region = self.export_region_dfg(node)
 
-                input_types = model.List([
+                input_types = [
                     model.Apply("core.ctrl", [
                         model.List([ type.to_model() for type in op.inputs ])
                     ])
-                ])
+                ]
 
                 other_output_types = [ type.to_model() for type in op.other_outputs ]
-                output_types = model.List([
+                output_types = [
                     model.Apply("core.ctrl", [model.List([
                         *[type.to_model() for type in row],
                         *other_output_types
                     ])])
                     for row in op.sum_ty.variant_rows
-                ])
+                ]
 
-                signature = model.Apply("core.fn", [input_types, output_types, model.ExtSet()])
+                signature = model.Apply("core.fn", [
+                    model.List(input_types),
+                    model.List(output_types),
+                    model.ExtSet()
+                ])
 
                 return model.Node(
                     operation = model.Block(),
@@ -317,8 +318,8 @@ def export_region_module(self, node: Node) -> model.Region:
     def export_region_dfg(self, node: Node) -> model.Region:
         node_data = self.hugr[node]
         children: list[model.Node] = []
-        source_types = model.Wildcard()
-        target_types = model.Wildcard()
+        source_types: model.Term = model.Wildcard()
+        target_types: model.Term = model.Wildcard()
         sources = []
         targets = []
 
@@ -355,8 +356,8 @@ def export_region_cfg(self, node: Node) -> model.Region:
 
         source = None
         targets = []
-        source_types = model.Wildcard()
-        target_types = model.Wildcard()
+        source_types: model.Term = model.Wildcard()
+        target_types: model.Term = model.Wildcard()
         children = []
 
         for child in node_data.children:
@@ -464,8 +465,10 @@ def _mangle_name(node: Node, name: str) -> str:
     # by adding the node id.
     return f"_{name}_{node.idx}"
 
-class UnionFind[T]:
-    def __init__(self):
+T = TypeVar('T')
+
+class UnionFind(Generic[T]):
+    def __init__(self) -> None:
         self.parents: dict[T, T] = {}
         self.sizes: dict[T, int] = {}
 

hugr-py/src/hugr/tys.py

@@ -682,11 +682,12 @@ def __eq__(self, value):
         return super().__eq__(value)
 
     def to_model(self) -> model.Term:
-        args = [arg.to_model() for arg in self.args]
-
-        # TODO: This cast is only neccessary because `Type` can both be an
+        # This cast is only neccessary because `Type` can both be an
         # actual type or a row variable.
-        args = cast(list[model.Term], args)
+        args = [
+            cast(model.Term, arg.to_model())
+            for arg in self.args
+        ]
 
         extension_name = self.type_def.get_extension().name
         type_name = self.type_def.name
@@ -739,11 +740,12 @@ def __str__(self) -> str:
         return _type_str(self.id, self.args)
 
     def to_model(self) -> model.Term:
-        args = [arg.to_model() for arg in self.args]
-
-        # TODO: This cast is only neccessary because `Type` can both be an
+        # This cast is only neccessary because `Type` can both be an
         # actual type or a row variable.
-        args = cast(list[model.Term], args)
+        args = [
+            cast(model.Term, arg.to_model())
+            for arg in self.args
+        ]
 
         return model.Apply(self.id, args)