tl_master.py

# for json
import json
from collections import OrderedDict

# for tl
import re
import struct
import binascii

__all__ = [
    'TLToken',
    'TLTokenJSON',
    'TLTokenTL',
    'TLTokenFileTL',
    'TLTokenFileJSON'
]

class TLToken:
    def __init__(self, name, crc, params, result, is_method=None):
        # name
        self.name = str(name)
        # aka id, should be already "converted" to int8 (struct.unpack'd for example)
        self.crc = int(crc)
        # [{'name': 'foo', 'type': 'long'}, {'name': 'bar', 'type': 'flags.1?int'}, ...]
        self.params = list(params)
        # type
        self.result = str(result)
        if is_method is None:
            is_method = '.' in name
        self.is_method = is_method

    def __str__(self):
        return 'TLToken: ' + self.name

class TLTokenJSON(TLToken):
    def __init__(self, json_obj):
        if not isinstance(json_obj, dict):
            json_obj = json.loads(json_obj)
        
        is_method = 'method' in json_obj
        super().__init__(
            name=json_obj['method' if is_method else 'predicate'],
            crc=json_obj['id'],
            params=json_obj['params'], 
            result=json_obj['type'], 
            is_method=is_method
        )

    @staticmethod
    def dump(inst, as_string=False): 
        if not isinstance(inst.crc, int):
            raise Exception('lol kek cheburek')
        dic = OrderedDict((
            ('id', str(inst.crc)),
            ('method' if inst.is_method else 'predicate', inst.name),
            ('params', inst.params),
            ('type', inst.result)
        ))
        return dic if not as_string else json.dumps(dic)

    def __repr__(self):
        return self.dump(self, True)


VECTOR = 'vector#1cb5c415 {t:Type} # [ t ] = Vector t;'
TL_line_regex = re.compile(r"([a-zA-Z\.0-9_]+)#([0-9a-f]+)([^=]*)=\s*([a-zA-Z\.<>0-9_]+);")
TL_param_regex = re.compile(r"\s([^:]+):(\S+)")

class TLTokenTL(TLToken):   
    def __init__(self, tl_line, is_method=None):
        line = TL_line_regex.match(tl_line)
        
        crc = struct.unpack('>i', 
            binascii.a2b_hex(
                '{:>08}'.format(line.group(2))
            )
        )[0]

        params = [
            {'name': name, 'type': type} 
            for name, type in 
            TL_param_regex.findall(line.group(3))
        ]
        super().__init__(
            name=line.group(1),
            crc=crc,
            params=params, 
            result=line.group(4),
            is_method=is_method
        )

    @staticmethod
    def dump(inst):
        """
        Prints prepending zeroes in crc
        and two spaces in case of empty params
        so output is slightly different from telegram's schemas
        but still correct
        """
        return '{name}#{crc} {params} = {result};'.format(
            name=inst.name,
            crc=str(binascii.b2a_hex(struct.pack('>i', inst.crc)), encoding='ascii'),
            params=' '.join('{name}:{type}'.format(**param) for param in inst.params),
            result=inst.result
        )

    def __repr__(self):
        return self.dump(self)
        


class TLTokenFileTL:
    @staticmethod
    def read(file, no_verify=True):
        def repack(line):
            return re.sub(r'#0+', '#', ' '.join(re.split(r' *', line)))
        def verify(line, token, is_method):
            tl_line = repr(token)
            return repack(line) == repack(tl_line)

        tokens = list()

        methods_now = False
        for l in file:
            l = l.strip()

            if not l or l.startswith('//'):
                # comment
                continue
            elif l == '---functions---':
                methods_now = True
                continue
            elif l == '---types---':
                methods_now = False
                continue
            elif l == VECTOR:
                # vector's line isn't supported, and has to be skipped
                continue
            else:
                try:
                    token = TLTokenTL(l, methods_now)
                    if not (no_verify and verify(l, token, methods_now)):
                        print('MISMATCH:')
                        print(l)
                        print(repack(repr(token)), end='\n\n')
                    else:
                        tokens.append(token)
                except Exception as e:
                    print('ERROR:')
                    print(e)
                    print(l, end='\n\n')
    
        return tokens

    @staticmethod
    def write(file, lst):
        file.write('---types---\n')
        for token in lst:
            if not token.is_method:
                file.write(TLTokenTL.dump(token))
                file.write('\n')

        file.write('---functions---\n')
        for token in lst:
            if token.is_method:
                file.write(TLTokenTL.dump(token))
                file.write('\n')


class TLTokenFileJSON:
    @staticmethod
    def read(file, no_verify=True):
        def process(inp, out, method):
            for token in inp:
                try:
                    out.append(TLTokenJSON(token))
                    assert no_verify or (out[-1].is_method == method)
                except Exception as e:
                    print('ERROR:')
                    print(token)
                    print(e, end='\n\n')
        
        data = json.load(file)

        constructors = list()
        process(data['constructors'], constructors, False)

        methods = list()
        process(data['methods'], methods, True)

        return constructors + methods

    def write(file, lst):
        json.dump(
            OrderedDict((
                ('constructors', [TLTokenJSON.dump(token, as_string=False) for token in lst if not token.is_method]),
                ('methods',      [TLTokenJSON.dump(token, as_string=False) for token in lst if     token.is_method])
            )),
        file)

"""
Usage example -- converting schema from TL to JSON:

from tl_master import *
with open('schema.tl') as tl, open('schema.json', 'w') as json:
    TLTokenFileJSON.write(json, TLTokenFileTL.read(tl))
"""