From a9bd4fe47679dd516ef3ab91106b48d7bf58c60b Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Fri, 5 Jul 2019 16:01:00 +0800
Subject: [PATCH] Add Bitlist and Bitvector dump/parse

---
 ssz/examples/zoo.py | 10 ++++++++++
 ssz/tools/dump.py   |  8 ++++++++
 ssz/tools/parse.py  | 13 +++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/ssz/examples/zoo.py b/ssz/examples/zoo.py
index 00e10c08..9f30f083 100644
--- 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
index 81987ffc..bdb12d1a 100644
--- 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
index 03a2aa24..ba86f3ab 100644
--- 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):