From 4880e4225db12af21b1c4ea7ec4797f090ba2cd0 Mon Sep 17 00:00:00 2001
From: Robert Jordens <rj@m-labs.hk>
Date: Thu, 9 Nov 2017 12:52:36 +0100
Subject: [PATCH] bit2bin: cleanup

---
 artiq/frontend/bit2bin.py | 66 ++++++++++++++++++++-------------------
 1 file changed, 34 insertions(+), 32 deletions(-)

diff --git a/artiq/frontend/bit2bin.py b/artiq/frontend/bit2bin.py
index 82a33f546..f511c0559 100755
--- a/artiq/frontend/bit2bin.py
+++ b/artiq/frontend/bit2bin.py
@@ -17,38 +17,39 @@ def flip32(data):
 
 
 def bit2bin(bit, bin, flip=False):
-    with open(bit, "rb") as bitfile:
-        l, = struct.unpack(">H", bitfile.read(2))
-        if l != 9:
-            raise ValueError("Missing <0009> header, not a bit file")
+    l, = struct.unpack(">H", bit.read(2))
+    if l != 9:
+        raise ValueError("Missing <0009> header, not a bit file")
+    _ = bit.read(l)  # unknown data
+    l, = struct.unpack(">H", bit.read(2))
+    if l != 1:
+        raise ValueError("Missing <0001> header, not a bit file")
 
-        bitfile.read(l)
-        d = bitfile.read(*struct.unpack(">H", bitfile.read(2)))
-        if d != b"a":
-            raise ValueError("Missing <a> header, not a bit file")
-
-        d = bitfile.read(*struct.unpack(">H", bitfile.read(2)))
-        print("Design name:", d)
-
-        while True:
-            key = bitfile.read(1)
-            if not key:
-                break
-            if key in b"bcd":
-                d = bitfile.read(*struct.unpack(">H", bitfile.read(2)))
-                name = {b"b": "Partname", b"c": "Date", b"d": "Time"}[key]
-                print(name, d)
-            elif key == b"e":
-                l, = struct.unpack(">I", bitfile.read(4))
-                print("found binary data length:", l)
-                d = bitfile.read(l)
-                if flip:
-                    d = flip32(d)
-                with open(bin, "wb") as f:
-                    f.write(d)
-            else:
-                d = bitfile.read(*struct.unpack(">H", bitfile.read(2)))
-                print("Unexpected key: ", key, d)
+    while True:
+        key = bit.read(1).decode()
+        if not key:
+            break
+        if key in "abcd":
+            d = bit.read(*struct.unpack(">H", bit.read(2)))
+            assert d.endswith(b"\x00")
+            d = d.decode()
+            name = {
+                    "a": "Design",
+                    "b": "Part name",
+                    "c": "Date",
+                    "d": "Time"
+                    }[key]
+            print("{}: {}".format(name, d))
+        elif key == "e":
+            l, = struct.unpack(">I", bit.read(4))
+            print("Bitstream payload length: {:#x}".format(l))
+            d = bit.read(l)
+            if flip:
+                d = flip32(d)
+            bin.write(d)
+        else:
+            d = bit.read(*struct.unpack(">H", bit.read(2)))
+            print("Unexpected key: {}: {}".format(key, d))
 
 
 if __name__ == "__main__":
@@ -64,4 +65,5 @@ if __name__ == "__main__":
                         help="Output bin file name")
     args = parser.parse_args()
 
-    bit2bin(args.bitfile, args.binfile, args.flip)
+    with open(args.bitfile, "rb") as f, open(args.binfile, "wb") as g:
+        bit2bin(f, g, args.flip)