tg-py

#!/usr/bin/env python

# -*- coding: utf-8 -*-

from __future__ import print_function

import argparse
import functools
import io
import json
import os
import platform
import re
import sys

if sys.version_info[0] < 3:
    import ConfigParser as configparser

    input = raw_input  # noqa:F821
else:
    import configparser

    unicode = str

try:
    from shlex import quote
except ImportError:
    from pipes import quote


HOME_DIR = os.path.expanduser("~")
CONFIG_DIR = os.path.join(HOME_DIR, ".tget")
CONFIG = os.path.join(CONFIG_DIR, "config.cfg")
STORE = os.path.join(CONFIG_DIR, "store.json")

config = configparser.SafeConfigParser()
config.read(CONFIG)

TOP_LEVEL_COMMANDS = {"ls", "search", "put", "del", "update", "reset", "uninstall", "-h", "--help"}
PROTECTED_COMMANDS = TOP_LEVEL_COMMANDS.union({"KEY"})

OS = platform.system()


def access_store(fn):
    """Add a `data` positional arg to the decorated function.
    """

    @functools.wraps(fn)
    def wrapped(*args, **kwargs):
        if not os.path.exists(STORE):
            _safe_write({})

        try:
            with open(STORE) as file_:
                data = json.load(file_)
        except ValueError:
            _echo("Unable to access the `store` - the file may be corrupted!", pre_nl=True)
            _echo("To recreate it anew, run the `tg reset` command.")
            _echo("Store location: {}".format(STORE))
            sys.exit(1)

        return fn(data, *args, **kwargs)

    return wrapped


@access_store
def ls(data):
    DIVIDER = "-" * 40

    _echo("")

    for key in sorted(data.keys()):
        _echo(key, post_nl=False)
        _echo(DIVIDER, post_nl=False)
        _echo(data[key])


@access_store
def search(data):
    keys = sorted(data.keys())
    echo_keys = 'echo "' + "\n".join(keys) + '"'

    fzf_height = r"50%"
    fzf_preview_command = r"tg-py {} | source /dev/stdin"
    fzf_preview_window = r"right:60%"
    show_available_keys = "fzf --header='Select the key' --height {} --reverse --preview '{}' --preview-window {}".format(
        fzf_height, fzf_preview_command, fzf_preview_window
    )

    surround_key = "sed 's/.*/tg & /'"
    trim_nl = r"tr -d '\n'"
    copy_to_clipboard = "pbcopy" if OS == "Darwin" else "xclip"

    statement = " | ".join((echo_keys, show_available_keys, surround_key, trim_nl, copy_to_clipboard))

    print(statement)
    _echo("Copied to your clipboard!", pre_nl=True)


@access_store
def get(data, key, execute, verbose, priority_vars):
    value = data.get(key)

    if value is None:
        _echo("Nothing found", pre_nl=True)
        sys.exit(1)

    value_placeholders = set(re.findall(r"\$\w+", value))

    try:
        priority_var_map = dict(priority_var.split("=", 1) for priority_var in priority_vars)
    except ValueError:
        _echo("Passed in ENV-VARs must be in the ENVVAR=value form.")
        sys.exit(1)

    for value_placeholder in value_placeholders:
        no_dollar_value_placeholder = value_placeholder[1:]

        if no_dollar_value_placeholder in priority_var_map:
            value = value.replace(value_placeholder, priority_var_map[no_dollar_value_placeholder])
            continue

        if no_dollar_value_placeholder in os.environ:
            value = value.replace(value_placeholder, os.environ[no_dollar_value_placeholder])

    if not execute or verbose:
        _echo(value, pre_nl=True)

    if execute:
        print(value)


@access_store
def put(data, key, value):
    if key in PROTECTED_COMMANDS:
        _echo("Bad Key!", pre_nl=True)
        sys.exit(1)

    if value is None:
        _echo("Value: (Hit <ENTER> twice to end input)", pre_nl=True, post_nl=True, flush=True)
        value = ""
        current_value = None
        previous_value = None
        while True:
            if current_value == "" and previous_value == "":
                break

            current_value = input()
            value += current_value
            value += "\n"
            previous_value = current_value

        # Remove `\n\n` from the value.
        value = value[:-2]

    if not value:
        _echo("Bad Value!")
        sys.exit(1)

    data[key] = value.strip()

    _safe_write(data)

    _echo("STORED")


@access_store
def delete(data, key):
    try:
        value = data.pop(key)
    except KeyError:
        _echo("Nothing found", pre_nl=True)
        sys.exit(1)

    _safe_write(data)

    _echo("The following was deleted:", pre_nl=True)
    _echo("Key: " + key)
    _echo("Value:\n " + value)


def update():
    print("curl https://raw.githubusercontent.com/oakhan3/tget/master/install | python")

    _echo("Updated!.", pre_nl=True)


