codegen.py

#!/usr/bin/env python

from __future__ import print_function

import os
import sys

sys.path.append(os.path.join("vendor", "rabbitmq-codegen"))
from amqp_codegen import *

AMQP_ACCEPTED_BY_UPDATE_JSON = "amqp-accepted-by-update.json"

BANNED_CLASSES = ['access', 'tx']
BANNED_FIELDS = {
    'ticket': 0,
    'nowait': 0,
    'capabilities': '',
    'insist': 0,
    'out_of_band': '',
    'known_hosts': '',
}

import codegen_helpers


def pyize(*args):
    a = ' '.join(args).replace('-', '_').replace(' ', '_')
    if a == 'global': a += '_'
    if a == 'type': a += '_'
    return a


def Pyize(*args):
    words = []
    for a in args:
        words.extend(a.split('-'))
    a = ''.join([a.title() for a in words]).replace('-', '').replace(' ', '')
    if a in ['SyntaxError', 'NotImplemented']:
        a = 'AMQP' + a
    return a


def PYIZE(*args):
    return ' '.join(args).replace('-', '_').replace(' ', '_').upper()


def print_constants(spec):
    for c in spec.allClasses():
        for m in c.allMethods():
            print("%-32s= 0x%08X \t# %i,%i %i" % (
                m.u,
                m.method_id,
                m.klass.index, m.index, m.method_id
            ))

    print()
    for c in spec.allClasses():
        if c.fields:
            print("%-24s= 0x%04X" % (
                c.u,
                c.index,))
    print()


def print_decode_methods_map(client_methods):
    print("METHODS = {")
    for m in client_methods:
        print("    %-32s%s," % (
            m.u + ':',
            m.decode,
        ))
    print("}")
    print()


def print_decode_properties_map(props_classes):
    print("PROPS = {")
    for c in props_classes:
        print("    %s: %s, \t# %d" % (
            c.u, c.decode, c.index))
    print("}")
    print()


def print_decode_method(m):
    print("class %s(Frame):" % (m.frame,))
    print("    name = '%s'" % (pyize(m.klass.name + '.' + m.name),))
    print("    method_id = %s" % (pyize('method', m.klass.name, m.name).upper(),))
    if m.hasContent:
        print("    has_content = True")
        print("    class_id = %s" % (m.klass.u,))
    print("")

    print("def %s(data, offset):" % (m.decode,))
    print("    frame = %s()" % (m.frame,))

    fields = codegen_helpers.UnpackWrapper()
    for i, f in enumerate(m.arguments):
        fields.add(f.n, f.t)

    fields.do_print(' ' * 4, "frame['%s']")
    print("    return frame, offset")
    print()


def print_decode_properties(c):
    print("def %s(data, offset):" % (c.decode,))
    print("    props = {}")
    print("    flags, = unpack_from('!H', data, offset)")
    print("    offset += 2")
    print("    assert (flags & 0x01) == 0")
    for i, f in enumerate(c.fields):
        print("    if (flags & 0x%04x): # 1 << %i" % (1 << (15 - i), 15 - i))
        fields = codegen_helpers.UnpackWrapper()
        fields.add(f.n, f.t)
        fields.do_print(" " * 8, "props['%s']")
    print("    return props, offset")
    print()


def _default_params(m):
    for f in m.arguments:
        yield "%s=%r" % (f.n, str(f.defaultvalue) if type(f.defaultvalue) == unicode else f.defaultvalue)
    if m.hasContent:
        yield "user_headers={}"
        yield "payload=''"
        yield "frame_size=None"


def _method_params_list(m):
    for f in m.arguments:
        if not f.banned:
            yield f.n
    if m.hasContent:
        yield 'user_headers'
        yield 'body'
        yield 'frame_size'


def print_encode_method(m):
    print("# %s" % (' '.join(_default_params(m)),))
    print("def %s(%s):" % (m.encode, ', '.join(_method_params_list(m)),))
    for f in [f for f in m.arguments if not f.banned and f.t in ['table']]:
        print("    %s_raw = table.encode(%s)" % (f.n, f.n))

    if m.hasContent:
        print("    props, headers = split_headers(user_headers, %s_PROPS_SET)" % (
            m.klass.name.upper(),))
        print("    if headers:")
        print("        props['headers'] = headers")

    fields = codegen_helpers.PackWrapper()
    fields.add(m.u, 'long')
    for f in m.arguments:
        fields.add(f.n, f.t)
    fields.close()

    if not m.hasContent:
        print("    return ( (0x01,")
        if fields.group_count() > 1:
            print("              join_as_bytes((")
            fields.do_print(' ' * 16, '%s')
            print("              ))")
        else:
            fields.do_print(' ' * 16, '%s')
        print("           ), )")
    else:
        print("    return [ (0x01,")
        print("              join_as_bytes((")
        fields.do_print(' ' * 16, '%s')
        print("              ))")
        print("           ),")
        print("           %s(len(body), props)," % (m.klass.encode,))
        print("        ] + encode_body(body, frame_size)")


