bencode.py

from typing import Iterator
from functools import wraps, singledispatch


def wrap(begin: bytes=None, end: bytes=b'e'):
    def outer_wrapper(func):
        @wraps(func)
        def inner_wrapper(*args, **kwargs):
            yield begin
            yield from func(*args, **kwargs)
            yield end
        return inner_wrapper
    return outer_wrapper


@singledispatch
def type_encode(x):
    raise NotImplementedError


@type_encode.register(int)
@wrap(b'i')
def encode_int(x: int) -> Iterator[bytes]:
    yield str(x).encode('utf-8')


@type_encode.register(float)
@wrap(b'f')
def encode_float(x: float) -> Iterator[bytes]:
    yield repr(x).replace('e', 'E').encode('utf-8')


@type_encode.register(bool)
def encode_bool(x: bool) -> Iterator[bytes]:
    yield from encode_int(int(bool(x)))


@type_encode.register(bytes)
def encode_bytes(x: bytes) -> Iterator[bytes]:
    yield b'%d:%s' % (len(x), x)


@type_encode.register(str)
def encode_string(x: str) -> Iterator[bytes]:
    yield from encode_bytes(x.encode('utf-8'))


@type_encode.register(list)
@wrap(b'l')
def encode_list(x: list) -> Iterator[bytes]:
    for i in x:
        yield from bencode(i)


@type_encode.register(dict)
@wrap(b'd')
def encode_dict(x: dict) -> Iterator[bytes]:
    for k, v in sorted(x.items()):
        yield from encode_string(str(k))
        yield from bencode(v)


def bencode(x) -> Iterator[bytes]:
    yield from type_encode.dispatch(type(x))(x)


print(b''.join(bencode(1)))
print(b''.join(bencode([1,2,3])))
print(b''.join(bencode({1:2,3:4})))
print(b''.join(bencode('test')))
print(b''.join(bencode([True, 1, 1])))