branch-description

#!/usr/bin/python3

import json
import os
import re
import sys
from lib import common

configs = common.Configs('branch-descriptions')

DESCRIPTIONS_JSON_FILE = 'descriptions.json'

def main():
    with common.DispatchingArgParser('short, per-branch descriptions') as parser:
        with parser.subcommand(cmd_set, 'set', help='sets the description for a branch') as cmd:
            cmd.add_argument('description', nargs='?')
            cmd.add_argument('-b', '--branch', help='the branch name to set a description for; defaults to current branch')
            cmd.add_argument('--color', help='The ASCII color code to display the description. For example, "33;48" for red on a gray background')
            cmd.add_argument('--gh', action='store_true', help="Append a GitHub PR URL, if there's exactly one")
        with parser.subcommand(cmd_get, 'get', help='gets the description for a branch') as cmd:
            cmd.add_argument('-b', '--branch', help='the branch name to set a description for; defaults to current branch')
            cmd.add_argument('--raw', action='store_true', help="don't linkify URLs")
        with parser.subcommand(cmd_unset, 'unset', help='unsets the description for a branch') as cmd:
            cmd.add_argument('-b', '--branch', help='the branch name to set a description for; defaults to current branch')
        with parser.subcommand(cmd_p10_prompt, 'p10k_prompt', help='outputs a zsh script to define a powerlevel10k prompt. feed this to eval') as cmd:
            pass
        with parser.subcommand(cmd_git_pager, 'git-pager', help=f"operates as a git pager: git config --global --add pager.branch '{common.SCRIPT_NAME} git-pager'"):
            pass

def cmd_p10_prompt():
    my_path = os.path.abspath(sys.argv[0])
    if "'" in my_path:
        raise Exception(f"invalid path to script: may not contain single quotes: {my_path}")
    print('''
        function prompt_branch_description() {
          local branch_description="$('%s' get)"
          [ -n "$branch_description" ] && p10k segment -f green -t "$branch_description"
        }
         ''' % my_path)

def cmd_git_pager():
    # If we can't get the repo info, just print stdout straight
    try:
        info = RepoInfo()
    except:
        for line in sys.stdin.readlines():
            print(line, end='') # line already contains newline
        return

    # okay, we have the repo info
    branch_descriptions = get_branch_descriptions(info)

    for line in sys.stdin.readlines():
        line = line.rstrip('\n\r')
        line_no_color = re.sub(r'\x1B\[[0-9;]*[JKmsu]', '', line)
        line_has_color = line != line_no_color

        branch_match = re.match('(\*)? *(\S+)$', line_no_color)
        if branch_match:
            def interleaving_print(s, file):
                '''Prints s to the file without a newline, and flushes s immediately.
                This lets us interleave writes to stdout and stderr on the same line.'''
                print(s, end='', file=file)
                file.flush()
            branch_active, branch_name = branch_match.groups()
            branch_description = branch_descriptions.get(branch_name, None)

            branch_sha = get_branch_sha(branch_name)

            interleaving_print('* ' if branch_active else '  ', sys.stdout)
            if branch_active:
                # paint over the "* " with a nicer version, in stderr
                interleaving_print('\r\033[33m▶︎ ', sys.stderr)
            if branch_sha:
                interleaving_print(f'({branch_sha}) ', sys.stderr)
            interleaving_print(branch_name, sys.stdout)
            if branch_active:
                interleaving_print('\033[0m', sys.stderr)
            if branch_description:
                color = branch_description.get('color', '96')
                desc = branch_description["description"]
                desc = _iterm2_links(desc)
                interleaving_print(f' \033[{color}m// {desc}\033[0m', file=sys.stderr)
            print('')
        else:
            print(line)


def cmd_set(description, branch, color, gh):
    info = RepoInfo(branch)

    if gh is not None and description is None:
        # gh is True or False; this is a CLI request
        if not gh:
            print('You must either provide a description, or use --gh', file=sys.stderr)
            exit(1)
        description = ''

    descriptions = configs.read_json(DESCRIPTIONS_JSON_FILE, default_value={})

    if info.repo_path not in descriptions:
        descriptions[info.repo_path] = {}
    if description is None:
        descriptions[info.repo_path].pop(info.current_branch, None)
    else:
        gh_url = _try_get_gh_url(info) if gh else None
        if gh_url:
            if description:
                if gh_url not in description:
                    description = f'{description} {gh_url}'
            else:
                description = gh_url

        descriptions[info.repo_path][info.current_branch] = {
            'description': description
        }
        if color:
            descriptions[info.repo_path][info.current_branch]['color'] = color

    clean_descriptions(descriptions)
    configs.write_json(DESCRIPTIONS_JSON_FILE, descriptions)

