From d0b74c55a07ce9e623578bf35d37a597e1a3c179 Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bsima.me>
Date: Wed, 8 Apr 2020 20:43:29 -0700
Subject: Refactor client, fixing some bugs

---
 Run/Que/client.py | 126 +++++++++++++++++++++++++++---------------------------
 1 file changed, 62 insertions(+), 64 deletions(-)

(limited to 'Run')

diff --git a/Run/Que/client.py b/Run/Que/client.py
index 6a91d0c..43517ff 100755
--- a/Run/Que/client.py
+++ b/Run/Que/client.py
@@ -4,34 +4,80 @@ simple client for que.run
 """
 
 import argparse
-import urllib.request as request
-import urllib.parse
-import time
 import subprocess
 import sys
+import time
+import urllib.parse
+import urllib.request as request
+
+MAX_TIMEOUT = 99999999  # basically never timeout
 
-# set to something ridiculously high so we don't run into timeouts while polling
-# or waiting for a message
-MAX_TIMEOUT = 100000000
 
+def send(args):
+    "Send a message to the que."
+    data = args.infile.read().encode("utf-8").strip()
+    req = request.Request(f"{args.host}/{args.target}")
+    req.add_header("User-AgenT", "Que/Client")
+    if args.serve:
+        while not time.sleep(1):
+            with request.urlopen(req, data=data, timeout=MAX_TIMEOUT) as req:
+                pass
+    else:
+        with request.urlopen(req, data=data, timeout=MAX_TIMEOUT) as req:
+            pass
+
+
+def recv(args):
+    "Receive a message from the que."
 
-def main(argv=None):
+    def _recv(_req):
+        msg = _req.readline().decode("utf-8").strip()
+        print(msg)
+        if args.then:
+            subprocess.run(
+                args.then.replace("\msg", msg).replace("\que", args.target), shell=True
+            )
+
+    params = urllib.parse.urlencode({"poll": args.poll})
+    req = request.Request(f"{args.host}/{args.target}?{params}")
+    req.add_header("User-Agent", "Que/Client")
+
+    with request.urlopen(req) as _req:
+        if args.poll:
+            while not time.sleep(1):
+                _recv(_req)
+        else:
+            _recv(_req)
+
+
+def get_args():
     cli = argparse.ArgumentParser(description=__doc__)
     cli.add_argument(
-        "--host", default="https://que.run", help="where que-server is running"
+        "--host", default="http://que.run", help="where que-server is running"
     )
     cli.add_argument(
         "--poll", default=False, action="store_true", help="stream data from the que"
     )
     cli.add_argument(
         "--then",
-        help="when polling, run this shell command after each response, replacing '\que' with the target and '\msg' with the body of the response",
+        help=" ".join(
+            [
+                "when polling, run this shell command after each response,",
+                "presumably for side effects,"
+                "replacing '\que' with the target and '\msg' with the body of the response",
+            ]
+        ),
     )
     cli.add_argument(
         "--serve",
         default=False,
         action="store_true",
-        help="when posting to the que, do so continuously in a loop. this can be used for serving a webpage or other file continuously",
+        help=" ".join(
+            [
+                "when posting to the que, do so continuously in a loop.",
+                "this can be used for serving a webpage or other file continuously",
+            ]
+        ),
     )
     cli.add_argument(
         "target", help="namespace and path of the que, like 'ns/path/subpath'"
@@ -42,66 +88,18 @@ def main(argv=None):
         type=argparse.FileType("r"),
         help="data to put on the que. Use '-' for stdin, otherwise should be a readable file",
     )
-
-    if argv is None:
-        args = cli.parse_args()
-    else:
-        args = cli.parse_args(argv)
-
-    if args.infile:
-        # send input data
-        data = args.infile.read().encode("utf-8").strip()
-        if args.serve:
-            # loop until ^C
-            while not time.sleep(1):
-                with request.urlopen(
-                    f"{args.host}/{args.target}",
-                    data=data,
-                    timeout=MAX_TIMEOUT,
-                    headers={"User-Agent": "Que/Client"},
-                ) as req:
-                    pass
-        else:
-            with request.urlopen(
-                f"{args.host}/{args.target}",
-                data=data,
-                timeout=MAX_TIMEOUT,
-                headers={"User-Agent": "Que/Client"},
-            ) as req:
-                pass
-    else:
-        # no input data, do a read instead
-        params = urllib.parse.urlencode({"poll": args.poll})
-        url = f"{args.host}/{args.target}?{params}"
-        with request.urlopen(url, headers={"User-Agent": "Que/Client"}) as req:
-            if args.poll:
-                while not time.sleep(1):
-                    msg = req.readline().decode("utf-8").strip()
-                    print(msg)
-                    if args.then:
-                        subprocess.run(
-                            args.then.replace("\msg", msg).replace("\que", args.target),
-                            shell=True,
-                        )
-            else:
-                msg = req.read().decode("utf-8").strip()
-                print(msg)
-                if args.then:
-                    subprocess.run(
-                        args.then.replace("\msg", msg).replace("\que", args.target),
-                        shell=True,
-                    )
+    return cli.parse_args()
 
 
 if __name__ == "__main__":
     try:
-        main()
+        args = get_args()
+        if args.infile:
+            send(args)
+        else:
+            recv(args)
     except KeyboardInterrupt:
-        print("Interrupted")
         sys.exit(0)
     except urllib.error.HTTPError as e:
         print(e)
         sys.exit(1)
-    except http.client.IncompleteRead as e:
-        print(e)
-        sys.exit(1)
-- 
cgit v1.2.3