diff --git a/src/pgraph/PGraph.py b/src/pgraph/PGraph.py index 168bcefb..82898b7c 100644 --- a/src/pgraph/PGraph.py +++ b/src/pgraph/PGraph.py @@ -10,7 +10,7 @@ import tempfile import subprocess import webbrowser -from typing import Any, Callable, ClassVar +from typing import Any, Callable, ClassVar, overload from numpy.typing import ArrayLike, NDArray from spatialmath.base.graphics import axes_logic @@ -416,19 +416,43 @@ def _require_cost(self, edge: Edge) -> float: ) return edge.cost - def add_edge(self, v1: BaseVertex | str, v2: BaseVertex | str, **kwargs: Any) -> Edge: + @overload + def add_edge(self, v1: Edge, **kwargs: Any) -> Edge: ... + + @overload + def add_edge( + self, v1: BaseVertex | str, v2: BaseVertex | str, **kwargs: Any + ) -> Edge: ... + + def add_edge( + self, + v1: Edge | BaseVertex | str, + v2: BaseVertex | str | None = None, + **kwargs: Any, + ) -> Edge: """ Add an edge to the graph (base class method) - :param v1: first vertex (start if a directed graph) - :type v1: BaseVertex subclass or str - :param v2: second vertex (end if a directed graph) - :type v2: BaseVertex subclass or str - :param kwargs: optional arguments to pass to ``BaseVertex.connect`` + :param v1: first vertex (start if a directed graph), or an existing + edge that already carries its two vertices + :type v1: BaseVertex subclass, str or Edge + :param v2: second vertex (end if a directed graph), must be omitted + if ``v1`` is an edge + :type v2: BaseVertex subclass or str, optional + :param kwargs: optional arguments to pass to ``BaseVertex.connect``, + not allowed if ``v1`` is an edge + :raises TypeError: ``v1`` or ``v2`` is of the wrong type, or ``v2``, + ``cost`` or ``data`` is given along with an edge + :raises ValueError: the edge has no vertices, or the vertices do not + belong to this graph :return: edge :rtype: Edge - Create an edge between a vertex pair and adds it to the graph. + - ``g.add_edge(v1, v2)`` creates an edge between a vertex pair and + adds it to the graph. + - ``g.add_edge(e)`` takes an edge, or an instance of an ``Edge`` + subclass, that was constructed with its two vertices and adds it to + the graph, just as ``g.add_vertex(v)`` takes an existing vertex. This is a graph centric way of creating an edge. The alternative is the ``connect`` method of a vertex. @@ -445,14 +469,53 @@ def add_edge(self, v1: BaseVertex | str, v2: BaseVertex | str, **kwargs: Any) -> >>> e2 = g.add_edge('v2', 'v3', cost=99) >>> print(e2) - :seealso: :meth:`Edge.connect` :meth:`BaseVertex.connect` - """ - v1 = self._resolve_vertex(v1, "v1") - v2 = self._resolve_vertex(v2, "v2") + Adding an instance of a user-defined ``Edge`` subclass: + + .. runblock:: pycon + + >>> from pgraph import UGraph, Edge + >>> class MyEdge(Edge): + ... pass + >>> g = UGraph() + >>> v1 = g.add_vertex(coord=[0,0], name='v1') + >>> v2 = g.add_vertex(coord=[1,1], name='v2') + >>> e = g.add_edge(MyEdge(v1, v2, data="a road")) + >>> print(type(e).__name__, e.data) + + .. note:: To supply an edge that does not yet have vertices use + ``g.add_edge(v1, v2, edge=e)`` or ``e.connect(v1, v2)``. + + :seealso: :meth:`add_vertex` :meth:`Edge.connect` :meth:`BaseVertex.connect` + """ + if isinstance(v1, Edge): + if v2 is not None: + raise TypeError("v2 must not be given when adding an Edge") + if kwargs: + raise TypeError( + "add_edge(edge) takes no other arguments, got " + + ", ".join(kwargs) + + " -- set them on the edge itself" + ) + edge = v1 + if edge.v1 is None or edge.v2 is None: + raise ValueError( + "edge does not have both vertices; use " + "add_edge(v1, v2, edge=e) or e.connect(v1, v2)" + ) + start, end = edge.v1, edge.v2 + kwargs = {"edge": edge} + else: + if v2 is None: + raise TypeError("add_edge() requires two vertices, or an Edge") + start = self._resolve_vertex(v1, "v1") + end = self._resolve_vertex(v2, "v2") + + if start._graph is not self or end._graph is not self: + raise ValueError("vertices must belong to this graph") if self._verbose: - print(f"New edge from {v1.name} to {v2.name}") - return v1.connect(v2, **kwargs) + print(f"New edge from {start.name} to {end.name}") + return start.connect(end, **kwargs) def remove_edge(self, edge: Edge) -> None: """ diff --git a/tests/test_add_edge_edge.py b/tests/test_add_edge_edge.py new file mode 100644 index 00000000..a9ac5a75 --- /dev/null +++ b/tests/test_add_edge_edge.py @@ -0,0 +1,109 @@ +""" +``graph.add_edge()`` accepts either two vertices (or names) or an ``Edge`` +that already carries its vertices, mirroring ``add_vertex()`` which accepts +coordinates or a vertex. +""" + +import unittest + +from pgraph import DGraph, Edge, UGraph, UVertex + + +class MyEdge(Edge): + pass + + +class _AddEdge: + Graph: type + + def setUp(self): + self.g = self.Graph() + self.a = self.g.add_vertex([0, 0], name="a") + self.b = self.g.add_vertex([3, 4], name="b") + self.c = self.g.add_vertex([6, 8], name="c") + + def test_two_vertices_unchanged(self): + e = self.g.add_edge(self.a, self.b) + self.assertIs(type(e), Edge) + self.assertEqual((e.v1, e.v2), (self.a, self.b)) + self.assertEqual(e.cost, 5.0) + + def test_two_names_unchanged(self): + e = self.g.add_edge("a", "b", cost=9) + self.assertEqual((e.v1, e.v2, e.cost), (self.a, self.b, 9)) + + def test_edge_subclass(self): + e = MyEdge(self.a, self.b, cost=2, data="hello") + r = self.g.add_edge(e) + + self.assertIs(r, e) + self.assertIsInstance(r, MyEdge) + self.assertEqual(self.g.ne, 1) + self.assertIn(e, self.g.edges()) + self.assertIs(self.a.edgeto(self.b), e) + self.assertEqual(e.data, "hello") + + def test_edge_cost_computed_at_construction(self): + r = self.g.add_edge(MyEdge(self.a, self.b)) + self.assertEqual(r.cost, 5.0) + + def test_edge_must_carry_vertices(self): + with self.assertRaises(ValueError) as cm: + self.g.add_edge(MyEdge(cost=1)) + self.assertIn("edge=", str(cm.exception)) # points at the alternative + with self.assertRaises(ValueError): + self.g.add_edge(MyEdge(self.a)) # only one end + self.assertEqual(self.g.ne, 0) + + def test_edge_and_second_vertex_is_an_error(self): + with self.assertRaises(TypeError): + self.g.add_edge(MyEdge(self.a, self.b), self.c) + self.assertEqual(self.g.ne, 0) + + def test_edge_with_cost_or_data_kwargs_is_an_error(self): + # connect() would silently ignore them when given an edge + e = MyEdge(self.a, self.b) + with self.assertRaises(TypeError): + self.g.add_edge(e, cost=3) + with self.assertRaises(TypeError): + self.g.add_edge(e, data="x") + self.assertEqual(self.g.ne, 0) + + def test_wrong_type(self): + with self.assertRaises(TypeError): + self.g.add_edge(42, self.b) + with self.assertRaises(TypeError): + self.g.add_edge(self.a, 42) + + def test_vertices_of_another_graph(self): + h = self.Graph() + x = h.add_vertex([0, 0], name="x") + y = h.add_vertex([1, 1], name="y") + with self.assertRaises(ValueError): + self.g.add_edge(MyEdge(x, y)) + with self.assertRaises(ValueError): + self.g.add_edge(x, y) + self.assertEqual(self.g.ne, 0) + self.assertEqual(h.ne, 0) + + def test_edge_between_ungraphed_vertices(self): + v, w = UVertex([0, 0]), UVertex([1, 1]) + with self.assertRaises(ValueError): + self.g.add_edge(MyEdge(v, w)) + + +class TestUGraphAddEdge(_AddEdge, unittest.TestCase): + Graph = UGraph + + +class TestDGraphAddEdge(_AddEdge, unittest.TestCase): + Graph = DGraph + + def test_direction_follows_edge(self): + self.g.add_edge(MyEdge(self.b, self.a)) + self.assertEqual(self.b.neighbours(), [self.a]) + self.assertEqual(self.a.neighbours(), []) + + +if __name__ == "__main__": + unittest.main()