From 72610201b0ac4ec5b38fdf42d8316a29074c2c39 Mon Sep 17 00:00:00 2001
From: Nikita Sokolov <faucct@yandex-team.ru>
Date: Mon, 12 Apr 2021 17:30:48 +0300
Subject: [PATCH] Python 3.8 support: Only dicts can be assigned to
 Pickler#memo

---
 src/main/python/ipystate/impl/memo.py     | 27 -----------------------
 src/main/python/ipystate/serialization.py |  8 +++----
 2 files changed, 4 insertions(+), 31 deletions(-)

diff --git a/src/main/python/ipystate/impl/memo.py b/src/main/python/ipystate/impl/memo.py
index a745117..e5ec68e 100644
--- a/src/main/python/ipystate/impl/memo.py
+++ b/src/main/python/ipystate/impl/memo.py
@@ -1,5 +1,3 @@
-from typing import Dict
-
 class ChunkedFile:
     def __init__(self):
         self._ba = bytearray()
@@ -12,28 +10,3 @@ def current_chunk(self) -> bytearray:
 
     def reset(self) -> None:
         self._ba = bytearray()
-
-
-class TransactionalDict(dict):
-    def __init__(self, original: Dict):
-        self._dirty = set()
-        self._prev = {}
-        super().__init__(original)
-
-    def __setitem__(self, key, value) -> None:
-        if key not in self._dirty:
-            self._dirty.add(key)
-            self._prev[key] = super().get(key, None)
-        super().__setitem__(key, value)
-
-    def rollback(self) -> None:
-        for k, v in self._prev.items():
-            if v is None:
-                super().__delitem__(k)
-            else:
-                super().__setitem__(k, v)
-        self.commit()
-
-    def commit(self) -> None:
-        self._dirty = set()
-        self._prev = {}
diff --git a/src/main/python/ipystate/serialization.py b/src/main/python/ipystate/serialization.py
index be2bb25..8fa917f 100644
--- a/src/main/python/ipystate/serialization.py
+++ b/src/main/python/ipystate/serialization.py
@@ -10,7 +10,7 @@
 
 from .impl.components_fuser import ComponentsFuser
 from .impl.dispatch.common import CommonDispatcher
-from .impl.memo import ChunkedFile, TransactionalDict
+from .impl.memo import ChunkedFile
 
 
 class Dump:
@@ -194,7 +194,7 @@ def _dump_pickle_component(self, component: Set[str], ns: Dict[str, object]) ->
 
         cf = ChunkedFile()
         pickler = Pickler(ns, self.configurable_dispatch_table, cf)
-        pickler.memo = TransactionalDict(pickler.memo)
+        committed_memo = pickler.memo.copy()
 
         comp_sorted_vars = self._sort_component_vars(component, ns)
         for var_name in comp_sorted_vars:
@@ -203,11 +203,11 @@ def _dump_pickle_component(self, component: Set[str], ns: Dict[str, object]) ->
             # TODO  and report to non serialized
             try:
                 pickler.dump(var_value)
-                pickler.memo.commit()
+                committed_memo = pickler.memo.copy()
                 chunk = cf.current_chunk()
                 serialized_vars.append((var_name, chunk))
             except Exception as e:
-                pickler.memo.rollback()
+                pickler.memo = committed_memo
                 non_serialized_var_names.add(var_name)
                 self._on_var_serialize_error(var_name, var_value, e)
             finally: