diff --git a/src/pgraph/PGraph.py b/src/pgraph/PGraph.py index 8727ef07..2edbaac7 100644 --- a/src/pgraph/PGraph.py +++ b/src/pgraph/PGraph.py @@ -16,6 +16,36 @@ from spatialmath.base.graphics import axes_logic +def _L1(v: NDArray) -> float: + """L1 (Manhattan) norm, the built-in ``"L1"`` metric""" + return float(np.linalg.norm(v, 1)) + + +def _L2(v: NDArray) -> float: + """L2 (Euclidean) norm, the built-in ``"L2"`` metric""" + return float(np.linalg.norm(v)) + + +def _SE2(v: NDArray) -> float: + """Mixed norm for (x, y, theta), the built-in ``"SE2"`` metric""" + if len(v) != 3: + raise ValueError( + f"SE2 metric requires a 3-element (x, y, theta) vector, got length {len(v)}" + ) + # wrap angle to range [-pi, pi) + v[2] = (v[2] + np.pi) % (2 * np.pi) - np.pi + return float(np.linalg.norm(v)) + + +#: the built-in metrics by name. They are module level, and pickled by name, +#: so that a graph using one can be pickled. +_METRICS: dict[str, Callable[[NDArray], float]] = { + "L1": _L1, + "L2": _L2, + "SE2": _SE2, +} + + class BaseGraph(ABC): #: concrete vertex type for this graph kind, provided by :class:`UGraph` @@ -509,6 +539,34 @@ def copy(self) -> BaseGraph: :return: deep copy :rtype: BaseGraph + The copy has its own vertices and edges, which keep their class, name, + coordinate, cost and any attributes the user has added, and in the + same order, so it can be modified without affecting the original. + + A graph can also be pickled, and copied with :func:`copy.deepcopy`, + which this method calls. Both work for a graph of any size. A + vertex or an edge cannot be pickled or copied on its own, since it has + no meaning apart from its graph: a :exc:`TypeError` is raised, and the + graph should be copied instead. + + .. note:: + + - A metric or heuristic given by name (``"L1"``, ``"L2"``, + ``"SE2"``) is pickled by name. A function you supply must be + importable at module level to be pickled; a ``lambda`` or a + function defined inside another function can be copied but not + pickled. + - Attributes that you add to vertices or edges must not refer to + other vertices or edges of the graph, since those cannot be + pickled on their own. + - A pickle is not a long-term storage format. Use :meth:`export` + for that. + + .. versionchanged:: 1.0 + Pickling a graph works, and copying one no longer fails with a + ``RecursionError`` for a graph with a path of about a hundred + vertices. + .. runblock:: pycon >>> from pgraph import UGraph @@ -517,9 +575,98 @@ def copy(self) -> BaseGraph: >>> g2 = g.copy() >>> g2 is g >>> g2[0] is g[0] + + :seealso: :meth:`vertex_copy` :meth:`export` """ return copy.deepcopy(self) + # pickling and copying + + #: the graph's own containers, which are saved as a flat description + _STRUCTURE = ("_vertexlist", "_vertexdict", "_edges") + + #: the references a vertex holds, which are rebuilt rather than saved + _VERTEX_LINKS = ("_edgelist", "_inlist", "_graph") + + def __getstate__(self) -> dict[str, Any]: + """ + State for pickling and deep copying (private method) + + The vertices and edges form a web of references, and following them + recursively overflows the stack for a graph with a path of a hundred + or so vertices. Instead the graph is described flatly: a list of + vertices, and a list of edges that refer to their end vertices by + index. The edges are listed in the order they were added, so + replaying them rebuilds each vertex's edge lists in their original + order. A built-in metric is saved by name. + """ + state = {k: v for k, v in self.__dict__.items() if k not in self._STRUCTURE} + for key in ("_metric", "_heuristic"): + for name, func in _METRICS.items(): + if state[key] is func: + state[key] = ("builtin", name) + index = {id(v): i for i, v in enumerate(self._vertexlist)} + state["_flat"] = { + "version": 1, + "vertices": [ + ( + type(v), + {k: x for k, x in v.__dict__.items() if k not in self._VERTEX_LINKS}, + ) + for v in self._vertexlist + ], + "edges": [ + ( + type(e), + index[id(e.v1)], + index[id(e.v2)], + {k: x for k, x in e.__dict__.items() if k not in ("_v1", "_v2")}, + ) + for e in self._edges + ], + } + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """ + Rebuild the graph from its flat description (private method) + + The vertex and edge objects are made without calling ``__init__``, so + subclasses with their own constructor signature are restored as they + were. Each edge is attached with the vertex's own hook, exactly as + :meth:`BaseVertex.connect` does. + """ + state = dict(state) + flat = state.pop("_flat") + if flat["version"] != 1: + raise ValueError(f"unknown pickled graph version {flat['version']!r}") + for key in ("_metric", "_heuristic"): + value = state[key] + if isinstance(value, tuple) and value[0] == "builtin": + state[key] = _METRICS[value[1]] + self.__dict__.update(state) + + self._vertexlist = [] + self._vertexdict = {} + for cls, attrs in flat["vertices"]: + v = cls.__new__(cls) + v.__dict__.update(attrs) + v._edgelist = [] + v._inlist = [] + v._graph = self + self._vertexlist.append(v) + if v._name is not None: + self._vertexdict[v._name] = v + + self._edges = {} + for cls, i, j, attrs in flat["edges"]: + e = cls.__new__(cls) + e.__dict__.update(attrs) + v1, v2 = self._vertexlist[i], self._vertexlist[j] + e._set_endpoints(v1, v2) + self._edges[e] = None + v1._attach(e, v2) + def add_vertex( self, coord: ArrayLike | BaseVertex | None = None, name: str | None = None ) -> BaseVertex: @@ -1063,30 +1210,11 @@ def _metricfunc(self, metric: Callable[[NDArray], float] | str) -> Callable[[NDA :seealso: :meth:`metric` :meth:`heuristic` """ - def L1(v): - return np.linalg.norm(v, 1) - - def L2(v): - return np.linalg.norm(v) - - def SE2(v): - if len(v) != 3: - raise ValueError( - f"SE2 metric requires a 3-element (x, y, theta) vector, got length {len(v)}" - ) - # wrap angle to range [-pi, pi) - v[2] = (v[2] + np.pi) % (2 * np.pi) - np.pi - return np.linalg.norm(v) - if callable(metric): return metric elif isinstance(metric, str): - if metric == "L1": - return L1 - elif metric == "L2": - return L2 - elif metric == "SE2": - return SE2 + if metric in _METRICS: + return _METRICS[metric] else: raise ValueError(f"unknown metric {metric!r}") else: @@ -3030,11 +3158,23 @@ def _set_endpoints(self, v1: BaseVertex | None, v2: BaseVertex | None) -> None: Only the library calls this: the constructor's own assignment, then :meth:`connect`, :meth:`BaseVertex.connect` for an edge built without - vertices, and :meth:`_BaseGraph.remove_edge`, which clears them. + vertices, and :meth:`BaseGraph.remove_edge`, which clears them. """ self._v1 = v1 self._v2 = v2 + def __reduce__(self) -> Any: + """ + Refuse to pickle or copy an edge on its own + + An edge is held by its vertices and its graph, so it has no meaning + apart from them. Pickle or copy the graph instead. + """ + raise TypeError( + "an Edge cannot be pickled or copied on its own, " + "pickle or copy its graph instead" + ) + def __repr__(self) -> str: """ Detailed representation of the edge @@ -3291,6 +3431,18 @@ def __init__(self, coord: ArrayLike | None = None, name: str | None = None): self._graph: BaseGraph | None = None # reference to owning graph # print('BaseVertex init', type(self)) + def __reduce__(self) -> Any: + """ + Refuse to pickle or copy a vertex on its own + + A vertex is held by its graph and by its edges, so it has no meaning + apart from them. Pickle or copy the graph instead. + """ + raise TypeError( + "a vertex cannot be pickled or copied on its own, " + "pickle or copy its graph instead" + ) + def __str__(self) -> str: """ Compact representation of the vertex diff --git a/tests/test_edge_endpoints.py b/tests/test_edge_endpoints.py index ed6a84f5..a31add97 100644 --- a/tests/test_edge_endpoints.py +++ b/tests/test_edge_endpoints.py @@ -106,15 +106,18 @@ def __init__(self, v1, v2, uv): a.connect(b, edge=o) self.assertEqual((o.v1, o.v2, o.uv), (a, b, (1, 2))) - def test_repr_and_copies_keep_the_endpoints(self): + def test_repr_shows_the_endpoints_and_an_edge_alone_is_not_copied(self): g = UGraph() a, b = g.add_vertex(name="a"), g.add_vertex(name="b") e = a.connect(b, cost=2) self.assertIn("a", str(e)) - for dup in (copy.copy(e), copy.deepcopy(e)): - self.assertEqual(dup.cost, 2) - self.assertIsNotNone(dup.v1) - self.assertIsNotNone(dup.v2) + # an edge has no meaning apart from its graph: copy the graph instead + for f in (copy.copy, copy.deepcopy): + with self.assertRaises(TypeError): + f(e) + dup = g.copy().edges()[0] + self.assertEqual(dup.cost, 2) + self.assertEqual((dup.v1.name, dup.v2.name), ("a", "b")) if __name__ == "__main__": diff --git a/tests/test_pickle.py b/tests/test_pickle.py new file mode 100644 index 00000000..17955723 --- /dev/null +++ b/tests/test_pickle.py @@ -0,0 +1,261 @@ +"""Pickling and copying a graph. + +A graph is a web of mutual references, which pickle and deepcopy follow +recursively: before the graph described itself flatly, ``copy()`` overflowed +the stack on a path of about 100 vertices, and a graph using a built-in metric +could not be pickled at all because the metric was a local function. +""" + +import copy +import pickle # nosec B403 nosemgrep - the tests pickle their own graphs +import unittest + +import numpy as np + +from pgraph import DGraph, DVertex, Edge, UGraph, UVertex + + +def dumps(obj, protocol=None): + return pickle.dumps(obj, protocol) # nosec nosemgrep + + +def loads(data): + return pickle.loads(data) # nosec B301 nosemgrep + + +def manhattan(v): + return float(np.abs(v).sum()) + + +class TaggedVertex(UVertex): + def __init__(self, tag, coord=None, name=None): # a different signature + super().__init__(coord=coord, name=name) + self.tag = tag + + +class TaggedDVertex(DVertex): + def __init__(self, tag, coord=None, name=None): + super().__init__(coord=coord, name=name) + self.tag = tag + + +class WeightedEdge(Edge): + def __init__(self, v1=None, v2=None, cost=None, data=None): + super().__init__(v1, v2, cost=cost, data=data) + self.weight = 7 + + +def chain(n, cls=UGraph): + g = cls() + vs = [g.add_vertex([i, 0]) for i in range(n)] + for a, b in zip(vs, vs[1:]): + g.add_edge(a, b) + return g + + +def snapshot(g): + """Everything observable about a graph, by name, including order""" + return { + "kind": type(g), + "vertices": [(v.name, type(v), v.coord.tolist() if v.coord is not None else None) for v in g], + "edges": [(e.v1.name, e.v2.name, e.cost, type(e)) for e in g.edges()], + "adjacent": {v.name: [w.name for w in v.adjacent()] for v in g}, + "out": {v.name: [(e.v1.name, e.v2.name) for e in v.out_edges()] for v in g}, + "in": {v.name: [(e.v1.name, e.v2.name) for e in v.in_edges()] for v in g}, + } + + +class _RoundTrip: + cls: type + + def make(self): + g = self.cls(dim=2) + a, b, c, d = (g.add_vertex([i, i * i], name=n) for i, n in enumerate("abcd")) + g.add_edge(a, b) + g.add_edge(b, c) + g.add_edge(a, b, cost=9.0) # parallel + g.add_edge(c, c) # loop + g.add_edge(d, a) + g.add_edge(c, a) + g.remove_edge(g.edges()[1]) # leave a gap in the order + g.add_edge(b, c, cost=2.5) + a.data = {"note": "hello"} + return g + + def check(self, g, h): + self.assertIsNot(g, h) + self.assertEqual(snapshot(g), snapshot(h)) + self.assertEqual(h.validate(), []) + self.assertEqual((g.n, g.ne), (h.n, h.ne)) + for v in g: + w = h[v.name] + self.assertIsNot(v, w) + self.assertIs(w._graph, h) + self.assertEqual((v.degree, v.in_degree, v.out_degree), (w.degree, w.in_degree, w.out_degree)) + for e in h.edges(): + self.assertIs(e.v1._graph, h) + self.assertIn(e, e.v1.out_edges()) + np.testing.assert_array_equal(g.adjacency(), h.adjacency()) + np.testing.assert_array_equal(g.incidence(), h.incidence()) + + def test_pickle(self): + g = self.make() + self.check(g, loads(dumps(g))) + + def test_pickle_every_protocol(self): + g = self.make() + for protocol in range(2, pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(protocol=protocol): + self.check(g, loads(dumps(g, protocol))) + + def test_copy(self): + g = self.make() + h = g.copy() + self.check(g, h) + self.assertIsNot(h["a"].coord, g["a"].coord) + self.assertEqual(h["a"].data, {"note": "hello"}) + self.assertIsNot(h["a"].data, g["a"].data) + + def test_copy_is_independent(self): + g = self.make() + h = g.copy() + h.remove_vertex(h["a"]) + h.add_vertex([5, 5], name="z") + self.assertEqual(g.n, 4) + self.assertIn("a", [v.name for v in g]) + self.assertEqual(g.validate(), []) + self.assertEqual(h.validate(), []) + + def test_restored_graph_is_usable(self): + h = loads(dumps(self.make())) + e = h.add_edge(h["d"], h["c"]) + self.assertEqual(h.ne, self.make().ne + 1) + h.remove_edge(e) + v = h.add_vertex([9, 9]) + self.assertEqual(v.name, "#4") + self.assertEqual(h.validate(), []) + path, length, _ = h.path_Astar(h["a"], h["c"]) + self.assertEqual(path[0], h["a"]) + + def test_settings_survive(self): + g = self.cls(metric="L1", heuristic="SE2", dim=3, verbose=False) + h = loads(dumps(g)) + self.assertEqual(h.metric([3, -4, 0]), 7) + with self.assertRaises(ValueError): + h.add_vertex([1, 2]) # dim is still 3 + self.assertIs(h.heuristic, g.heuristic) + + def test_builtin_metrics_pickle_by_name(self): + for name, value in (("L1", 7.0), ("L2", 5.0)): + with self.subTest(metric=name): + g = self.cls(metric=name, heuristic=name) + h = loads(dumps(g)) + self.assertEqual(h.metric(np.r_[3.0, -4.0]), value) + self.assertIs(h.metric, g.metric) + self.assertIs(h.heuristic, g.heuristic) + + def test_default_metric_pickles(self): + g = self.cls() + h = loads(dumps(g)) + self.assertAlmostEqual(h.metric(np.r_[3.0, 4.0]), 5.0) + + def test_se2_metric_pickles(self): + g = self.cls(metric="SE2") + h = loads(dumps(g)) + self.assertIs(h.metric, g.metric) + + def test_user_module_level_metric(self): + g = self.cls(metric=manhattan) + h = loads(dumps(g)) + self.assertIs(h.metric, manhattan) + + def test_lambda_metric_cannot_be_pickled_but_can_be_copied(self): + g = self.cls(metric=lambda v: 1.0) + with self.assertRaises((pickle.PicklingError, AttributeError)): + dumps(g) + self.assertEqual(g.copy().metric(np.r_[1.0]), 1.0) + + def test_empty_graph(self): + g = self.cls() + h = loads(dumps(g)) + self.assertEqual((h.n, h.ne), (0, 0)) + + def test_long_chain_pickles_and_copies(self): + g = chain(5000, self.cls) + for h in (loads(dumps(g)), g.copy(), copy.deepcopy(g)): + self.assertEqual((h.n, h.ne), (5000, 4999)) + self.assertEqual(h[2500].adjacent()[0].name, g[2500].adjacent()[0].name) + self.assertEqual(h.validate(), []) + + def test_vertex_alone_is_refused(self): + g = self.make() + for f in (dumps, copy.deepcopy, copy.copy): + with self.subTest(f=f.__name__): + with self.assertRaises(TypeError): + f(g["a"]) + + def test_edge_alone_is_refused(self): + g = self.make() + for f in (dumps, copy.deepcopy, copy.copy): + with self.subTest(f=f.__name__): + with self.assertRaises(TypeError): + f(g.edges()[0]) + + def test_unknown_version_is_refused(self): + state = self.make().__getstate__() + state["_flat"]["version"] = 99 + g = self.cls.__new__(self.cls) + with self.assertRaises(ValueError): + g.__setstate__(state) + + def test_data_referring_to_a_vertex_is_refused(self): + g = self.make() + g["a"].data = g["b"] + with self.assertRaises(TypeError): + dumps(g) + + +class TestUGraphPickle(_RoundTrip, unittest.TestCase): + cls = UGraph + + def test_loop_listed_twice_at_vertex(self): + h = loads(dumps(self.make())) + loop = [e for e in h.edges() if e.v1 is e.v2][0] + self.assertEqual(h["c"].edges().count(loop), 2) + + +class TestDGraphPickle(_RoundTrip, unittest.TestCase): + cls = DGraph + + def test_in_edge_order_survives(self): + g = self.make() + h = loads(dumps(g)) + for v in g: + self.assertEqual( + [(e.v1.name, e.v2.name, e.cost) for e in v.in_edges()], + [(e.v1.name, e.v2.name, e.cost) for e in h[v.name].in_edges()], + ) + + +class TestSubclasses(unittest.TestCase): + def test_vertex_and_edge_classes(self): + for gcls, vcls in ((UGraph, TaggedVertex), (DGraph, TaggedDVertex)): + with self.subTest(graph=gcls.__name__): + g = gcls(vertex_class=vcls, edge_class=WeightedEdge) + a = g.add_vertex([0, 0], name="a") + b = g.add_vertex([1, 1], name="b") + a.tag, b.tag = "first", "second" + e = g.add_edge(a, b) + e.weight = 11 + for h in (loads(dumps(g)), g.copy()): + self.assertIsInstance(h["a"], vcls) + self.assertEqual((h["a"].tag, h["b"].tag), ("first", "second")) + self.assertIsInstance(h.edges()[0], WeightedEdge) + self.assertEqual(h.edges()[0].weight, 11) + self.assertIs(h.vertex_class, vcls) + self.assertIs(h.edge_class, WeightedEdge) + self.assertEqual(h.validate(), []) + + +if __name__ == "__main__": + unittest.main()