From 03a9506176e7811d33dac2cb492d37336dd5d6ba Mon Sep 17 00:00:00 2001
From: Jeff Newman <jeffnewman@camsys.com>
Date: Thu, 8 Apr 2021 10:17:40 -0500
Subject: [PATCH] fix copy and pickle so they work with freezing

---
 addict/addict.py |  8 +++++++-
 test_addict.py   | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/addict/addict.py b/addict/addict.py
index 55e02d1..3a7c88e 100644
--- a/addict/addict.py
+++ b/addict/addict.py
@@ -119,10 +119,16 @@ def __getnewargs__(self):
         return tuple(self.items())
 
     def __getstate__(self):
-        return self
+        state = self.to_dict()
+        isFrozen = (hasattr(self, '__frozen') and
+                    object.__getattribute__(self, '__frozen'))
+        state['__addict__frozen__'] = isFrozen
+        return state
 
     def __setstate__(self, state):
+        shouldFreeze = state.pop('__addict__frozen__', False)
         self.update(state)
+        self.freeze(shouldFreeze)
 
     def __or__(self, other):
         if not isinstance(other, (Dict, dict)):
diff --git a/test_addict.py b/test_addict.py
index b4af646..2b96bf4 100644
--- a/test_addict.py
+++ b/test_addict.py
@@ -394,6 +394,38 @@ def __init__(self):
         b.child = "new stuff"
         self.assertTrue(isinstance(a.child, self.dict_class))
 
+    def test_copy_with_freeze(self):
+        class MyMutableObject(object):
+
+            def __init__(self):
+                self.attribute = None
+
+        foo = MyMutableObject()
+        foo.attribute = True
+
+        a = self.dict_class()
+        a.child.immutable = 42
+        a.child.mutable = foo
+        a.freeze()
+
+        b = a.copy()
+
+        # immutable object should not change
+        b.child.immutable = 21
+        self.assertEqual(a.child.immutable, 21)
+
+        # mutable object should change
+        b.child.mutable.attribute = False
+        self.assertEqual(a.child.mutable.attribute, b.child.mutable.attribute)
+
+        # changing child of b should not affect a
+        b.child = "new stuff"
+        self.assertTrue(isinstance(a.child, self.dict_class))
+
+        # b should be frozen
+        with self.assertRaises(KeyError):
+            b.missing
+
     def test_deepcopy(self):
         class MyMutableObject(object):
             def __init__(self):
@@ -450,6 +482,16 @@ def test_pickle(self):
         a = self.dict_class(TEST_DICT)
         self.assertEqual(a, pickle.loads(pickle.dumps(a)))
 
+    def test_pickle_with_freeze(self):
+        a = self.dict_class(TEST_DICT)
+        a.freeze()
+        aa = pickle.loads(pickle.dumps(a))
+        self.assertEqual(a, aa)
+        with self.assertRaises(KeyError):
+            a.missing
+        with self.assertRaises(KeyError):
+            aa.missing
+
     def test_add_on_empty_dict(self):
         d = self.dict_class()
         d.x.y += 1