Merge pull request #16 from alexras/alexras/multiple-key-enums

Allow multiple enum keys to map to the same value

bread/__init__.py

@@ -11,7 +11,7 @@
 from .lifecycle import *
 
 __title__ = 'bread'
-__version__ = '3.0.2'
+__version__ = '3.1.0'
 __author__ = 'Alex Rasmussen'
 __license__ = 'MIT'
 __copyright__ = 'Copyright 2015 Alex Rasmussen'

bread/enum.py

@@ -1,32 +1,68 @@
 from .integer import intX
 
 
+def _validate_key_types(enum_values):
+    for key in enum_values.keys():
+        if type(key) == int:
+            continue
+
+        valid = False
+
+        if type(key) in (tuple, list):
+            for alternative in key:
+                valid = (type(alternative) == int)
+
+        if not valid:
+            raise ValueError("Keys in an enum's values dict should "
+                             "be ints, int tuples, or int lists (%s)" % (enum_values))
+
+
 def enum(length, values, default=None):
     def make_enum_field(parent, **field_options):
         enum_field = intX(length, signed=False)(parent, **field_options)
 
         old_encode_fn = enum_field._encode_fn
         old_decode_fn = enum_field._decode_fn
 
-        keys = {v: k for k, v in list(values.items())}
+        keys = {}
+        flattened_values = {}
+
+        _validate_key_types(values)
+
+        for k, v in values.items():
+            if type(k) in (tuple, list):
+                for alternative in k:
+                    if type(alternative) is not int:
+                        raise ValueError("Keys in an enum's values dict "
+                                         "should be ints, tuples, or lists (%s)" % (values))
+
+                    flattened_values[alternative] = v
+
+                # If presented with a list of options for an enum's integer representation, always pick the first one
+                keys[v] = k[0]
+            elif type(k) == int:
+                keys[v] = k
+                flattened_values[k] = v
+            else:
+                raise ValueError("Keys in an enum's values dict should be ints, tuples, or lists (%s)" % (values))
 
         def encode_enum(key):
             if key not in keys:
-                raise ValueError('%s is not a valid enum value; valid values %s' % (key, keys))
+                raise ValueError('%s is not a valid enum value; valid values %s' % (key, keys.keys()))
 
             return old_encode_fn(keys[key])
 
         def decode_enum(encoded):
             decoded_value = old_decode_fn(encoded)
 
-            if decoded_value not in values:
+            if decoded_value not in flattened_values:
                 if default is not None:
                     return default
                 else:
                     raise ValueError(
-                        '%d is not a valid enum value; valid values %s' % (decoded_value, values))
+                        '%d is not a valid enum value; valid values %s' % (decoded_value, flattened_values))
 
-            return values[decoded_value]
+            return flattened_values[decoded_value]
 
         enum_field._encode_fn = encode_enum
         enum_field._decode_fn = decode_enum

docs/source/conf.py

@@ -48,9 +48,9 @@
 # built documents.
 #
 # The short X.Y version.
-version = '3.0.2'
+version = '3.1.0'
 # The full version, including alpha/beta/rc tags.
-release = '3.0.2'
+release = '3.1.0'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.

docs/source/spec_language.rst

@@ -73,6 +73,22 @@ Here is an example of a 2-bit field representing a card suit: ::
          3: "clubs"
      }))
 
+If multiple bit sequences can represent the same value, you can use a ``tuple``
+or ``list`` of keys in the enum's specification, like so: ::
+
+    import bread as b
+    ('waveform', b.enum(8, {
+        0: 'triangle',
+        1: 'saw down',
+        2: 'saw up',
+        3: 'square',
+        4: 'sine',
+        (5, 7): 'sample and hold'
+    }))
+
+When constructing a new struct, if multiple bit sequences denote the same
+value, the first value in the tuple or list is used.
+
 Arrays
 ~~~~~~
 

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup
 
 setup(name='bread',
-      version='3.0.2',
+      version='3.1.0',
       description='Binary format parsing made easier',
       url='https://github.com/alexras/bread',
       author='Alex Rasmussen',

test.py

@@ -689,6 +689,40 @@ def test_enum_default():
     assert result.suit == "spades"
 
 
+def test_enum_multiple_values_same_key():
+    enum_test = [
+        ('waveform', b.enum(8, {
+            0: 'triangle',
+            1: 'saw down',
+            2: 'saw up',
+            3: 'square',
+            4: 'sine',
+            (5, 7): 'sample and hold'
+        }))
+    ]
+
+    data = bytearray([7])
+    result = b.parse(data, enum_test)
+    assert result.waveform == 'sample and hold'
+
+    dumped = b.write(result, enum_test)
+    assert dumped == bytearray([7])
+
+    data = bytearray([5])
+    result = b.parse(data, enum_test)
+    assert result.waveform == 'sample and hold'
+
+    data = bytearray([2])
+    result = b.parse(data, enum_test)
+    assert result.waveform == 'saw up'
+
+    test_struct = b.new(enum_test)
+    test_struct.waveform = 'sample and hold'
+
+    dumped = b.write(test_struct, enum_test)
+    assert dumped == bytearray([5])
+
+
 def test_enum_set_invalid_value():
     with pytest.raises(ValueError):
         enum_test = [