turn bits into a class

1 files changed, 16 insertions, 2 deletions

diff --git a/subv.py b/subv.py
index 4f38988..dd45b16 100644
--- a/subv.py
+++ b/subv.py
@@ -1,4 +1,5 @@
 import re
+import bits
 white = re.compile(r'[ \t\.\n]+')
 hex = re.compile(r'^\-?(0x)?[0-9a-f]+$')
 num = re.compile(r'^\d+$')
@@ -230,6 +231,11 @@ def untag(part, expect=None):
     >>> untag(('$label', 'imm20'), expect='imm20')
     '$label'
 
+    >>> untag((2, 'num'), expect=['num', 'imm20'])
+    2
+    >>> untag(('$label', 'imm20'), expect=['num', 'imm20'])
+    '$label'
+
     >>> untag((2, 'imm12'), expect='imm20')
     Traceback (most recent call last):
         ...
@@ -238,9 +244,15 @@ def untag(part, expect=None):
     Traceback (most recent call last):
         ...
     ValueError: expected ('$label', 'imm20') to be labelled off12
+    >>> untag((2, 'imm12'), expect=['num', 'off12'])
+    Traceback (most recent call last):
+        ...
+    ValueError: expected (2, 'imm12') to be labelled one of ['num', 'off12']
     """
-    if expect and part[1] != expect:
+    if isinstance(expect, str) and part[1] != expect:
         raise ValueError("expected {} to be labelled {}".format(part, expect))
+    elif isinstance(expect, list) and part[1] not in expect:
+        raise ValueError("expected {} to be labelled one of {}".format(part, expect))
     return part[0]
 
 def format_reference(ref):
@@ -287,7 +299,9 @@ def format_part(part):
     >>> format_part(('$label:suff', 'tag'))
     '$label:suff/tag'
     """
-    if not is_reference(part):
+    if isinstance(part, bits.Bitfield):
+        part = ('{:02x}'.format(part.val), part.size)
+    elif not is_reference(part):
         first = '{:02x}'.format(part[0])
         part = (first, *part[1:])
     return '/'.join([str(p) for p in part])