Merge pull request #964 from mitmproxy/flat-structure

Flat structure

1 files changed, 67 insertions, 0 deletions

diff --git a/pathod/language/writer.py b/pathod/language/writer.py
new file mode 100644
index 00000000..1a27e1ef
--- /dev/null
+++ b/pathod/language/writer.py
@@ -0,0 +1,67 @@
+import time
+from netlib.exceptions import TcpDisconnect
+import netlib.tcp
+
+BLOCKSIZE = 1024
+# It's not clear what the upper limit for time.sleep is. It's lower than the
+# maximum int or float. 1 year should do.
+FOREVER = 60 * 60 * 24 * 365
+
+
+def send_chunk(fp, val, blocksize, start, end):
+    """
+        (start, end): Inclusive lower bound, exclusive upper bound.
+    """
+    for i in range(start, end, blocksize):
+        fp.write(
+            val[i:min(i + blocksize, end)]
+        )
+    return end - start
+
+
+def write_values(fp, vals, actions, sofar=0, blocksize=BLOCKSIZE):
+    """
+        vals: A list of values, which may be strings or Value objects.
+
+        actions: A list of (offset, action, arg) tuples. Action may be "pause"
+        or "disconnect".
+
+        Both vals and actions are in reverse order, with the first items last.
+
+        Return True if connection should disconnect.
+    """
+    sofar = 0
+    try:
+        while vals:
+            v = vals.pop()
+            offset = 0
+            while actions and actions[-1][0] < (sofar + len(v)):
+                a = actions.pop()
+                offset += send_chunk(
+                    fp,
+                    v,
+                    blocksize,
+                    offset,
+                    a[0] - sofar - offset
+                )
+                if a[1] == "pause":
+                    time.sleep(
+                        FOREVER if a[2] == "f" else a[2]
+                    )
+                elif a[1] == "disconnect":
+                    return True
+                elif a[1] == "inject":
+                    send_chunk(fp, a[2], blocksize, 0, len(a[2]))
+            send_chunk(fp, v, blocksize, offset, len(v))
+            sofar += len(v)
+        # Remainders
+        while actions:
+            a = actions.pop()
+            if a[1] == "pause":
+                time.sleep(a[2])
+            elif a[1] == "disconnect":
+                return True
+            elif a[1] == "inject":
+                send_chunk(fp, a[2], blocksize, 0, len(a[2]))
+    except TcpDisconnect:  # pragma: no cover
+        return True