Implemented the wait in the zauto (before the close) and fixed the bl…

…ocking next.

py/zero/__init__.py

@@ -404,9 +404,11 @@ def __iter__(self):
         return self
 
     def next(self):
-        ''' Receives a message. Blocks regardless of self.setup.block. If method is rep, must send
-            reply before going to next(). The message is unmarshalled and returned.
+        ''' Receives a message. If method is rep, must send reply before going to next(). The
+            message is unmarshalled and returned.
         '''
+        if not self.setup.block and not self.sock.poll(timeout=100): # Milliseconds; 0.1s
+            raise StopIteration()
         res = self._decode(self.sock.recv())
         self.setup.debug('Received %r from %s', res, self.setup.point)
         if self.active:
@@ -447,7 +449,7 @@ def opposite(self):
         return Zero(self.setup.opposite())
 
 
-def zauto(zero, loops):
+def zauto(zero, loops, wait=False):
     'Keep listening and sending until the loop ends. All received objects are yielded.'
     try:
         if zero.setup.replies:
@@ -467,6 +469,8 @@ def zauto(zero, loops):
     except StopIteration:
         zero.setup.debug('Loop ended')
     finally:
+        if wait:
+            raw_input('Press enter when done.')
         zero.setup.debug('Closing: %r', zero)
         zero.close()
 

py/zero/__main__.py

@@ -22,7 +22,7 @@ def main():
         setup, loop = ZeroSetup.argv()
         zero = Zero(setup)
 
-        for msg in zauto(zero, loop):
+        for msg in zauto(zero, loop, setup.args['--wait']):
             sys.stdout.write(json.dumps(msg) + '\n')
             sys.stdout.flush()
     except UnsupportedZmqMethod, e:
@@ -45,9 +45,9 @@ def main():
         else:
             # Something happened...
             raise e
-    if setup.args['--wait']:
-        raw_input('Press enter when done.')
-    zero.close()
+        if setup.args['--wait']:
+            raw_input('Press enter when done.')
+        zero.close()
 
 if __name__ == '__main__':
     main()