Merge pull request #31 from petercorke/feat/graph-vertex-edge-class

feat: UGraph/DGraph take vertex_class and edge_class

Author: petercorke

src/pgraph/PGraph.py

@@ -29,6 +29,8 @@ def __init__(
         heuristic: Callable[[NDArray], float] | str | None = None,
         verbose: bool = False,
         dim: int | None = None,
+        vertex_class: type[BaseVertex] | None = None,
+        edge_class: type[Edge] | None = None,
     ):
         """
         Abstract base class for graphs
@@ -44,15 +46,53 @@ def __init__(
             None (unconstrained -- vertices may have coordinates of any
             length, or none at all)
         :type dim: int, optional
+        :param vertex_class: class of the vertices this graph creates, defaults
+            to ``UVertex`` for a :class:`UGraph` and ``DVertex`` for a
+            :class:`DGraph`
+        :type vertex_class: subclass of ``UVertex`` or ``DVertex``, optional
+        :param edge_class: class of the edges this graph creates, defaults to
+            :class:`Edge`
+        :type edge_class: subclass of ``Edge``, optional
         :raises ValueError: ``dim`` is given but is not a positive integer
+        :raises TypeError: ``vertex_class`` is not a subclass of this graph's
+            own vertex kind, or ``edge_class`` is not a subclass of ``Edge``
 
         This is the common base class of :class:`UGraph` and :class:`DGraph`
         and should not be instantiated directly.
 
+        ``vertex_class`` and ``edge_class`` let user subclasses be built by the
+        graph, so they never need to be constructed by hand::
+
+            g = UGraph(vertex_class=MyVertex, edge_class=MyEdge)
+            v = g.add_vertex([0, 0])  # a MyVertex
+            e = g.add_edge(v, w)      # a MyEdge
+
+        The classes are called as ``vertex_class(coord, name=name)`` and
+        ``edge_class(v1, v2, cost=cost, data=data)``, so a subclass that
+        overrides ``__init__`` must accept and pass on those arguments.
+        A graph given a ``vertex_class`` only accepts vertices of that class
+        in :meth:`add_vertex`, and one given an ``edge_class`` only accepts
+        edges of that class.
+
         :seealso: :class:`UGraph` :class:`DGraph` :meth:`add_vertex`
+            :meth:`add_edge` :attr:`vertex_class` :attr:`edge_class`
         """
         if dim is not None and dim <= 0:
             raise ValueError(f"dim must be a positive integer, got {dim!r}")
+        kind = type(self)._vertex_cls
+        if vertex_class is None:
+            vertex_class = kind
+        elif not (isinstance(vertex_class, type) and issubclass(vertex_class, kind)):
+            raise TypeError(
+                f"vertex_class must be a subclass of {kind.__name__} for a "
+                f"{type(self).__name__}, got {vertex_class!r}"
+            )
+        if edge_class is None:
+            edge_class = Edge
+        elif not (isinstance(edge_class, type) and issubclass(edge_class, Edge)):
+            raise TypeError(f"edge_class must be a subclass of Edge, got {edge_class!r}")
+        self._vertex_class: type[BaseVertex] = vertex_class
+        self._edge_class: type[Edge] = edge_class
         # we use a list and a dict, the list respects the order of adding
         self._vertexlist: list[BaseVertex] = []
         self._vertexdict: dict[str, BaseVertex] = {}
@@ -100,13 +140,17 @@ def __repr__(self) -> str:
         return str(self)
 
     @classmethod
-    def Dict(cls, d: dict, reverse: bool = False) -> _BaseGraph:
+    def Dict(cls, d: dict, reverse: bool = False, **kwargs: Any) -> _BaseGraph:
         """
         Create graph from parent/child dictionary
 
         :param d: dictionary that maps from ``BaseVertex`` subclass to ``BaseVertex`` subclass
         :type d: dict
         :param reverse: reverse link direction, defaults to False
+        :param kwargs: optional arguments passed to the graph constructor, for
+            example ``vertex_class`` or ``edge_class``
+        :raises TypeError: ``vertex_class`` or ``edge_class`` is not a
+            suitable class, see :meth:`_BaseGraph.__init__`
         :return: graph
         :rtype: UGraph or DGraph
 
@@ -127,7 +171,7 @@ def Dict(cls, d: dict, reverse: bool = False) -> _BaseGraph:
         :seealso: :meth:`Adjacency`
         """
 
-        g = cls()
+        g = cls(**kwargs)
 
         for vertex, parent in d.items():
             if isinstance(vertex, str):
