summaryrefslogtreecommitdiff
blob: b992bec0a011ec8a23d2fcb5bc24a0f6f0384294 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
--- a/rdflib/store.py
+++ b/rdflib/store.py
@@ -114,26 +114,38 @@ class TripleRemovedEvent(Event):
 
 
+class _Pickler(Pickler):
+    def __init__(self, file, node):
+        super(_Pickler, self).__init__(file)
+        self._node = node
+
+    def persistent_id(self, key: Any) -> Optional[str]:
+        try:
+            return self._node._ids.get(key)
+        except TypeError:
+            return None
+
+
+class _Unpickler(Unpickler):
+    def __init__(self, file, node):
+        super(_Unpickler, self).__init__(file)
+        self._node = node
+
+    def persistent_load(self, pid):
+        return self._node._get_object(pid)
+
+
 class NodePickler:
     def __init__(self) -> None:
         self._objects: Dict[str, Any] = {}
         self._ids: Dict[Any, str] = {}
         self._get_object = self._objects.__getitem__
 
-    def _get_ids(self, key: Any) -> Optional[str]:
-        try:
-            return self._ids.get(key)
-        except TypeError:
-            return None
-
     def register(self, object: Any, id: str) -> None:
         self._objects[id] = object
         self._ids[object] = id
 
     def loads(self, s: bytes) -> Node:
-        up = Unpickler(BytesIO(s))
-        # NOTE on type error: https://github.com/python/mypy/issues/2427
-        # type error: Cannot assign to a method
-        up.persistent_load = self._get_object  # type: ignore[assignment]
+        up = _Unpickler(BytesIO(s), self)
         try:
             return up.load()
         except KeyError as e:
@@ -143,10 +155,7 @@ class NodePickler:
         self, obj: "Node", protocol: Optional[Any] = None, bin: Optional[Any] = None
     ):
         src = BytesIO()
-        p = Pickler(src)
-        # NOTE on type error: https://github.com/python/mypy/issues/2427
-        # type error: Cannot assign to a method
-        p.persistent_id = self._get_ids  # type: ignore[assignment]
+        p = _Pickler(src, self)
         p.dump(obj)
         return src.getvalue()