.iterm2_shell_integration.xsh

import base64

try:
    from prompt_toolkit import print_formatted_text
    from prompt_toolkit.formatted_text import FormattedText
except:
    pass

def iterm2_status(status):
    return "\033]133;D;{}\007".format(status)

# Mark start of prompt
ITERM2_PROMPT_MARK = "\033]133;A\007"

# Mark end of prompt
ITERM2_PROMPT_END = "\033]133;B\007"

# Tell terminal to create a mark at this location
ITERM2_PREEXEC = "\033]133;C;\007"

# Usage: iterm2_set_user_var key value
# These variables show up in badges (and later in other places). For example
# iterm2_set_user_var currentDirectory "$PWD"
# Gives a variable accessible in a badge by \(user.currentDirectory)
# Calls to this go in iterm2_print_user_vars.
def iterm2_set_user_var(key, value):
    # Replace any \n with spaces to avoid breaking the shell.
    key = key.replace("\n", " ")
    return "\033]1337;SetUserVar={}={}\007" % (key, base64.b64encode(value.encode('utf-8')))

# Inform iTerm2 that command starts here
@events.on_precommand
def iterm2_precmd(cmd=None):
    print(ITERM2_PREEXEC, end = '')

# Inform iTerm2 of command success/failure here
@events.on_postcommand
def iterm2_postcmd(cmd, rtn, out, ts):
    print(iterm2_status(rtn), end = '')

PRE_INSERT = ""

@events.on_pre_prompt
def preprompt():
    global PRE_INSERT
    pre_insert_arr = []

    $PROMPT = $PROMPT.replace(PRE_INSERT, "")

    # Inform iTerm2 that a prompt starts here. This will be used to insert the
    # prompt mark.
    pre_insert_arr.append(ITERM2_PROMPT_MARK)

    # Inform iTerm2 where a prompt ends. I believe iTerm2 uses this to figure
    # out what the commandline being run is.
    if $SHELL_TYPE == "prompt_toolkit":
        if not ITERM2_PROMPT_END in $PROMPT:
            $PROMPT += ITERM2_PROMPT_END
    elif $SHELL_TYPE == "readline":
        if not ITERM2_PROMPT_END in $PROMPT:
            $PROMPT += "\001" + ITERM2_PROMPT_END + "\002"
    else:
        # Error
        pass

    pre_insert_arr.append("\033]1337;RemoteHost={}@{}\007".format($USER, iterm2_hostname))
    pre_insert_arr.append("\033]1337;CurrentDir={}\007".format($PWD))

    # TODO: Users can define a function called iterm2_print_user_vars.
    # It should call iterm2_set_user_var and produce no other output.
    #if functions -q -- iterm2_print_user_vars
    #  iterm2_print_user_vars
    #end

    # Set the prompt.
    if $SHELL_TYPE == "prompt_toolkit":
        PRE_INSERT = "".join(pre_insert_arr)
        $PROMPT = PRE_INSERT + $PROMPT
    elif $SHELL_TYPE == "readline":
        PRE_INSERT = "\001" + "".join(pre_insert_arr) + "\002"
        $PROMPT = PRE_INSERT + $PROMPT
    else:
        pass

import socket
iterm2_hostname = socket.gethostname()

# Remove starting/ending whitespace, and any potential \n
iterm2_hostname = iterm2_hostname.strip().replace("\n", " ")

# Show dummy prompt
iterm2_precmd()
print("\033]1337;ShellIntegrationVersion=5;shell=xonsh\007", end='')