Add Bitlist and Bitvector dump/parse

ethereum · Jul 5, 2019 · 25037b9 · 25037b9
1 parent bdbb43c
commit 25037b9
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/ssz/examples/zoo.py b/ssz/examples/zoo.py
@@ -1,4 +1,6 @@
 from ssz.sedes import (
+    Bitlist,
+    Bitvector,
     List,
     Serializable,
     SignedSerializable,
@@ -15,6 +17,8 @@ class ClockInRecords(SignedSerializable):
     fields = [
         ('epoch', uint64),
         ('bio_id_scan', bytes32),
+        ('poo_log_bits', Bitlist(32)),
+        ('wash_log_bits', Bitvector(32)),
         ('signature', bytes96),
     ]
 
@@ -42,11 +46,15 @@ class Zoo(Serializable):
             epoch=123,
             bio_id_scan=b'\xab' * 32,
             signature=b'\xab' * 96,
+            poo_log_bits=(True,) * 16 + (False,) * 16,
+            wash_log_bits=(False,) * 16 + (True,) * 16,
         ),
         ClockInRecords(
             epoch=124,
             bio_id_scan=b'\xab' * 32,
             signature=b'\xab' * 96,
+            poo_log_bits=(True,) * 16 + (False,) * 16,
+            wash_log_bits=(False,) * 16 + (True,) * 16,
         )
     ),
     vaccinated=True,
@@ -59,6 +67,8 @@ class Zoo(Serializable):
             epoch=125,
             bio_id_scan=b'\xcd' * 32,
             signature=b'\xcd' * 96,
+            poo_log_bits=(True,) * 16 + (False,) * 16,
+            wash_log_bits=(False,) * 16 + (True,) * 16,
         ),
     ),
     vaccinated=True,

diff --git a/ssz/tools/dump.py b/ssz/tools/dump.py
@@ -8,6 +8,8 @@
 
 from ssz.sedes import (
     BaseSedes,
+    Bitlist,
+    Bitvector,
     Boolean,
     ByteVector,
     Container,
@@ -46,6 +48,8 @@ def dump(value, sedes=None, codec=DefaultCodec):
         return dump_list(value, sedes, codec)
     elif isinstance(sedes, Vector):
         return dump_vector(value, sedes, codec)
+    elif isinstance(sedes, (Bitlist, Bitvector)):
+        return dump_bits(value, sedes, codec)
     elif isinstance(sedes, Container):
         return dump_container(value, sedes, codec)
     elif isinstance(sedes, MetaSerializable):
@@ -84,6 +88,10 @@ def dump_vector(value, sedes, codec):
     return tuple(dump(element, sedes.element_sedes, codec) for element in value)
 
 
+def dump_bits(value, sedes, codec):
+    return codec.encode_bytes(sedes.serialize(value), sedes)
+
+
 @to_tuple
 def dump_container(value, sedes, codec):
     if not isinstance(value, Sequence):

diff --git a/ssz/tools/parse.py b/ssz/tools/parse.py
@@ -11,6 +11,8 @@
 
 from ssz.sedes import (
     BaseSedes,
+    Bitlist,
+    Bitvector,
     Boolean,
     ByteVector,
     Container,
@@ -44,6 +46,8 @@ def parse(value, sedes, codec=DefaultCodec):
         return parse_vector(value, sedes, codec)
     elif isinstance(sedes, Container):
         return parse_container(value, sedes, codec)
+    elif isinstance(sedes, (Bitlist, Bitvector)):
+        return parse_bits(value, sedes, codec)
     elif isinstance(sedes, MetaSerializable):
         return parse_serializable(value, sedes, codec)
     elif isinstance(sedes, BaseSedes):
@@ -80,6 +84,15 @@ def parse_vector(value, sedes, codec):
     return tuple(parse(element, sedes.element_sedes, codec) for element in value)
 
 
+def parse_bits(value, sedes, codec):
+    if not isinstance(value, Sequence):
+        raise ValueError(f"Expected Sequence, got {type(value)}")
+
+    data = parse_bytes(value, sedes, codec)
+    deserialized = sedes.deserialize(data)
+    return tuple(parse_boolean(element, sedes, codec) for element in deserialized)
+
+
 @to_tuple
 def parse_container(value, sedes, codec):
     if not isinstance(value, Sequence):