Feature: Add support for restricted cuba key in DataContainer

simphony/core/data_container.py

@@ -1,13 +1,14 @@
-from simphony.core.cuba import CUBA
+import sys as _sys
 
-_CUBA_MEMBERS = CUBA.__members__
+from simphony.core.cuba import CUBA
 
 
 class DataContainer(dict):
     """ A DataContainer instance
 
     The DataContainer object is implemented as a python dictionary whose keys
-    are restricted to be members of the CUBA enum class.
+    are restricted to the instance's `restricted_keys`, default to the CUBA
+    enum members.
 
     The data container can be initialized like a typical python dict
     using the mapping and iterables where the keys are CUBA enum members.
@@ -22,66 +23,140 @@ class DataContainer(dict):
     # Memory usage optimization.
     __slots__ = ()
 
+    # These are the allowed CUBA keys (faster to convert to set for lookup)
+    restricted_keys = frozenset(CUBA)
+
+    # Map CUBA enum name to CUBA enum
+    # Used by assigning key using keyword name
+    _restricted_mapping = CUBA.__members__
+
     def __init__(self, *args, **kwargs):
         """ Constructor.
         Initialization follows the behaviour of the python dict class.
         """
-        self._check_arguments(args, kwargs)
-        if len(args) == 1 and not hasattr(args[0], 'keys'):
-            super(DataContainer, self).__init__()
-            for key, value in args[0]:
-                self.__setitem__(key, value)
-        elif len(args) == 1:
-            mapping = args[0]
-            if not isinstance(mapping, DataContainer):
-                if any(not isinstance(key, CUBA) for key in mapping):
-                    non_cuba_keys = [
-                        key for key in mapping if not isinstance(key, CUBA)]
-                    message = \
-                        "Key(s) {!r} are not in the approved CUBA keywords"
-                    raise ValueError(message.format(non_cuba_keys))
-            super(DataContainer, self).__init__(mapping)
-        super(DataContainer, self).update(
-            {CUBA[kwarg]: value for kwarg, value in kwargs.viewitems()})
+
+        super(DataContainer, self).__init__()
+        self.update(*args, **kwargs)
 
     def __setitem__(self, key, value):
         """ Set/Update the key value only when the key is a CUBA key.
 
         """
-        if isinstance(key, CUBA):
+        if isinstance(key, CUBA) and key in self.restricted_keys:
             super(DataContainer, self).__setitem__(key, value)
         else:
-            message = "Key {!r} is not in the approved CUBA keywords"
+            message = "Key {!r} is not in the supported CUBA keywords"
             raise ValueError(message.format(key))
 
     def update(self, *args, **kwargs):
         self._check_arguments(args, kwargs)
-        if len(args) == 1 and not hasattr(args[0], 'keys'):
+
+        if args and not hasattr(args[0], 'keys'):
+            # args is an iterator
             for key, value in args[0]:
-                self.__setitem__(key, value)
-        elif len(args) == 1:
+                self[key] = value
+        elif args:
             mapping = args[0]
-            if not isinstance(mapping, DataContainer):
-                if any(not isinstance(key, CUBA) for key in mapping):
-                    non_cuba_keys = [
-                        key for key in mapping if not isinstance(key, CUBA)]
-                    message = \
-                        "Key(s) {!r} are not in the approved CUBA keywords"
-                    raise ValueError(message.format(non_cuba_keys))
-            super(DataContainer, self).update(mapping)
+            if (isinstance(mapping, DataContainer) and
+                    mapping.restricted_keys == self.restricted_keys):
+                super(DataContainer, self).update(mapping)
+            else:
+                self._check_mapping(mapping)
+                super(DataContainer, self).update(mapping)
+
         super(DataContainer, self).update(
-            {CUBA[kwarg]: value for kwarg, value in kwargs.viewitems()})
+            {self._restricted_mapping[kwarg]: value
+             for kwarg, value in kwargs.viewitems()})
 
     def _check_arguments(self, args, kwargs):
         """ Check for the right arguments.
 
         """
         # See if there are any non CUBA keys in the keyword arguments
-        if any(key not in _CUBA_MEMBERS for key in kwargs):
-            non_cuba_keys = kwargs.viewkeys() - _CUBA_MEMBERS.viewkeys()
-            message = "Key(s) {!r} are not in the approved CUBA keywords"
-            raise ValueError(message.format(non_cuba_keys))
+        invalid_keys = [key for key in kwargs
+                        if key not in self._restricted_mapping]
+        if invalid_keys:
+            message = "Key(s) {!r} are not in the supported CUBA keywords"
+            raise ValueError(message.format(invalid_keys))
         # Only one positional argument is allowed.
         if len(args) > 1:
             message = 'DataContainer expected at most 1 arguments, got {}'
             raise TypeError(message.format(len(args)))
+
+    def _check_mapping(self, mapping):
+        ''' Check if the keys in the mappings are all supported CUBA keys
+
+        Parameters
+        ----------
+        mapping : Mapping
+
+        Raises
+        ------
+        ValueError
+            if any of the keys in the mappings is not supported
+        '''
+        invalid_keys = [key for key in mapping
+                        if (not isinstance(key, CUBA) or
+                            key not in self.restricted_keys)]
+        if invalid_keys:
+            message = 'Key(s) {!r} are not in the supported CUBA keywords'
+            raise ValueError(message.format(invalid_keys))
+
+
+def create_data_container(restricted_keys,
+                          class_name='RestrictedDataContainer'):
+    ''' Create a DataContainer subclass with the given
+    restricted keys
+
+    Note
+    ----
+    For pickling to work, the created class needs to be assigned
+    a name `RestrictedDataContainer` in the module where it is
+    created
+
+    Parameters
+    ----------
+    restricted_keys : sequence
+        CUBA IntEnum
+
+    class_name : str
+        Name of the returned class
+
+    Returns
+    -------
+    RestrictedDataContainer : DataContainer
+        subclass of DataContainer with the given `restricted_keys`
+
+    Examples
+    --------
+    >>> RestrictedDataContainer = create_data_container((CUBA.NAME,
+                                                         CUBA.VELOCITY))
+    >>> container = RestrictedDataContainer()
+    >>> container.restricted_keys
+    frozenset({<CUBA.VELOCITY: 55>, <CUBA.NAME: 61>})
+    >>> container[CUBA.NAME] = 'name'
+    >>> container[CUBA.POSITION] = (1.0, 1.0, 1.0)
+    ...
+    ValueError: Key <CUBA.POSITION: 119> is not in the supported CUBA keywords
+    '''
+    # Make sure all restricted keys are CUBA keys
+    if any(not isinstance(key, CUBA) for key in restricted_keys):
+        raise ValueError('All restricted keys should be CUBA IntEnum')
+
+    mapping = {key.name: key for key in restricted_keys}
+
+    new_class = type(class_name, (DataContainer,),
+                     {'__doc__': DataContainer.__doc__,
+                      '__slots__': (),
+                      'restricted_keys': frozenset(restricted_keys),
+                      '_restricted_mapping': mapping})
+
+    # For the dynamically generated class to have a chance to be
+    # pickleable.  Bypass this step for some Python implementations
+    try:
+        new_class.__module__ = _sys._getframe(1).f_globals.get(
+            '__name__', '__main__')
+    except AttributeError:
+        pass
+
+    return new_class

simphony/core/tests/test_data_container.py

@@ -1,7 +1,9 @@
+import cPickle
+from io import BytesIO
 import unittest
 
 from simphony.core.cuba import CUBA
-from simphony.core.data_container import DataContainer
+from simphony.core.data_container import DataContainer, create_data_container
 
 
 class TestDataContainer(unittest.TestCase):
@@ -63,6 +65,16 @@ def test_initialization_with_a_dictionary_of_ints(self):
         with self.assertRaises(ValueError):
             DataContainer(data)
 
+    def test_initialization_with_generator(self):
+        generator = ((key, key + 3) for key in CUBA)
+        container = DataContainer(generator)
+        self.assertEqual(len(container), len(CUBA))
+
+    def test_initialization_with_non_cuba_generator(self):
+        generator = (('foo'+str(i), i) for i in range(5))
+        with self.assertRaises(ValueError):
+            DataContainer(generator)
+
     def test_update_with_a_dictionary(self):
         container = DataContainer()
         data = {key: key + 3 for key in CUBA}
@@ -139,5 +151,66 @@ def test_setitem_with_non_cuba_key(self):
             container[100] = 29
 
 
+# Create a container class here for testing pickling
+# As with all dynamically created classes, it needs to be
+# properly named in a module in order for pickling to work
+RestrictedDataContainer = create_data_container(CUBA)
+
+
+class TestRestrictedDataContainer(unittest.TestCase):
+
+    def setUp(self):
+        self.maxDiff = None
+        iter_cuba = iter(CUBA)
+        # The first 9 keys are supported keys
+        self.valid_keys = tuple(iter_cuba.next() for i in range(1, 10))
+        # The rest are not supported
+        self.invalid_keys = tuple(key for key in iter_cuba)
+
+    def test_setitem_with_valid_key(self):
+        container = create_data_container(self.valid_keys)()
+        container[self.valid_keys[0]] = 20
+        self.assertIsInstance(container.keys()[0], CUBA)
+        self.assertEqual(container[self.valid_keys[0]], 20)
+
+    def test_setitem_with_invalid_key(self):
+        container = create_data_container(self.valid_keys)()
+
+        for key in self.invalid_keys:
+            with self.assertRaises(ValueError):
+                container[key] = 1
+
+    def test_update_with_valid_keys(self):
+        data = {key: key+3 for key in self.valid_keys}
+        container = create_data_container(self.valid_keys)(data)
+        self.assertTrue(all(key in self.valid_keys for key in container))
+
+    def test_update_with_some_invalid_keys(self):
+        data = {key: key+3 for key in self.valid_keys}
+        data[self.invalid_keys[0]] = 20
+
+        with self.assertRaises(ValueError):
+            create_data_container(self.valid_keys)(data)
+
+    def test_error_with_non_cuba_keys(self):
+        with self.assertRaises(ValueError):
+            create_data_container((1, 2))
+
+    def test_data_container_can_be_pickled(self):
+        # Create a container with data
+        container = RestrictedDataContainer()
+        for key in CUBA:
+            container[key] = key+3
+
+        # Pickle and write to a buffer
+        stream = BytesIO()
+        cPickle.dump(container, stream)
+        stream.seek(0)
+
+        # Restore the data
+        pickled_data = cPickle.load(stream)
+        self.assertDictEqual(pickled_data, container)
+
+
 if __name__ == '__main__':
     unittest.main()