netron/tools/keras_script.py

''' Keras metadata script '''

import json
import os
import pydoc
import re
import sys

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

def _metadata():

    def parse_docstring(docstring):
        headers = []
        lines = docstring.splitlines()
        indentation = min(filter(lambda s: s > 0, map(lambda s: len(s) - len(s.lstrip()), lines)))
        lines = list((s[indentation:] if len(s) > len(s.lstrip()) else s) for s in lines)
        docstring = '\n'.join(lines)
        tag_re = re.compile('(?<=\n)(Args|Arguments|Variables|Fields|Yields|Call arguments|Raises|Examples|Example|Usage|Input shape|Output shape|Returns|References):\n', re.MULTILINE)
        parts = tag_re.split(docstring)
        headers.append(('', parts.pop(0)))
        while len(parts) > 0:
            headers.append((parts.pop(0), parts.pop(0)))
        return headers

    def parse_arguments(arguments):
        result = []
        item_re = re.compile(r'^   ? ?(\*?\*?\w[\w.]*?\s*):\s', re.MULTILINE)
        content = item_re.split(arguments)
        if content.pop(0) != '':
            raise Exception('')
        while len(content) > 0:
            result.append((content.pop(0), content.pop(0)))
        return result

    def convert_code_blocks(description):
        lines = description.splitlines()
        output = []
        while len(lines) > 0:
            line = lines.pop(0)
            if line.startswith('>>>') and len(lines) > 0 and \
                (lines[0].startswith('>>>') or lines[0].startswith('...')):
                output.append('```')
                output.append(line)
                while len(lines) > 0 and lines[0] != '':
                    output.append(lines.pop(0))
                output.append('```')
            else:
                output.append(line)
        return '\n'.join(output)

    def remove_indentation(value):
        lines = value.splitlines()
        indentation = min(map(lambda s: len(s) - len(s.lstrip()), \
            filter(lambda s: len(s) > 0, lines)))
        lines = list((s[indentation:] if len(s) > 0 else s) for s in lines)
        return '\n'.join(lines).strip()

    def update_argument(schema, name, description):
        if not 'attributes' in schema:
            schema['attributes'] = []
        attribute = next((_ for _ in schema['attributes'] if _['name'] == name), None)
        if not attribute:
            attribute = {}
            attribute['name'] = name
            schema['attributes'].append(attribute)
        attribute['description'] = remove_indentation(description)

    def update_input(schema, description):
        if not 'inputs' in schema:
            schema['inputs'] = [ { 'name': 'input' } ]
        parameter = next((_ for _ in schema['inputs'] \
            if (_['name'] == 'input' or _['name'] == 'inputs')), None)
        if parameter:
            parameter['description'] = remove_indentation(description)
        else:
            raise Exception('')

    def update_output(schema, description):
        if not 'outputs' in schema:
            schema['outputs'] = [ { 'name': 'output' } ]
        parameter = next((param for param in schema['outputs'] if param['name'] == 'output'), None)
        if parameter:
            parameter['description'] = remove_indentation(description)
        else:
            raise Exception('')

    def update_examples(schema, value):
        if 'examples' in schema:
            del schema['examples']
        value = convert_code_blocks(value)
        lines = value.splitlines()
        code = []
        summary = []
        while len(lines) > 0:
            line = lines.pop(0)
            if len(line) > 0:
                if line.startswith('```'):
                    while len(lines) > 0:
                        line = lines.pop(0)
                        if line == '```':
                            break
                        code.append(line)
                else:
                    summary.append(line)
            if len(code) > 0:
                example = {}
                if len(summary) > 0:
                    example['summary'] = '\n'.join(summary)
                example['code'] = '\n'.join(code)
                if not 'examples' in schema:
                    schema['examples'] = []
                schema['examples'].append(example)
                code = []
                summary = []

    def update_references(schema, value):
        if 'references' in schema:
            del schema['references']
        references = []
        reference = ''
        lines = value.splitlines()
        for line in lines:
            if line.lstrip().startswith('- '):
                if len(reference) > 0:
                    references.append(reference)
                reference = line.lstrip().lstrip('- ')
            else:
                if line.startswith('  '):
                    line = line[2:]
                reference = ' '.join([ reference, line.strip() ])
        if len(reference) > 0:
            references.append(reference)
        for reference in references:
            if not 'references' in schema:
                schema['references'] = []
            schema['references'].append({ 'description': reference })

    json_path = os.path.join(os.path.dirname(__file__), '../source/keras-metadata.json')
    with open(json_path, 'r', encoding='utf-8') as file:
        json_root = json.loads(file.read())

    for schema in json_root:
        name = schema['name']
        if 'module' in schema:
            class_name = schema['module'] + '.' + name
            class_definition = pydoc.locate(class_name)
            if not class_definition:
                raise Exception('\'' + class_name + '\' not found.')
            if not class_definition.__doc__:
                raise Exception('\'' + class_name + '\' missing __doc__.')
            headers = parse_docstring(class_definition.__doc__)
            for header in headers:
                key = header[0]
                value = header[1]
                if key == '':
                    description = convert_code_blocks(value)
                    schema['description'] = remove_indentation(description)
                elif key in ('Args', 'Arguments'):
                    arguments = parse_arguments(value)
                    for argument in arguments:
                        update_argument(schema, argument[0], argument[1])
                elif key == 'Call arguments':
                    pass
                elif key == 'Returns':
                    pass
                elif key == 'Input shape':
                    update_input(schema, value)
                elif key == 'Output shape':
                    update_output(schema, value)
                elif key in ('Example', 'Examples', 'Usage'):
                    update_examples(schema, value)
                elif key == 'References':
                    update_references(schema, value)
                elif key == 'Variables':
                    pass
                elif key == 'Raises':
                    pass
                else:
                    raise Exception('')

    with open(json_path, 'w', encoding='utf-8') as file:
        json_data = json.dumps(json_root, sort_keys=False, indent=2)
        for line in json_data.splitlines():
            file.write(line.rstrip() + '\n')

def main(): # pylint: disable=missing-function-docstring
    command_table = { 'metadata': _metadata }
    command = sys.argv[1]
    command_table[command]()

if __name__ == '__main__':
    main()