Resolve pickling and deepcopying bug with dynamically added traits

traits/has_traits.py

@@ -1386,6 +1386,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
@@ -1425,9 +1428,12 @@ 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 )
             self._post_init_trait_listeners()
             self.traits_init()
+            for attr, trait in inst_traits.iteritems():
+                self.add_trait(attr, trait)
 
         self.traits_inited( True )
 
@@ -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, trait in inst_traits.iteritems():
+            new.add_trait(attr, trait)
         new.copy_traits( self, traits, memo, copy, **metadata )
         new._post_init_trait_listeners()
         new.traits_init()

traits/tests/test_regression.py

@@ -3,6 +3,7 @@
 import sys
 
 from ..has_traits import HasTraits, Property, on_trait_change
+from ..trait_errors import TraitError
 from ..trait_types import Bool, DelegatesTo, Instance, Int, List
 from ..testing.unittest_tools import unittest
 
@@ -189,6 +190,22 @@ def test_delegation_refleak(self):
         # All the counts should be the same.
         self.assertEqual(counts[warmup:-1], counts[warmup+1:])
 
+    def test_hastraits_pickle_deepcopy(self):
+        from pickle import dumps, loads
+        from copy import deepcopy
+        a = HasTraits()
+        a.add_trait('foo', Int)
+        a.foo = 1
+        with self.assertRaises(TraitError):
+            a.foo = 'a'
+        pkld_a = dumps(a)
+        unpkld_a = loads(pkld_a)
+        with self.assertRaises(TraitError):
+            unpkld_a.foo = 'a'
+        copied_a = deepcopy(a)
+        with self.assertRaises(TraitError):
+            copied_a.foo = 'a'
+
 
 if __name__ == '__main__':
     unittest.main()