Merge pull request #5 from enkore/master

Refactored IO

i3pystatus/__init__.py

@@ -5,6 +5,7 @@ import json
 import urllib.request, urllib.error, urllib.parse
 from threading import Thread
 import time
+from contextlib import contextmanager
 
 class BaseModule:
     output = None
@@ -38,30 +39,23 @@ class IntervalModule(AsyncModule):
             self.run()
             time.sleep(self.interval)
 
-class I3statusHandler:
+class IOHandler:
-    modules = []
+    def __init__(self, inp=sys.stdin, out=sys.stdout):
+        self.inp = inp
+        self.out = out
 
-    def __init__(self):
+    def write(self, message):
-        pass
-
-    def register(self, module):
-        """Register a new module."""
-
-        self.modules.append(module)
-        module.registered(self)
-
-    def print_line(self, message):
         """Unbuffered printing to stdout."""
 
-        sys.stdout.write(message + "\n")
+        self.out.write(message + "\n")
-        sys.stdout.flush()
+        self.out.flush()
 
-    def read_line(self):
+    def read(self):
         """Interrupted respecting reader for stdin."""
 
         # try reading a line, removing any extra whitespace
         try:
-            line = sys.stdin.readline().strip()
+            line = self.inp.readline().decode("utf-8").strip()
             # i3status sends EOF, or an empty line
             if not line:
                 sys.exit(3)
@@ -70,26 +64,48 @@ class I3statusHandler:
         except KeyboardInterrupt:
             sys.exit()
 
+class JSONIO:
+    def __init__(self, io):
+        self.io = io
+        self.io.write(self.io.read())
+        self.io.write(self.io.read())
+
+    @contextmanager
+    def read(self):
+        line, prefix = self.io.read(), ""
+
+        # ignore comma at start of lines
+        if line.startswith(","):
+            line, prefix = line[1:], ","
+
+        j = json.loads(line)
+        yield j
+
+        self.io.write(prefix + json.dumps(j))
+
+class I3statusHandler:
+    modules = []
+
+    def __init__(self, fd=None):
+        if fd is None:
+            fd = sys.stdin
+
+        self.io = IOHandler(fd)
+
+    def register(self, module):
+        """Register a new module."""
+
+        self.modules.append(module)
+        module.registered(self)
+
     def run(self):
-        self.print_line(self.read_line())
+        jio = JSONIO(self.io)
-        self.print_line(self.read_line())
 
         while True:
-            line, prefix = self.read_line(), ""
+            with jio.read() as j:
+                for module in self.modules:
+                    module.tick()
 
-            # ignore comma at start of lines
+                    output = module.output
-            if line.startswith(","):
+                    if output:
-                line, prefix = line[1:], ","
+                        j.insert(0, output)
-
-            j = json.loads(line)
-
-            for module in self.modules:
-                module.tick()
-
-                output = module.output
-
-                if output:
-                    j.insert(0, output)
-
-            # and echo back new encoded json
-            self.print_line(prefix+json.dumps(j))