feat(hugr-py): only require input type annotations when building

hugr-py/src/hugr/_cfg.py

@@ -84,9 +84,10 @@ def entry(self) -> Node:
         return self._entry_block.root
 
     def _entry_op(self) -> ops.DataflowBlock:
-        dop = self.hugr[self.entry].op
-        assert isinstance(dop, ops.DataflowBlock)
-        return dop
+        return self.hugr._get_typed_op(self.entry, ops.DataflowBlock)
+
+    def _exit_op(self) -> ops.ExitBlock:
+        return self.hugr._get_typed_op(self.exit, ops.ExitBlock)
 
     def add_entry(self, sum_rows: Sequence[TypeRow], other_outputs: TypeRow) -> Block:
         # update entry block types

hugr-py/src/hugr/_dfg.py

@@ -55,14 +55,10 @@ def new_nested(cls, root_op: DP, hugr: Hugr, parent: ToNode | None = None) -> Se
         return new
 
     def _input_op(self) -> ops.Input:
-        dop = self.hugr[self.input_node].op
-        assert isinstance(dop, ops.Input)
-        return dop
+        return self.hugr._get_typed_op(self.input_node, ops.Input)
 
     def _output_op(self) -> ops.Output:
-        dop = self.hugr[self.output_node].op
-        assert isinstance(dop, ops.Output)
-        return dop
+        return self.hugr._get_typed_op(self.output_node, ops.Output)
 
     def root_op(self) -> DP:
         return cast(DP, self.hugr[self.root].op)

hugr-py/src/hugr/_hugr.py

@@ -12,6 +12,7 @@
     TypeVar,
     cast,
     overload,
+    Type as PyType,
 )
 
 from typing_extensions import Self
@@ -131,6 +132,7 @@ def to_serial(self, node: Node, hugr: Hugr) -> SerialOp:
 
 P = TypeVar("P", InPort, OutPort)
 K = TypeVar("K", InPort, OutPort)
+OpVar = TypeVar("OpVar", bound=Op)
 
 
 @dataclass(frozen=True, eq=True, order=True)
@@ -183,6 +185,11 @@ def __iter__(self):
     def __len__(self) -> int:
         return self.num_nodes()
 
+    def _get_typed_op(self, node: ToNode, cl: PyType[OpVar]) -> OpVar:
+        op = self[node].op
+        assert isinstance(op, cl)
+        return op
+
     def children(self, node: ToNode | None = None) -> list[Node]:
         node = node or self.root
         return self[node].children