Refactor barrier to make importable

.gitignore

@@ -0,0 +1 @@
+__pycache__/

barrier.py

@@ -3,27 +3,49 @@
 import argparse
 import socket
 
-def barrier(host, port, wait_for):
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        s.bind((host, port))
-        s.listen()
+class Barrier:
+    def __init__(self, host, port, wait_for):
+        self.host = host
+        self.port = port
+        self.wait_for = wait_for
 
+    def listen(self):
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self.sock.bind((self.host, self.port))
+        self.sock.listen()
+
+    def waitSingle(self):
         connections = []
         addresses = []
 
         while True:
-            conn, addr = s.accept()
+            conn, addr = self.sock.accept()
             connections.append(conn)
             addresses.append(addr)
-            print("Connection from {}".format(addr))
+            yield addr
 
-            if len(connections) == wait_for:
+            if len(connections) == self.wait_for:
                 break
 
         for conn in connections:
             conn.close()
 
+        return None
+
+    def wait(self):
+        g = self.waitSingle()
+
+        conn = []
+        while True:
+            try:
+                conn.append(next(g))
+            except StopIteration:
+                break
+
+        return conn
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
@@ -52,5 +74,16 @@ if __name__ == "__main__":
 
     results = parser.parse_args()
 
+    barrier = Barrier(results.host, results.port, results.processes)
+    barrier.listen()
     print("Barrier listens on {}:{} and waits for {} processes".format(results.host, results.port, results.processes))
-    barrier(results.host, results.port, results.processes)
+
+    # connectedAddr = barrier.wait()
+    connectedAddrGen = barrier.waitSingle()
+
+    while True:
+        try:
+            connectedAddr = next(connectedAddrGen)
+            print("Connection from {}".format(connectedAddr))
+        except StopIteration:
+            break