Resolve pickling and deepcopying bug with dynamically added traits

traits/has_traits.py

@@ -1372,6 +1372,9 @@ def __getstate__(self):
                              for name in self.trait_names( type = 'delegate',
                                                            transient = False )
                              if name in dic ] ) )
+        # Add all instance traits
+        inst_traits = self._instance_traits()
+        result['__instance_traits__'] = inst_traits
 
         # If this object implements ISerializable, make sure that all
         # contained HasTraits objects in its persisted state also implement
@@ -1411,7 +1414,10 @@ def __setstate__ ( self, state, trait_change_notify = True ):
         else:
             # Otherwise, apply the Traits 3.0 restore logic:
             self._init_trait_listeners()
+            inst_traits = state.pop('__instance_traits__', {})
             self.trait_set( trait_change_notify = trait_change_notify, **state )
+            for attr in inst_traits:
+                self.add_trait(attr, inst_traits[attr])
             self._post_init_trait_listeners()
             self.traits_init()
 
@@ -1779,6 +1785,9 @@ def clone_traits ( self, traits = None, memo = None, copy = None,
         new = self.__new__( self.__class__ )
         memo[ id( self ) ] = new
         new._init_trait_listeners()
+        inst_traits = self._instance_traits()
+        for attr in inst_traits:
+            new.add_trait(attr, inst_traits[attr])
         new.copy_traits( self, traits, memo, copy, **metadata )
         new._post_init_trait_listeners()
         new.traits_init()

traits/tests/test_regression.py

@@ -13,6 +13,7 @@
     from traits.trait_numeric import Array
 
 from traits.has_traits import HasTraits, Property, on_trait_change
+from traits.trait_errors import TraitError
 from traits.trait_types import Bool, DelegatesTo, Either, Instance, Int, List
 from traits.testing.unittest_tools import unittest
 
@@ -211,6 +212,31 @@ def test_delegation_refleak(self):
         # All the counts should be the same.
         self.assertEqual(counts[warmup:-1], counts[warmup+1:])
 
+    def test_hastraits_deepcopy(self):
+        # Regression test for enthought/traits#2 and enthought/traits#16
+        from copy import deepcopy
+        a = HasTraits()
+        a.add_trait('foo', Int)
+        a.foo = 1
+        with self.assertRaises(TraitError):
+            a.foo = 'a'
+        copied_a = deepcopy(a)
+        with self.assertRaises(TraitError):
+            copied_a.foo = 'a'
+
+    def test_hastraits_pickle(self):
+        # Regression test for enthought/traits#2 and enthought/traits#16
+        from pickle import dumps, loads
+        a = HasTraits()
+        a.add_trait('foo', Int)
+        a.foo = 1
+        with self.assertRaises(TraitError):
+            a.foo = 'a'
+        pickled_a = dumps(a)
+        unpickled_a = loads(pickled_a)
+        with self.assertRaises(TraitError):
+            unpickled_a.foo = 'a'
+
     @unittest.skipUnless(numpy_available, "test requires NumPy")
     def test_exception_from_numpy_comparison_ignored(self):
         # Regression test for enthought/traits#376.