@@ -162,6 +206,7 @@ def Adjacency(
         A: NDArray,
         coords: NDArray | None = None,
         names: list[str] | None = None,
+        **kwargs: Any,
     ) -> _BaseGraph:
         """
         Create graph from adjacency matrix
@@ -172,14 +217,25 @@ def Adjacency(
         :type coords: ndarray(N,M), optional
         :param names: names of vertices, defaults to None
         :type names: list(N) of str, optional
+        :param kwargs: optional arguments passed to the graph constructor, for
+            example ``vertex_class`` or ``edge_class``
+        :raises ValueError: ``A`` is not square, the length of ``names`` or
+            the number of rows of ``coords`` does not match ``A``, or ``A``
+            has a non-zero diagonal element and the graph is a ``DGraph``
+        :raises TypeError: ``vertex_class`` or ``edge_class`` is not a
+            suitable class, see :meth:`_BaseGraph.__init__`
         :return: graph
         :rtype: UGraph or DGraph
 
         Create a directed or undirected graph where non-zero elements ``A[i,j]``
         correspond to edges from vertex ``i`` to vertex ``j``.
 
         .. warning:: For undirected graph ``A`` should be symmetric but this
-            is not checked.  Only the upper triangular part is used.
+            is not checked.  Only the strictly upper triangular part is used,
+            so the diagonal is ignored and no self-loops are created. A
+            directed graph rejects a non-zero diagonal with ``ValueError``.
+            Edge counts greater than one are not turned into parallel edges;
+            a non-zero element is a single edge with that element as its cost.
 
         .. runblock:: pycon
 
@@ -199,7 +255,7 @@ def Adjacency(
         if coords is not None and coords.shape[0] != A.shape[0]:
             raise ValueError("coords must have same number of rows as adjacency matrix")
 
-        g = cls()
+        g = cls(**kwargs)
 
         name = None
         coord = None
@@ -227,6 +283,30 @@ def Adjacency(
 
         return g
 
+    @property
+    def vertex_class(self) -> type[BaseVertex]:
+        """
+        Class of the vertices this graph creates
+
+        :return: vertex class
+        :rtype: subclass of ``UVertex`` or ``DVertex``
+
+        :seealso: :meth:`add_vertex` :attr:`edge_class`
+        """
+        return self._vertex_class
+
+    @property
+    def edge_class(self) -> type[Edge]:
+        """
+        Class of the edges this graph creates
+
+        :return: edge class
+        :rtype: subclass of ``Edge``
+
+        :seealso: :meth:`add_edge` :attr:`vertex_class`
+        """
+        return self._edge_class
+
     def copy(self) -> _BaseGraph:
         """
         Deepcopy of graph
