docgen

#!/usr/bin/env python3
"""Generate markdown from terraform module."""

import re
import argparse
import os
from pprint import pprint


class TFDoc:
    """stateful class to parse terraform docstrings."""

    def __init__(self):
        """Set private vars to default."""
        self._variables = {}
        self._outputs = {}
        self._sections = {}
        self._outline = {'title': "", 'text': ""}
        self._footers = {}

        self._block = None
        self._unclassified_text = ""
        self._indent = ""

        self._multiline_comment = False
        self._in_var = False
        self._in_type = False
        self.debug = False

    def set_indent(self, indent):
        """
        Set current default indent for text to 'indent'
        """
        self._indent = indent

    def variable(self, name):
        """
        Add a variable or Argument pass the variable name.

        This starts a variable block so calls to description, type, default, or add text will be
        associated with this variable.
        """
        self._variables[name] = {'description': "", 'type': "", 'default': None, 'text': "", 'object': {} }
        self._block = self._variables[name]
        self._indent = 2

    def description(self, description):
        """Add description text to the active variable."""
        self._block['description'] = description

    def type(self, type):
        """Set the variable type."""
        if '{' in type:
            self._block['type'] = type[:-2] + ')'
            self._in_type = True
        else:
            self._block['type'] = type

    def object(self, key, _type, description=''):
        """Add object definition for complex types"""
        self._block['object'][key] = {'type': _type, 'description': description}

    def default(self, default):
        """Define the default value for the current variable."""
        self._block['default'] = default

    def output(self, name):
        """
        Add output or Attribute.

        this starts an output block so calls to add_txt are associated with this Attribute.
        """
        self._outputs[name] = {'text': ""}
        self._block = self._outputs[name]

    def section(self, name):
        """
        Start an additional section with headding name.

        These sections appear between the Argument and attribute sections.
        """
        self._sections[name] = {'text': ""}
        self._block = self._sections[name]

    def footer(self, name):
        """
        Start a footer block.

        Footer blocks appear after the Attribute block.
        """
        self._footers[name] = {'text': ""}
        self._block = self._footers[name]

    def outline(self, title):
        """
        Start the outline block.

        The outline block appears first in the docs.
        The title provided will be the title of the whole document.
        """
        self._outline['title'] = title
        self._block = self._outline

    def add_text(self, text):
        """Add text tothe current block or the unclissified block if no block is active."""
        if not self._block:
            self._unclassified_text += text + '\n'
        else:
            self._block['text'] += " " * self._indent + text + '\n'

    def end_block(self):
        """
        End the current block.

        prevents stray text being added to this block and resets indent.
        """
        self._block = None
        self._indent = 0

    def __str__(self):
        """Output the masterpiece."""
        text = f'''# {self._outline['title']}\n\n{self._outline['text']}'''

        text += '## Argument Reference\n\n'
        for var in sorted(self._variables.keys()):
            attribute = self._variables[var]
            default = f'''(optional, default: {attribute['default']})''' if attribute['default'] else '(Required)'
            _object = attribute['object']
            text += f'''- ```{var}``` - {default} [{attribute['type']}] {attribute['description']} '''
            if _object:
                text += '\n##### object'
                text += '''
| key | type | description |
|-----|------|-------------|
'''
                for key, definition in _object.items():
                    text += f'''|`{key}` | {definition['type']} | {definition['description']} |\n'''
            text += f'''{attribute['text']}\n\n'''

        for (section, attribute) in self._sections.items():
            text += f'''## {section}\n\n{attribute['text']}\n'''
        if len(self._outputs) > 0:
            text += '## Attributes Reference\n\n'
            for var in sorted(self._outputs.keys()):
                attribute=self._outputs[var]
                text += f'''- ```{var}``` {attribute['text']}\n'''

        for (section, attribute) in self._footers.items():
            text += f'''## {section}\n\n{attribute['text']}\n'''

        text += self._unclassified_text + '\n'

        return text

    def __add__(self, other):
        if type(other) == str:
            self.process_line(other)
            return self
        elif type(other) == list:
            for line in other:
                self.process_line(line)
            return self
        else:
            return NotImplemented

    def process_line(self, raw_line):
        """
        Process a single line from a terraform file
        """
        line = raw_line.strip()
        not self.debug or print(line)

        if self._in_type:
            if '}' in line:
                self._in_type = False
                return
            key, rest = line.split('=', 1)
            (_type, description) = rest.split('#=') if '#=' in rest else (rest,'')
            self.object(key.strip(), _type.strip(), description.strip())
            return

        if words := line.split(None, 1):

            if words[0] == '#=':
                self.add_text(line[3:] if len(words) > 1 else "\n")
                return

            if words[0] == '#=OUTLINE=':
                self.outline(words[1])
                return

            if words[0] == '#=SECTION=':
                self.section(words[1])
                return

            if words[0] == '#=FOOTER=':
                self.footer(words[1])
                return

            if words[0] == "#=INDENT=":
                self.set_indent(int(words[1]))

            if words[0] == '#==' or words[0] == "}":
                self.end_block()
                self._in_var = False
                return

            if words[0] == '/*=' and '*/' not in line:
                self._multiline_comment = True
                return

            if self._multiline_comment and '*/' in line:
                self._multiline_comment = False
                return

            if self._multiline_comment:
                self.add_text(raw_line)
                return

            # Language components
            # The following words might incidentally exist in multiline comments
            # So must go last
            if words[0] == 'variable':
                name = re.search('"(.*)"', words[1])
                self.variable(name.group(1))
                self._in_var = True

            if words[0] == 'description' and self._in_var:
                description = re.search('"(.*)"', words[1])
                self.description(description.group(1))

            # TODO Handle multiline types
            if words[0] == 'type' and self._in_var:
                self.type(words[1].split('=')[1].strip())

            if words[0] == 'default' and self._in_var:
                self.default(words[1])

            if words[0] == 'output':
                name = re.search('"(.*)"', words[1])
                self.output(name.group(1))
                in_var = True

        elif self._multiline_comment:
            self.add_text(raw_line)


parser = argparse.ArgumentParser(description='Generate terraform docs', prog='PROG', usage='%(prog)s [options]')
parser.add_argument('source', nargs='+')
parser.add_argument('--out', nargs=1)
parser.add_argument('--verbose', '-v', action='store_true')
parser.add_argument('--debug', '-d', action='store_true')

args = parser.parse_args()

docs = TFDoc()
docs.debug = args.debug

for file in args.source:
    if os.path.isdir(file):
        continue
    if args.verbose: print(f'>>> {file}')
    with open(file) as f:
        docs += f.readlines()


if args.out:
    file = open(args.out[0], 'w+')
    print(docs, file=file)
    file.close()

else:
    print(docs)