fizzbuzz

fizzbuzz.py

import socket
from collections import namedtuple
from enum import Enum
import struct
import argparse
import asyncio
import sys


class StructMeta(type):
    def __new__(*args, **kwargs):
        result = type.__new__(*args, **kwargs)
        annotations = getattr(result, '__annotations__', {})
        codes = ''.join(t.value for t in annotations.values())
        result.format = f'{result.byte_order.value}{codes}'
        result.size = struct.calcsize(result.format)
        result.namedtuple = namedtuple(result.__name__, annotations.keys())
        return result


class Struct(metaclass=StructMeta):
    class ByteOrder(Enum):
        Native = '@'
        LittleEndian = '<'
        BigEndian = '>'
        Network = '!'
    byte_order = ByteOrder.Native

    class Types(Enum):
        Pad = 'x'
        Char = 'c'
        SignedChar = 'b'
        UnsignedChar = 'B'
        Bool = '?'
        Short = 'h'
        UnsignedShort = 'H'
        Int = 'i'
        UnsignedInt = 'I'
        Long = 'l'
        UnsignedLong = 'L'
        LongLong = 'q'
        UnsignedLongLong = 'Q'
        SSizeT = 'n'
        SizeT = 'N'
        Float = 'f'
        Double = 'd'
        Void = 'P'

        @staticmethod
        def String(size):
            return namedtuple('String', ['name', 'value'])('String', f'{size}s')

    @classmethod
    def new(cls, *args, **kwargs):
        return cls.namedtuple(*args, **kwargs)

    @classmethod
    def decode(cls, data):
        return cls.namedtuple(*struct.unpack(cls.format, data))

    @classmethod
    def encode(cls, data):
        return struct.pack(cls.format, *data)


class GeneratorRequest(Struct):
    byte_order = Struct.ByteOrder.Network
    amount: Struct.Types.UnsignedInt


class GeneratorResponse(Struct):
    byte_order = Struct.ByteOrder.Network
    value: Struct.Types.UnsignedInt
    last: Struct.Types.Bool
    error: Struct.Types.Bool


async def generator(reader, writer):
    data = await reader.read(GeneratorRequest.size)
    request = GeneratorRequest.decode(data)
    if request.amount == 0:
        response = GeneratorResponse.new(value=0, last=True, error=True)
        writer.write(GeneratorResponse.encode(response))
        await writer.drain()
        return
    current = 1
    while current < request.amount:
        response = GeneratorResponse.new(value=current, last=False, error=False)
        writer.write(GeneratorResponse.encode(response))
        await writer.drain()
        current += 1
    response = GeneratorResponse.new(
        value=request.amount, last=True, error=False
    )
    writer.write(GeneratorResponse.encode(response))
    await writer.drain()
    writer.close()


def run_generator(args):
    loop = asyncio.get_event_loop()
    server = loop.run_until_complete(
        asyncio.start_server(generator, args.interface, args.port)
    )
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.close()
        loop.run_until_complete(server.wait_closed())
        loop.close()


class FizzBuzzRequest(Struct):
    byte_order = Struct.ByteOrder.Network
    number: Struct.Types.UnsignedInt
    done: Struct.Types.Bool


class FizzBuzzResponse(Struct):
    byte_order = Struct.ByteOrder.Network
    value: Struct.Types.String(4)
    done: Struct.Types.Bool
    error: Struct.Types.Bool


async def fizz_buzzer(reader, writer):
    while True:
        data = await reader.read(FizzBuzzRequest.size)
        request = FizzBuzzRequest.decode(data)
        if request.done:
            writer.close()
            return
        if request.number == 0:
            response = FizzBuzzResponse.new(value='', error=True, done=True)
            writer.write(FizzBuzzResponse.encode(response))
            await writer.drain()
            return
        responded = False
        if request.number % 3 == 0:
            response = FizzBuzzResponse.new(value=b'Fizz', done=False, error=False)
            writer.write(FizzBuzzResponse.encode(response))
            responded = True
        if request.number % 5 == 0:
            response = FizzBuzzResponse.new(value=b'Buzz', done=False, error=False)
            writer.write(FizzBuzzResponse.encode(response))
            responded = True
        if not responded:
            response = FizzBuzzResponse.new(
                value=bytes(str(request.number), 'utf-8'), done=False, error=False
            )
            writer.write(FizzBuzzResponse.encode(response))
        response = FizzBuzzResponse.new(
            value=b'', done=True, error=False
        )
        writer.write(FizzBuzzResponse.encode(response))
        await writer.drain()


def run_fizz_buzzer(args):
    loop = asyncio.get_event_loop()
    server = loop.run_until_complete(
        asyncio.start_server(fizz_buzzer, args.interface, args.port)
    )
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.close()
        loop.run_until_complete(server.wait_closed())
        loop.close()


def parse_address(address):
    host, port = address.split(':')
    return host, int(port)


def run_client(args):
    gen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    gen_sock.connect(parse_address(args.generator))
    fb_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    fb_sock.connect(parse_address(args.fizzbuzzer))
    gen_request = GeneratorRequest.new(amount=args.amount)
    gen_sock.send(GeneratorRequest.encode(gen_request))
    while True:
        data = gen_sock.recv(GeneratorResponse.size)
        gen_response = GeneratorResponse.decode(data)
        if gen_response.error:
            raise Exception('generator error')
        fb_request = FizzBuzzRequest.new(number=gen_response.value, done=False)
        fb_sock.send(FizzBuzzRequest.encode(fb_request))
        while True:
            data = fb_sock.recv(FizzBuzzResponse.size)
            fb_response = FizzBuzzResponse.decode(data)
            if fb_response.error:
                raise Exception('fizz-buzzer error')
            if fb_response.done:
                break
            sys.stdout.write(fb_response.value.decode('utf-8'))

        sys.stdout.write('\n')
        sys.stdout.flush()
        
        if gen_response.last:
            break

    fb_request = FizzBuzzRequest.new(number=0, done=True)
    fb_sock.send(FizzBuzzRequest.encode(fb_request))
    gen_sock.close()
    fb_sock.close()


def main():
    parser = argparse.ArgumentParser()
    sub_parsers = parser.add_subparsers(dest='cmd')
    sub_parsers.required = True

    generator_parser = sub_parsers.add_parser('generator')
    generator_parser.add_argument('-i', '--interface', default='127.0.0.1')
    generator_parser.add_argument('port')
    generator_parser.set_defaults(func=run_generator)

    fizz_buzzer_parser = sub_parsers.add_parser('fizz-buzzer')
    fizz_buzzer_parser.add_argument('-i', '--interface', default='127.0.0.1')
    fizz_buzzer_parser.add_argument('port')
    fizz_buzzer_parser.set_defaults(func=run_fizz_buzzer)

    client_parser = sub_parsers.add_parser('client')
    client_parser.add_argument('generator')
    client_parser.add_argument('fizzbuzzer')
    client_parser.add_argument('amount', type=int)
    client_parser.set_defaults(func=run_client)

    args = parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()