From f891beb13906cd2fcc33a3d7a88c627628166167 Mon Sep 17 00:00:00 2001 From: Seyon Sivarajah Date: Fri, 28 Jun 2024 12:14:19 +0100 Subject: [PATCH] docs: add hugr.py add namespace in conftest.oy to allow `Dfg` use in HUGR doctests --- hugr-py/src/hugr/conftest.py | 21 +++ hugr-py/src/hugr/hugr.py | 270 ++++++++++++++++++++++++++++++++++- 2 files changed, 284 insertions(+), 7 deletions(-) create mode 100644 hugr-py/src/hugr/conftest.py diff --git a/hugr-py/src/hugr/conftest.py b/hugr-py/src/hugr/conftest.py new file mode 100644 index 000000000..1ad1a1e5c --- /dev/null +++ b/hugr-py/src/hugr/conftest.py @@ -0,0 +1,21 @@ +import pytest +import hugr.hugr as hugr +import hugr.node_port as node_port +import hugr.dfg as dfg +import hugr.ops as ops +import hugr.tys as tys +import hugr.val as val + + +@pytest.fixture(autouse=True) +def add_np(doctest_namespace): + doctest_namespace.update( + { + "hugr": hugr, + "node_port": node_port, + "dfg": dfg, + "ops": ops, + "tys": tys, + "val": val, + } + ) diff --git a/hugr-py/src/hugr/hugr.py b/hugr-py/src/hugr/hugr.py index c8fb5bf95..73c371a82 100644 --- a/hugr-py/src/hugr/hugr.py +++ b/hugr-py/src/hugr/hugr.py @@ -1,3 +1,5 @@ +"""Core data structures for HUGR.""" + from __future__ import annotations from collections.abc import Mapping @@ -26,11 +28,15 @@ @dataclass() class NodeData: + """Node weights in HUGR graph. Defined by an operation and parent node.""" + + #: The operation of the node. op: ops.Op + #: The parent node, or None for the root node. parent: Node | None - _num_inps: int = 0 - _num_outs: int = 0 - children: list[Node] = field(default_factory=list) + _num_inps: int = field(default=0, repr=False) + _num_outs: int = field(default=0, repr=False) + children: list[Node] = field(default_factory=list, repr=False) def to_serial(self, node: Node) -> SerialOp: o = self.op.to_serial(self.parent if self.parent else node) @@ -48,7 +54,11 @@ def to_serial(self, node: Node) -> SerialOp: class ParentBuilder(ToNode, Protocol[OpVar]): + """Abstract interface implemented by builders of nodes that contain child HUGRs.""" + + #: The child HUGR. hugr: Hugr[OpVar] + # Unique parent node. parent_node: Node def to_node(self) -> Node: @@ -56,21 +66,39 @@ def to_node(self) -> Node: @property def parent_op(self) -> OpVar: + """The parent node's operation.""" return cast(OpVar, self.hugr[self.parent_node].op) @dataclass() class Hugr(Mapping[Node, NodeData], Generic[OpVar]): + """The core HUGR datastructure. + + Args: + root_op: The operation for the root node. Defaults to a Module. + + Examples: + >>> h = Hugr() + >>> h.root_op() + Module() + >>> h[h.root].op + Module() + """ + + #: Root node of the HUGR. root: Node + # List of nodes, with None for deleted nodes. _nodes: list[NodeData | None] + # Bidirectional map of links between ports. _links: BiMap[_SO, _SI] + # List of free node indices, populated when nodes are deleted. _free_nodes: list[Node] - def __init__(self, root_op: OpVar) -> None: + def __init__(self, root_op: OpVar | None = None) -> None: self._free_nodes = [] self._links = BiMap() self._nodes = [] - self.root = self._add_node(root_op, None, 0) + self.root = self._add_node(root_op or ops.Module(), None, 0) def __getitem__(self, key: ToNode) -> NodeData: key = key.to_node() @@ -94,6 +122,20 @@ def _get_typed_op(self, node: ToNode, cl: PyType[OpVar2]) -> OpVar2: return op def children(self, node: ToNode | None = None) -> list[Node]: + """The child nodes of a given `node`. + + Args: + node: Parent node. Defaults to the Hugr root. + + Returns: + List of child nodes. + + Examples: + >>> h = Hugr() + >>> n = h.add_node(ops.Const(val.TRUE)) + >>> h.children(h.root) + [Node(1)] + """ node = node or self.root return self[node].children @@ -123,13 +165,55 @@ def add_node( parent: ToNode | None = None, num_outs: int | None = None, ) -> Node: + """Add a node to the HUGR. + + Args: + op: Operation of the node. + parent: Parent node of added node. Defaults to HUGR root if None. + num_outs: Number of output ports expected for this node. Defaults to None. + + Returns: + Handle to the added node. + """ parent = parent or self.root return self._add_node(op, parent, num_outs) def add_const(self, value: val.Value, parent: ToNode | None = None) -> Node: + """Add a constant node to the HUGR. + + Args: + value: Value of the constant. + parent: Parent node of added node. Defaults to HUGR root if None. + + Returns: + Handle to the added node. + + Examples: + >>> h = Hugr() + >>> n = h.add_const(val.TRUE) + >>> h[n].op + Const(TRUE) + """ return self.add_node(ops.Const(value), parent) def delete_node(self, node: ToNode) -> NodeData | None: + """Delete a node from the HUGR. + + Args: + node: Node to delete. + + Returns: + The deleted node data, or None if the node was not found. + + Examples: + >>> h = Hugr() + >>> n = h.add_const(val.TRUE) + >>> deleted = h.delete_node(n) + >>> deleted.op + Const(TRUE) + >>> len(h) + 1 + """ node = node.to_node() parent = self[node].parent if parent: @@ -156,6 +240,19 @@ def _unused_sub_offset(self, port: P) -> _SubPort[P]: return sub_port def add_link(self, src: OutPort, dst: InPort) -> None: + """Add a link (edge) between two nodes to the HUGR, from an outgoing port to an incoming + port. + + Args: + src: Source port. + dst: Destination port. + + Examples: + >>> df = dfg.Dfg(tys.Bool) + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(0)) + >>> list(df.hugr.linked_ports(df.input_node[0])) + [InPort(Node(2), 0)] + """ src_sub = self._unused_sub_offset(src) dst_sub = self._unused_sub_offset(dst) # if self._links.get_left(dst_sub) is not None: @@ -166,6 +263,12 @@ def add_link(self, src: OutPort, dst: InPort) -> None: self[dst.node]._num_inps = max(self[dst.node]._num_inps, dst.offset + 1) def delete_link(self, src: OutPort, dst: InPort) -> None: + """Delete a link (edge) between two nodes from the HUGR. + + Args: + src: Source port. + dst: Destination port. + """ try: sub_offset = next( i for i, inp in enumerate(self.linked_ports(src)) if inp == dst @@ -176,12 +279,45 @@ def delete_link(self, src: OutPort, dst: InPort) -> None: # TODO make sure sub-offset is handled correctly def root_op(self) -> OpVar: + """The operation of the root node. + + Examples: + >>> h = Hugr() + >>> h.root_op() + Module() + """ return cast(OpVar, self[self.root].op) def num_nodes(self) -> int: + """The number of nodes in the HUGR. + + Examples: + >>> h = Hugr() + >>> n = h.add_const(val.TRUE) + >>> h.num_nodes() + 2 + """ return len(self._nodes) - len(self._free_nodes) def num_ports(self, node: ToNode, direction: Direction) -> int: + """The number of ports of a node in a given direction. + Not necessarily the number of connected ports - if port `i` is + connected, then all ports `0..i` are assumed to exist. + + Args: + node: Node to query. + direction: Direction of ports to count. + + Examples: + >>> h = Hugr() + >>> n1 = h.add_const(val.TRUE) + >>> n2 = h.add_const(val.FALSE) + >>> h.add_link(n1.out(0), n2.inp(2)) # not a valid link! + >>> h.num_ports(n1, Direction.OUTGOING) + 1 + >>> h.num_ports(n2, Direction.INCOMING) + 3 + """ return ( self.num_in_ports(node) if direction == Direction.INCOMING @@ -189,9 +325,11 @@ def num_ports(self, node: ToNode, direction: Direction) -> int: ) def num_in_ports(self, node: ToNode) -> int: + """The number of incoming ports of a node. See :meth:`num_ports`.""" return self[node]._num_inps def num_out_ports(self, node: ToNode) -> int: + """The number of outgoing ports of a node. See :meth:`num_ports`.""" return self[node]._num_outs def _linked_ports( @@ -208,6 +346,21 @@ def linked_ports(self, port: OutPort) -> Iterable[InPort]: ... @overload def linked_ports(self, port: InPort) -> Iterable[OutPort]: ... def linked_ports(self, port: OutPort | InPort): + """Return an iterable of In(Out)Ports linked to given Out(In)Port. + + Args: + port: Given port. + + Returns: + Iterator over linked ports. + + Examples: + >>> df = dfg.Dfg(tys.Bool) + >>> df.set_outputs(df.input_node[0]) + >>> list(df.hugr.linked_ports(df.input_node[0])) + [InPort(Node(2), 0)] + + """ match port: case OutPort(_): return self._linked_ports(port, self._links.fwd) @@ -217,9 +370,33 @@ def linked_ports(self, port: OutPort | InPort): # TODO: single linked port def outgoing_order_links(self, node: ToNode) -> Iterable[Node]: + """Iterator over nodes connected by an outgoing state order link from a + given node. + + Args: + node: Source node of state order link. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(-1), df.output_node.inp(-1)) + >>> list(df.hugr.outgoing_order_links(df.input_node)) + [Node(2)] + """ return (p.node for p in self.linked_ports(node.out(-1))) def incoming_order_links(self, node: ToNode) -> Iterable[Node]: + """Iterator over nodes connected by an incoming state order link to a + given node. + + Args: + node: Destination node of state order link. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(-1), df.output_node.inp(-1)) + >>> list(df.hugr.incoming_order_links(df.output_node)) + [Node(1)] + """ return (p.node for p in self.linked_ports(node.inp(-1))) def _node_links( @@ -235,25 +412,86 @@ def _node_links( yield port, list(self._linked_ports(port, links)) def outgoing_links(self, node: ToNode) -> Iterable[tuple[OutPort, list[InPort]]]: + """Iterator over outgoing links from a given node. + + Args: + node: Node to query. + + Returns: + Iterator of pairs of outgoing port and the incoming ports connected + to that port. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(0)) + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(1)) + >>> list(df.hugr.outgoing_links(df.input_node)) + [(OutPort(Node(1), 0), [InPort(Node(2), 0), InPort(Node(2), 1)])] + """ return self._node_links(node, self._links.fwd) def incoming_links(self, node: ToNode) -> Iterable[tuple[InPort, list[OutPort]]]: + """Iterator over incoming links to a given node. + + Args: + node: Node to query. + + Returns: + Iterator of pairs of incoming port and the outgoing ports connected + to that port. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(0)) + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(1)) + >>> list(df.hugr.incoming_links(df.output_node)) + [(InPort(Node(2), 0), [OutPort(Node(1), 0)]), (InPort(Node(2), 1), [OutPort(Node(1), 0)])] + """ return self._node_links(node, self._links.bck) def num_incoming(self, node: Node) -> int: - # connecetd links + """The number of incoming links to a `node`. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(0)) + >>> df.hugr.num_incoming(df.output_node) + 1 + """ return sum(1 for _ in self.incoming_links(node)) def num_outgoing(self, node: ToNode) -> int: - # connecetd links + """The number of outgoing links from a `node`. + + Examples: + >>> df = dfg.Dfg() + >>> df.hugr.add_link(df.input_node.out(0), df.output_node.inp(0)) + >>> df.hugr.num_outgoing(df.input_node) + 1 + """ return sum(1 for _ in self.outgoing_links(node)) # TODO: num_links and _linked_ports def port_kind(self, port: InPort | OutPort) -> tys.Kind: + """The kind of a `port`. + + Examples: + >>> df = dfg.Dfg(tys.Bool) + >>> df.hugr.port_kind(df.input_node.out(0)) + ValueKind(Bool) + """ return self[port.node].op.port_kind(port) def port_type(self, port: InPort | OutPort) -> tys.Type | None: + """The type of a `port`, if the kind is + :class:`ValueKind `, else None. + + Examples: + >>> df = dfg.Dfg(tys.Bool) + >>> df.hugr.port_type(df.input_node.out(0)) + Bool + """ op = self[port.node].op if isinstance(op, ops.DataflowOp): return op.port_type(port) @@ -264,6 +502,22 @@ def port_type(self, port: InPort | OutPort) -> tys.Type | None: return None def insert_hugr(self, hugr: Hugr, parent: ToNode | None = None) -> dict[Node, Node]: + """Insert a HUGR in to this HUGR. + + Args: + hugr: HUGR to insert. + parent: Parent for root of inserted HUGR, defaults to None. + + Returns: + Mapping from node indices in inserted to HUGR to their new indices + in this HUGR. + + Examples: + >>> d = dfg.Dfg() + >>> h = Hugr() + >>> h.insert_hugr(d.hugr) + {Node(0): Node(1), Node(1): Node(2), Node(2): Node(3)} + """ mapping: dict[Node, Node] = {} for idx, node_data in enumerate(hugr._nodes): @@ -285,6 +539,7 @@ def insert_hugr(self, hugr: Hugr, parent: ToNode | None = None) -> dict[Node, No return mapping def to_serial(self) -> SerialHugr: + """Serialise the HUGR.""" node_it = (node for node in self._nodes if node is not None) def _serialise_link( @@ -317,6 +572,7 @@ def _constrain_offset(self, p: P) -> int: @classmethod def from_serial(cls, serial: SerialHugr) -> Hugr: + """Load a HUGR from a serialised form.""" assert serial.nodes, "Empty Hugr is invalid" hugr = Hugr.__new__(Hugr)