Ticket #51: test_socket.py

import yaml
import socket
from optparse import OptionParser
import time

class SockStream:
    def __init__(self, sock):
        self.sock = sock
    def read(self, bytes):
        return self.sock.recv(bytes)
    def write(self, string):
        return self.sock.send(string)
    def __getattr__(self, name):
        return getattr(self.sock, name)


op = OptionParser()
op.add_option("--server", action="store_true", default=False)
op.add_option("--port", type="int", default=8000)
opts, args = op.parse_args()

if opts.server:
    serv = socket.socket()
    serv.bind(('127.0.0.1', opts.port))
    serv.listen(1)
    s,addr = serv.accept()
    s = SockStream(s)
    for item in yaml.load_all(s):
        print item
else:
    s = socket.socket()
    s.connect(('127.0.0.1',opts.port))
    s = SockStream(s)
    def generator():
        for i in range(100):
            time.sleep(1)
            print i
            yield i
    yaml.dump_all(generator(), s, explicit_start=True, explicit_end=True)