Patch IDEA run configurations with environment variables from a file

patch_idea_env.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import unicode_literals


import argparse
import re
import os
import subprocess
import sys
import codecs
import shlex
import shutil

from xml.etree import ElementTree
from datetime import datetime


ENV_VAR_RE = re.compile(r'(?:export\s+)?(?P<key>[A-Z_][A-Z0-9_]*)=(?P<value>.*)')

BOLD = subprocess.check_output('tput bold', shell=True)
RED_FG = subprocess.check_output('tput setaf 1', shell=True)
CYAN_FG = subprocess.check_output('tput setaf 6', shell=True)
GREEN_FG = subprocess.check_output('tput setaf 2', shell=True)
RESET = subprocess.check_output('tput sgr0', shell=True)


def find_env_file(env_file):
    if os.path.exists(env_file):
        return env_file


def find_idea_dir(env_file, hint):
    if hint and os.path.isdir(hint):
        return hint

    candidate = os.path.realpath(env_file)
    previous = None

    for _ in range(0, 100):
        candidate = os.path.dirname(candidate)

        if candidate == previous:
            break

        previous = candidate
        idea = os.path.join(candidate, '.idea')

        if os.path.isdir(idea):
            return idea


def parse_env(env_file):
    out = {}

    with codecs.open(env_file, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            m = ENV_VAR_RE.match(line)

            if m:
                try:
                    val = shlex.split(m.group('value'))[0]
                except IndexError:
                    val = ''

                out[m.group('key')] = val

    return out


def patch_workspaces(idea_path, filter_, env, dry_run):
    ws_path = os.path.join(idea_path, 'workspace.xml')
    config_path = './/component[@name="RunManager"]/configuration[@default="false"]'

    try:
        tree = ElementTree.parse(ws_path)
    except IOError:
        print_err(RED_FG + b'file(' + BOLD + ws_path.encode() + RESET + RED_FG +
                  b') does not exist' + RESET)
        exit(1)
    except ElementTree.ParseError:
        print_err(RED_FG + b'file(' + BOLD + ws_path.encode() + RESET + RED_FG +
                  b') could not be parsed' + RESET)
        exit(1)
    else:
        print_std(BOLD + b'[ Patching Workspace ]' + RESET)

        for node in tree.findall(config_path):
            name = node.attrib.get('name')

            if name and (not filter_ or filter_ in name):
                print_std('  → patching run config('.encode() + BOLD + name.encode() + RESET + b')')

                envs = node.find('envs')

                if not envs:
                    envs = ElementTree.SubElement(node, 'envs')

                for k, v in env.items():
                    env_node = envs.find('env[@name="{}"]'.format(k))
                    attr = {'name': k, 'value': v}

                    if not env_node:
                        ElementTree.SubElement(envs, 'env', attr)
                    else:
                        env_node.attrib.update(attr)

        if not dry_run:
            backup = ws_path + '.bak.' + datetime.now().isoformat()
            shutil.copy2(ws_path, backup)
            tree.write(ws_path)


def parse():
    parser = argparse.ArgumentParser(description='Patches .idea run configuration with an '
                                                 'environment file')

    parser.add_argument('env_file', help='Path to an environment file. The syntax is something '
                                         'you could source in a shell in order to load environment '
                                         'variables.',
                        default='./env')

    parser.add_argument('-f', '--filter', help='Only affect run configurations that contain this '
                                               'text')

    parser.add_argument('-i', '--idea', help='Manually specify .idea location')

    parser.add_argument('-d', '--dry-run', action='store_true', help='Do not really make changes')

    return parser.parse_args()


def print_err(s):
    sys.stdout.buffer.write(s + b'\n')
    sys.stdout.flush()


def print_std(s):
    sys.stdout.buffer.write(s + b'\n')
    sys.stdout.flush()


def main():
    args = parse()

    env_file = find_env_file(args.env_file)

    if not env_file:
        print_err(RED_FG + b'Could not find env file' + RESET)
        exit(1)

    idea_dir = find_idea_dir(env_file, args.idea)

    if not idea_dir:
        print_err(RED_FG + b'Could not find .idea dir' + RESET)

    print_std(BOLD + b'[ Found Files ]' + RESET)
    print_std(b'  - env = ' + BOLD + env_file.encode() + RESET)
    print_std(b'  - .idea = ' + BOLD + idea_dir.encode() + RESET)
    print_std(b'')

    env = parse_env(env_file)
    patch_workspaces(idea_dir, args.filter, env, args.dry_run)
    print_std(b'')

    if args.dry_run:
        print_std(CYAN_FG + BOLD + b'Done!' + RESET + CYAN_FG + b' (dry)' + RESET)
    else:
        print_std(GREEN_FG + BOLD + b'Done!' + RESET)


if __name__ == '__main__':
    main()