def cmd_unset(branch):
    cmd_set(None, branch, None, None)

def cmd_get(branch, raw):
    try:
        info = RepoInfo(branch)
    except:
        return

    branch_descriptions = get_branch_descriptions(info)
    branch_description = branch_descriptions.get(info.current_branch, {})
    desc = branch_description.get('description', '')
    color = branch_description.get('color')
    if not raw:
        desc = _iterm2_links(desc)
    if color:
        desc = f'\033[{color}m{desc}\033[0m'
    print(desc)

def get_branch_descriptions(repo_info=None):
    if repo_info is None:
        repo_info = RepoInfo()
    descriptions = configs.read_json(DESCRIPTIONS_JSON_FILE, default_value={})
    if clean_descriptions(descriptions):
        configs.write_json(DESCRIPTIONS_JSON_FILE, descriptions)
    return descriptions.get(repo_info.repo_path, {})

def clean_descriptions(descriptions):
    cleaned_any = False
    # when we iterate, we first call list() on the iterator to reify the keys. That way, it's safe to delete them from the dict
    # as we iterate
    for repo_path in list(descriptions.keys()): # defensive copy so we can delete in the iteration
        if not os.path.isdir(repo_path):
            descriptions.pop(repo_path)
            cleaned_any = True
        else:
            repo_descriptions = descriptions[repo_path]
            try:
                repo_branch_names = common.simple_exec("git for-each-ref --format=%(refname:short) refs/heads/".split(), cwd=repo_path)
            except:
                # If git gives an error, just ignore this dir. It could be that it's not a git dir anymore, but that seems
                # unlikely. More likely it's some config issue with git.
                continue
            repo_branch_names = set(repo_branch_names.split('\n'))
            for branch_name in list(repo_descriptions.keys()):
                if branch_name not in repo_branch_names:
                    repo_descriptions.pop(branch_name)
                    cleaned_any = True
            if not repo_descriptions:
                descriptions.pop(repo_path)
    return cleaned_any

def get_branch_sha(branch_name='HEAD', length=8):
    return common.simple_exec(['git', 'rev-parse', f'--short={length}', branch_name])


def _is_iterm2():
    return bool(os.environ.get('ITERM_SESSION_ID'))

def _iterm2_links(desc):
    if not _is_iterm2():
        return desc
    found_any_urls = False
    splits = desc.split()
    for idx, word in enumerate(splits):
        as_url = _try_parse_url(word)
        if as_url and as_url.scheme and as_url.path:
            url_label = as_url.path.split('/')[-1]
            icon = ''
            if re.search('github\\.com/\\w+/\\w+/pull/', as_url.geturl()):
                url_label = f'#{url_label}'
                icon = '\uea64'
            if re.search('tulipmfg\\.atlassian\\.net', as_url.geturl()):
                icon = '\ueaaf'
            if url_label:
                if icon:
                    icon = f'{icon} '
                word = f'{icon}\033]8;;{word}\033\\{url_label}\033]8;;\033\\' # nerdfont nf-cod-git_pull_request
                splits[idx] = word
                found_any_urls = True
    if found_any_urls:
        desc = ' '.join(splits)
    return desc

def _try_parse_url(s):
    from urllib.parse import urlparse
    try:
        return urlparse(s)
    except:
        return None

def _try_get_gh_url(info):
    branch_sha = info.get_sha()
    urls_json = common.simple_exec(['gh', 'pr', 'list', '--search', branch_sha, '--json', 'url'])
    urls_list = json.loads(urls_json)
    if not urls_list:
        print("Couldn't find a GitHub PR for current sha", file=sys.stderr)
    else:
        urls = [o['url'] for o in urls_list]
        if len(urls) > 1:
            print("Found too many GitHub PRs for current sha:", file=sys.stderr)
            for u in urls:
              print(f'- {u}', file=sys.stderr)
        else:
            return urls[0]
    return None

class RepoInfo:
    def __init__(self, branch_name=None):
        self.repo_path = common.simple_exec('git rev-parse --show-toplevel'.split())
        self.current_branch = branch_name or common.simple_exec('git rev-parse --abbrev-ref HEAD'.split())

    def get_sha(self):
        return common.simple_exec(['git', 'rev-parse', self.current_branch])

if __name__ == '__main__':
    main()