def print_encode_properties(c):
    print("%s_PROPS_SET = {" % (c.name.upper(),))
    for f in c.fields:
        print('    "%s", %s # %s' % (f.n, ' ' * (16 - len(f.n)), f.t))
    print("}")
    print()
    print("ENCODE_PROPS_%s = {" % (c.name.upper(),))
    for i, f in enumerate(c.fields):
        pn = pyize(f.name)
        print("    '%s': (" % (pn,))
        print("        %i," % (i,))
        print("        0x%04x, # (1 << %i)" % (1 << (15 - i), 15 - i,))

        fields = codegen_helpers.PackWrapper()
        if f.t not in ['table']:
            fields.add("val", f.t)
        else:
            fields.add("table.encode(val)", f.t, nr='%s')
        fields.close()
        if len(fields.fields) > 1:
            print(' ' * 8 + "lambda val: join_as_bytes((")
            fields.do_print(' ' * 16, '%s')
            print(' ' * 8 + ')) ),')
        else:
            print('  ' * 4 + 'lambda val:', end=' ')
            fields.do_print('', '%s', comma=False)
            print('        ),')

    print("}")
    print()
    print("def %s(body_size, props):" % (c.encode,))
    print("    pieces = ['']*%i" % (len(c.fields),))
    print("    flags = 0")
    print("    enc = ENCODE_PROPS_%s" % (c.name.upper(),))
    print()
    print("    for key in %s_PROPS_SET & set(futils.viewkeys(props)):" % \
          (c.name.upper(),))
    print("        i, f, fun = enc[key]")
    print("        flags |= f")
    print("        pieces[i] = fun(props[key])")
    print("")
    print("    return (0x02, join_as_bytes((")
    print("        pack('!HHQH',")
    print("              %s, 0, body_size, flags)," % (c.u,))
    print("        join_as_bytes(pieces),")
    print("        ))")
    print("        )")


def GetAmqpSpec(spec_path, accepted_by_udate):
    spec = AmqpSpec(spec_path)

    for c in spec.allClasses():
        c.banned = bool(c.name in BANNED_CLASSES)
        c.u = PYIZE('CLASS', c.name)

    spec.classes = filter(lambda c: not c.banned, spec.classes)

    for c in spec.allClasses():
        for m in c.allMethods():
            m.u = PYIZE('METHOD', m.klass.name, m.name)
            m.method_id = m.klass.index << 16 | m.index
            m.decode = pyize('decode', m.klass.name, m.name)
            m.encode = pyize('encode', m.klass.name, m.name)
            m.frame = Pyize('frame', m.klass.name, m.name)

            try:
                m.accepted_by = accepted_by_udate[c.name][m.name]
            except KeyError:
                print(" [!] Method %s.%s unknown! Assuming " \
                      "['server', 'client']" % (c.name, m.name), file=sys.stderr)
                m.accepted_by = ['server', 'client']

            for f in m.arguments:
                f.t = spec.resolveDomain(f.domain)
                f.n = pyize(f.name)
                f.banned = bool(f.name in BANNED_FIELDS)

    for c in spec.allClasses():
        if c.fields:
            c.decode = pyize('decode', c.name, 'properties')
            c.encode = pyize('encode', c.name, 'properties')
            for f in c.fields:
                f.t = spec.resolveDomain(f.domain)
                f.n = pyize(f.name)
    return spec


def generate_spec(spec_path):
    accepted_by_udate = json.loads(file(AMQP_ACCEPTED_BY_UPDATE_JSON).read())
    return GetAmqpSpec(spec_path, accepted_by_udate)


def main(spec_path):
    spec = generate_spec(spec_path)
    print("""# Autogenerated - do not edit
from __future__ import absolute_import
import calendar
import datetime
from struct import pack, unpack_from

import future.utils as futils

from . import table
from .compat import as_str, join_as_bytes

""")
    print("PREAMBLE = b'AMQP\\x00\\x%02x\\x%02x\\x%02x'" % (
        spec.major, spec.minor, spec.revision))
    print()
    print_constants(spec)
    print()

    props_classes = [c for c in spec.allClasses() if c.fields]

    client_methods = [m for m in spec.allMethods() if 'client' in m.accepted_by]
    print('''
class Frame(dict):
    has_content = False
    is_error = False

''')
    for m in client_methods:
        print_decode_method(m)
        print()
    print()
    print_decode_methods_map(client_methods)
    print()
    for c in props_classes:
        print_decode_properties(c)
        print()
    print_decode_properties_map(props_classes)
    print()

    server_methods = [m for m in spec.allMethods() if 'server' in m.accepted_by]
    for m in server_methods:
        print_encode_method(m)
        print()

    for c in props_classes:
        print_encode_properties(c)
        print()

    print("""
def split_headers(user_headers, properties_set):
    props = {}
    headers = {}
    for key, value in futils.viewitems(user_headers):
        if key in properties_set:
            props[key] = value
        else:
            headers[key] = value
    return props, headers

def encode_body(body, frame_size):
    limit = frame_size - 7 - 1   # spec is broken...
    r = []
    while body:
        payload, body = body[:limit], body[limit:]
        r.append( (0x03, payload) )
    return r
""", end="")


def spec_exceptions(spec_path):
    spec = generate_spec(spec_path)
    print("""# Autogenerated - do not edit
""")
    err_constants = [(name, value, klass)
                     for name, value, klass in spec.constants
                     if klass in ('hard-error', 'soft-error')]
    print("class AMQPError(Exception): pass")
    print("class AMQPSoftError(AMQPError): pass")
    print("class AMQPHardError(AMQPError): pass")
    print()
    for name, value, klass in err_constants:
        print("class %s(AMQP%s):" % (Pyize(name),
                                     Pyize(klass)))
        print("    reply_code = %s" % (value,))
    print()
    print("ERRORS = {")
    for name, value, klass in err_constants:
        print("    %i: %s," % (value, Pyize(name)))
    print("}")


if __name__ == "__main__":
    do_main_dict({"spec": main, 'spec_exceptions': spec_exceptions})