websockets: more compact and legible human_readable

2 files changed, 38 insertions, 25 deletions

diff --git a/netlib/utils.py b/netlib/utils.py
index 905d948f..7e539977 100644
--- a/netlib/utils.py
+++ b/netlib/utils.py
@@ -70,11 +70,12 @@ def getbit(byte, offset):
 class BiDi:
     """
         A wee utility class for keeping bi-directional mappings, like field
-        constants in protocols:
+        constants in protocols. Names are attributes on the object, dict-like
+        access maps values to names:
 
         CONST = BiDi(a=1, b=2)
         assert CONST.a == 1
-        assert CONST[1] == "a"
+        assert CONST.get_name(1) == "a"
     """
     def __init__(self, **kwargs):
         self.names = kwargs
@@ -89,5 +90,21 @@ class BiDi:
             return self.names[k]
         raise AttributeError("No such attribute: %s", k)
 
-    def __getitem__(self, k):
-        return self.values[k]
+    def get_name(self, n, default=None):
+        return self.values.get(n, default)
+
+
+def pretty_size(size):
+    suffixes = [
+        ("B", 2**10),
+        ("kB", 2**20),
+        ("MB", 2**30),
+    ]
+    for suf, lim in suffixes:
+        if size >= lim:
+            continue
+        else:
+            x = round(size/float(lim/2**10), 2)
+            if x == int(x):
+                x = int(x)
+            return str(x) + suf
diff --git a/netlib/websockets.py b/netlib/websockets.py
index d358ed53..1d02d684 100644
--- a/netlib/websockets.py
+++ b/netlib/websockets.py
@@ -162,16 +162,21 @@ class FrameHeader:
             raise ValueError("Masking key must be 4 bytes.")
 
     def human_readable(self):
-        return "\n".join([
-            ("fin                   - " + str(self.fin)),
-            ("rsv1                  - " + str(self.rsv1)),
-            ("rsv2                  - " + str(self.rsv2)),
-            ("rsv3                  - " + str(self.rsv3)),
-            ("opcode                - " + str(self.opcode)),
-            ("mask                  - " + str(self.mask)),
-            ("length_code           - " + str(self.length_code)),
-            ("masking_key           - " + repr(str(self.masking_key))),
-        ])
+        vals = [
+            "wf:",
+            OPCODE.get_name(self.opcode, hex(self.opcode)).lower()
+        ]
+        flags = []
+        for i in ["fin", "rsv1", "rsv2", "rsv3", "mask"]:
+            if getattr(self, i):
+                flags.append(i)
+        if flags:
+            vals.extend([":", "|".join(flags)])
+        if self.masking_key:
+            vals.append(":key=%s"%repr(self.masking_key))
+        if self.payload_length:
+            vals.append(" %s"%utils.pretty_size(self.payload_length))
+        return "".join(vals)
 
     def to_bytes(self):
         first_byte = utils.setbit(0, 7, self.fin)
@@ -308,17 +313,8 @@ class Frame(object):
         return cls.from_file(io.BytesIO(bytestring))
 
     def human_readable(self):
-        return "\n".join([
-            ("fin                   - " + str(self.header.fin)),
-            ("rsv1                  - " + str(self.header.rsv1)),
-            ("rsv2                  - " + str(self.header.rsv2)),
-            ("rsv3                  - " + str(self.header.rsv3)),
-            ("opcode                - " + str(self.header.opcode)),
-            ("mask                  - " + str(self.header.mask)),
-            ("length_code           - " + str(self.header.length_code)),
-            ("masking_key           - " + repr(str(self.header.masking_key))),
-            ("payload               - " + repr(str(self.payload))),
-        ])
+        hdr = self.header.human_readable()
+        return hdr + "\n" + repr(self.payload)
 
     def to_bytes(self):
         """