diff --git a/simphony/core/data_container.py b/simphony/core/data_container.py index 69302fc7..6a9ebce4 100644 --- a/simphony/core/data_container.py +++ b/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({, }) + >>> container[CUBA.NAME] = 'name' + >>> container[CUBA.POSITION] = (1.0, 1.0, 1.0) + ... + ValueError: Key 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 diff --git a/simphony/core/tests/test_data_container.py b/simphony/core/tests/test_data_container.py index 84295ea6..6e157c17 100644 --- a/simphony/core/tests/test_data_container.py +++ b/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()