generate_snippets.py

#!/usr/bin/env python3

import json
import os
import subprocess
import io
import pathlib
import sys
import argparse

import jinja2

from make_providers_dirs import make_providers_dirs
from color import Color


JINJA_TEMPLATE = """# -*- mode: snippet -*-
# name: [[thing_short]]_[[name]]
# key: [[thing_short]]_[[name]]
# --

[[thing]] "[[name]]" "${1:this}" {
[% for pos, a in enumerated_attributes %]
    [[a]] = "${[[pos]]:[[a]]}"
[% endfor %]
}
"""


def get_terraform_schema(provider_path):

    cmd = ["terraform", "providers", "schema", "-json"]

    try:

        out_bytes = subprocess.check_output(
            cmd,
            cwd=provider_path,
            stderr=subprocess.STDOUT,
        )

        return out_bytes

    except subprocess.CalledProcessError as e:

        if "Could not load plugin" in e.output.decode(sys.stdout.encoding):

            try:
                print()
                print(f"{Color.YELLOW}Running {Color.BOLD}`terraform init´{Color.END}")
                subprocess.check_output(
                    ["terraform", "init"], cwd=provider_path, stderr=subprocess.STDOUT
                )
            except subprocess.CalledProcessError as e:
                print(f"{Color.RED}This one failed...{Color.END}")
                print(e.output.decode(sys.stdout.encoding))
                return

            try:
                print(
                    f"{Color.YELLOW}Trying to run `terraform providers schema -json` again{Color.END}"
                )
                out_bytes = subprocess.check_output(
                    cmd,
                    cwd=provider_path,
                    stderr=subprocess.STDOUT,
                )
            except subprocess.CalledProcessError as e:
                print(f"{Color.RED}This one failed...{Color.END}")
                print(e.output.decode(sys.stdout.encoding))
                return

            return out_bytes


def make_jinja2_template():

    environment = jinja2.Environment(
        loader=jinja2.BaseLoader(),
        block_start_string="[%",
        block_end_string="%]",
        variable_start_string="[[",
        variable_end_string="]]",
        autoescape=False,
        trim_blocks=True,
    )

    return environment.from_string(JINJA_TEMPLATE)


def get_saved_schema(path, current_version):

    for dir_entry in os.scandir(path):
        if current_version in dir_entry.name:
            with open(dir_entry.path) as f:
                return json.load(f).get("provider_schemas")
    return False


def current_version(path):

    import re

    # version = "0.0.4"
    with open(path, "r") as f:
        for line in f:
            maybe_match = re.search(r'\s+version\s*=\s*"([^"]+)"', line)
            if maybe_match:
                return maybe_match.group(1)

    return False


def generate_snippets(args):
    """
    Scans through the providers directory (generated by `make_providers_dirs`). Runs `terraform init` if needed. Runs
    `terraform providers schema -json`. Puts the data through a Jinja2 template and writes the snippet to disk.
    """

    for dir_entry in os.scandir("providers"):

        print()
        print("=" * 10)
        print(
            f"About to run {Color.BOLD}`terraform providers schema -json`{Color.END} in {dir_entry.path}"
        )

        version = current_version(os.path.join(dir_entry.path, dir_entry.name + ".tf"))

        provider_schemas = get_saved_schema(dir_entry.path, version)

        if not provider_schemas:
            json_schema_bytes = get_terraform_schema(dir_entry.path)
            if json_schema_bytes:
                f = io.BytesIO(json_schema_bytes)
                provider_schemas = json.load(f).get("provider_schemas")

                schema_filename = os.path.join(
                    dir_entry.path, f"{dir_entry.name}-schema-{version}.json"
                )

                with open(schema_filename, "wb") as s:
                    print(f"Writing {schema_filename} to disk.")
                    s.write(f.getbuffer())

        if provider_schemas:

            # Get the first dict key
            schemas = next(iter(provider_schemas.values()))

            resource_schemas = schemas.get("resource_schemas", {})

            data_source_schemas = schemas.get("data_source_schemas", {})

            template = make_jinja2_template()

            print(
                f"Loaded schema JSON for {dir_entry.name} - templating and writing snippet files"
            )

            for resource_name, resource_schema in resource_schemas.items():

                attributes = resource_schema.get("block").get("attributes")
                attributes = [
                    a for a in attributes if attributes.get(a).get("required") == True
                ]

                snippet = template.render(
                    thing="resource",
                    thing_short="r",
                    name=resource_name,
                    enumerated_attributes=enumerate(attributes, start=2),
                )

                filename = f"snippets/terraform-mode/{dir_entry.name}/resource/r_{resource_name}"
                os.makedirs(os.path.dirname(filename), exist_ok=True)
                with open(filename, "w") as f:
                    f.write(snippet)

            for data_source_name, data_source_schema in data_source_schemas.items():

                attributes = data_source_schema.get("block").get("attributes")
                attributes = [
                    a for a in attributes if attributes.get(a).get("required") == True
                ]

                snippet = template.render(
                    thing="data",
                    thing_short="d",
                    name=data_source_name,
                    enumerated_attributes=enumerate(attributes, start=2),
                )

                filename = f"snippets/terraform-mode/{dir_entry.name}/data/d_{data_source_name}"
                os.makedirs(os.path.dirname(filename), exist_ok=True)
                with open(filename, "w") as f:
                    f.write(snippet)

        print("=" * 10)

    pathlib.Path("snippets/terraform-mode/.yas-make-groups").touch()


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.set_defaults(func=lambda args: parser.print_help())
    subparsers = parser.add_subparsers()

    # create the parser for the "generate_snippets" command
    parser_foo = subparsers.add_parser("generate_snippets")
    parser_foo.set_defaults(func=generate_snippets)

    # create the parser for the "make_provider_dirs" command
    parser_bar = subparsers.add_parser("make_provider_dirs")
    parser_bar.set_defaults(func=make_providers_dirs)

    # parse the args and call whatever function was selected
    args = parser.parse_args()
    args.func(args)