@@ -257,26 +337,29 @@ def add_vertex(
         :type coord: array-like or BaseVertex subclass, optional
         :param name: name of vertex, defaults to "#i"
         :type name: str, optional
-        :raises TypeError: ``coord`` is a ``BaseVertex`` of the wrong kind
+        :raises TypeError: ``coord`` is a ``BaseVertex`` that is not an
+            instance of this graph's :attr:`vertex_class`
         :raises ValueError: the graph was constructed with ``dim``, and
             ``coord`` is given but its length doesn't match
         :return: the added vertex
         :rtype: BaseVertex subclass
 
         - ``g.add_vertex()`` creates a new vertex with optional ``coord`` and
           ``name``.
-        - ``g.add_vertex(v)`` takes an instance or subclass of this graph's
-          own vertex kind and adds it to the graph.
+        - ``g.add_vertex(v)`` takes an instance of this graph's
+          :attr:`vertex_class` (by default any ``UVertex`` for a
+          :class:`UGraph`, any ``DVertex`` for a :class:`DGraph`) and adds it
+          to the graph.
 
         If the vertex has no name and ``name`` is None give it a default name
         ``#N`` where ``N`` is a consecutive integer.
 
         The vertex is placed into a dictionary with a key equal to its name.
 
         This single implementation, shared by :class:`UGraph` and
-        :class:`DGraph`, is parameterized by each subclass's
-        :attr:`_vertex_cls` rather than duplicated per subclass -- see
-        :doc:`policy` for why that matters.
+        :class:`DGraph`, is parameterized by the graph's :attr:`vertex_class`
+        rather than duplicated per subclass -- see :doc:`policy` for why that
+        matters.
 
         .. runblock:: pycon
 
@@ -301,15 +384,15 @@ def add_vertex(
 
         :seealso: :meth:`vertex_copy`
         """
-        if isinstance(coord, self._vertex_cls):
+        if isinstance(coord, self._vertex_class):
             vertex = coord
         elif isinstance(coord, BaseVertex):
             raise TypeError(
-                f"expecting {self._vertex_cls.__name__} or coordinate data, "
+                f"expecting {self._vertex_class.__name__} or coordinate data, "
                 f"got {type(coord).__name__}"
             )
         else:
-            vertex = self._vertex_cls(coord, name=name)
+            vertex = self._vertex_class(coord, name=name)
 
         if (
             self._dim is not None
@@ -2297,6 +2380,11 @@ class UGraph(_BaseGraph):
     """
     Class for undirected graphs
 
+    Vertices are :class:`UVertex` objects, or instances of the subclass given
+    as ``vertex_class``. Edges are :class:`Edge` objects, or instances of the
+    subclass given as ``edge_class``; see :class:`_BaseGraph` for the
+    constructor arguments and the exceptions they raise.
+
     .. inheritance-diagram:: UGraph
 
     :seealso: :class:`_BaseGraph` :class:`DGraph`
@@ -2351,6 +2439,11 @@ class DGraph(_BaseGraph):
     """
     Class for directed graphs
 
+    Vertices are :class:`DVertex` objects, or instances of the subclass given
+    as ``vertex_class``. Edges are :class:`Edge` objects, or instances of the
+    subclass given as ``edge_class``; see :class:`_BaseGraph` for the
+    constructor arguments and the exceptions they raise.
+
     .. inheritance-diagram:: DGraph
 
     :seealso: :class:`_BaseGraph` :class:`UGraph`
@@ -2451,6 +2544,9 @@ class Edge:
           edgelist of _each_ vertex.
         - This class can be inherited to provide user objects with graph capability.
         - Inheritance is an alternative to providing arbitrary user data.
+        - Pass the subclass as ``edge_class`` when constructing the graph and
+          the graph will create instances of it itself, see
+          :class:`_BaseGraph`.
 
     An Edge points to a pair of vertices.  At ``connect`` time the vertices
     get references back to the Edge object.
@@ -2571,6 +2667,8 @@ def connect(self, v1: BaseVertex, v2: BaseVertex) -> None:
         :type v2: BaseVertex subclass
         :raises ValueError: either vertex has not been added to a graph, or
             the vertices belong to different graphs
+        :raises TypeError: this edge is not an instance of the graph's
+            :attr:`_BaseGraph.edge_class`
 
         The edge connects vertices ``v1`` and ``v2``, and is added to the
         graph that those vertices belong to.
@@ -2953,9 +3051,10 @@ def connect(
 
         :param dest: The vertex to connect to
         :type dest: ``BaseVertex`` subclass
-        :param edge: Use this as the edge object, otherwise a new ``Edge``
-                     object is created from the vertices being connected,
-                     and the ``cost`` and ``edge`` parameters, defaults to None
+        :param edge: Use this as the edge object, otherwise a new object of the
+                     graph's :attr:`_BaseGraph.edge_class` is created from the
+                     vertices being connected, and the ``cost`` and ``data``
+                     parameters, defaults to None
         :type edge: ``Edge`` subclass, optional
         :param cost: the cost to traverse this edge, defaults to None
         :type cost: float, optional
@@ -2965,6 +3064,8 @@ def connect(
         :raises ValueError: either vertex has not been added to a graph, the
             vertices belong to different graphs, or ``edge`` already connects
             other vertices
+        :raises TypeError: ``edge`` is not an instance of the graph's
+            :attr:`_BaseGraph.edge_class`
         :return: the edge connecting the vertices
         :rtype: Edge
 
@@ -3002,10 +3103,15 @@ def connect(
         elif self._graph is not dest._graph:
             raise ValueError("vertices must belong to the same graph")
         elif isinstance(edge, Edge):
+            if not isinstance(edge, self._graph._edge_class):
+                raise TypeError(
+                    f"expecting an edge of class {self._graph._edge_class.__name__}, "
+                    f"got {type(edge).__name__}"
+                )
             e = edge
             self._bind_edge(e, dest)
         else:
-            e = Edge(self, dest, cost=cost, data=data)
+            e = self._graph._edge_class(self, dest, cost=cost, data=data)
 
         self._graph._edgelist.add(e)
         self._graph._connectivitychange = True
@@ -3336,8 +3442,9 @@ def connect(
 
         :param dest: vertex to connect to
         :type dest: BaseVertex subclass
-        :param edge: Use this as the edge object, otherwise a new ``Edge``
-                     object is created, defaults to None
+        :param edge: Use this as the edge object, otherwise a new object of the
+                     graph's :attr:`_BaseGraph.edge_class` is created,
+                     defaults to None
         :type edge: ``Edge`` subclass, optional
         :param cost: the cost to traverse this edge, defaults to None
         :type cost: float, optional
@@ -3347,6 +3454,8 @@ def connect(
         :raises ValueError: either vertex has not been added to a graph, the
             vertices belong to different graphs, or ``edge`` already connects
             other vertices
+        :raises TypeError: ``edge`` is not an instance of the graph's
+            :attr:`_BaseGraph.edge_class`
         :return: the edge connecting the vertices
         :rtype: Edge
 
@@ -3412,8 +3521,9 @@ def connect(
 
         :param dest: vertex to connect to
         :type dest: BaseVertex subclass
-        :param edge: Use this as the edge object, otherwise a new ``Edge``
-                     object is created, defaults to None
+        :param edge: Use this as the edge object, otherwise a new object of the
+                     graph's :attr:`_BaseGraph.edge_class` is created,
+                     defaults to None
         :type edge: ``Edge`` subclass, optional
         :param cost: the cost to traverse this edge, defaults to None
         :type cost: float, optional
@@ -3423,6 +3533,8 @@ def connect(
         :raises ValueError: either vertex has not been added to a graph, the
             vertices belong to different graphs, or ``edge`` already connects
             other vertices
+        :raises TypeError: ``edge`` is not an instance of the graph's
+            :attr:`_BaseGraph.edge_class`
         :return: the edge connecting the vertices
         :rtype: Edge