def reset():
    _echo("This will ERASE any currently stored data, type RESET to confirm: ", pre_nl=True, post_nl=False, flush=True)

    value = input()

    if value == "RESET":
        _safe_write({})
        _echo("Deleted all stored key-value pairs.", pre_nl=True)
    else:
        _echo("Aborted.", pre_nl=True)


def uninstall():
    executable_dir = os.path.join(HOME_DIR, ".local/bin/tg-py")
    print("yes | rm {}".format(executable_dir))

    _echo("Make sure to:", pre_nl=True)
    _echo(" - Delete your configuration located at: '{}'".format(CONFIG_DIR))
    _echo(" - Delete the 'tg' function in your shell config file.")
    _echo("Successfully uninstalled 'tg'")


# Internal Functions
def _safe_write(data):
    with io.open(STORE, "w", encoding="utf8") as outfile:
        data_str = json.dumps(data, indent=4, separators=(",", ": "))
        outfile.write(unicode(data_str))


def _echo(value, pre_nl=False, post_nl=True, flush=False):
    if pre_nl:
        print("echo")

    print("echo " + quote(value))

    if post_nl:
        print("echo")

    if flush:
        sys.stdout.flush()


def _create_main_parser(description):
    # Main Parser
    main_parser = argparse.ArgumentParser(prog="tg", description=description)
    _echo_help(main_parser)

    sub_parsers = main_parser.add_subparsers(metavar="command")

    # GET
    # NOTE: The `get_help_parser` sub-parser is never actually used, it only serves to automatically produce documentation for the main parser.
    get_help_parser = sub_parsers.add_parser("KEY", help="Retrieve a value. See tg KEY -h for more info.")
    _add_get_parser_args(get_help_parser)
    _echo_help(get_help_parser)

    # PUT
    put_parser = sub_parsers.add_parser("put", help="Store a new key-value pair.")
    put_parser.add_argument("key", help="Key used to retrieve a value.")
    put_parser.add_argument("-v", "--value", help="Non-interactive way to set a value.")
    put_parser.set_defaults(fn=put)
    _echo_help(put_parser)

    # DEL
    delete_parser = sub_parsers.add_parser("del", help="Delete a snippet.")
    delete_parser.add_argument("key")
    delete_parser.set_defaults(fn=delete)
    _echo_help(delete_parser)

    # LS
    ls_parser = sub_parsers.add_parser("ls", help="See all stored snippets.")
    ls_parser.set_defaults(fn=ls)
    _echo_help(ls_parser)

    # SEARCH
    search_parser = sub_parsers.add_parser("search", help="Search all stored snippets.")
    search_parser.set_defaults(fn=search)
    _echo_help(search_parser)

    # UPDATE
    update_parser = sub_parsers.add_parser("update", help="Update tg!")
    update_parser.set_defaults(fn=update)
    _echo_help(update_parser)

    # RESET
    reset_parser = sub_parsers.add_parser(
        "reset", help="Re-initialize the data-store. This will ERASE any currently stored data, use with care!"
    )
    reset_parser.set_defaults(fn=reset)
    _echo_help(reset_parser)

    # UNINSTALL
    uninstall_parser = sub_parsers.add_parser("uninstall")
    uninstall_parser.set_defaults(fn=uninstall)
    _echo_help(uninstall_parser)

    return main_parser


def _create_get_parser(description):
    get_parser = argparse.ArgumentParser(prog="tg", description=description)
    _add_get_parser_args(get_parser)
    _echo_help(get_parser)

    return get_parser


def _add_get_parser_args(parser):
    parser.add_argument("key", metavar="KEY")
    parser.add_argument(
        "-e", "--exec", help="Execute the retrieved value in the current shell.", dest="execute", action="store_true"
    )
    parser.add_argument(
        "-v",
        "--verbose",
        help="Echo the executed command when the -e option is passed",
        dest="verbose",
        action="store_true",
    )
    parser.add_argument(
        "priority_vars",
        metavar="ENVVAR=value",
        nargs="*",
        help="Optional Environment Variables to interpolate into the output. These will take precedence over any ACTUAL env-vars in your current session.",
    )


def _echo_help(parser):
    parser.print_help = lambda: _echo(parser.format_help(), post_nl=False)


if __name__ == "__main__":
    description = "A super mini-utility to store and retrieve random snippets in your terminal!"
    main_parser = _create_main_parser(description)
    get_parser = _create_get_parser(description)

    if not os.path.isdir(CONFIG_DIR):
        os.mkdir(CONFIG_DIR)

    if not os.path.exists(CONFIG):
        with open(CONFIG, "w") as config_file:
            config_file.write("[get]\ninterpolate_env_vars=True")

    args = sys.argv[1:]

    if len(args):
        if args[0] in TOP_LEVEL_COMMANDS:
            parsed_args = vars(main_parser.parse_args(args))
            fn = parsed_args.pop("fn")
            fn(**parsed_args)
            sys.exit()

        parsed_args = vars(get_parser.parse_args(args))
        get(**parsed_args)
        sys.exit()

    else:
        main_parser.print_help()
        sys.exit(1)