CODEWITHSUNDEEP
pythonzshansi-escape
Jared
Jared

Reputation: 4810

How can I prevent ANSI escape sequences in my Python script from messing with my zsh RPROMPT and cursor positions?

I've been working on a Python script that generates an RPROMPT for zsh about the state of a git repo in the current working directory. It's invoked in my .zshrc file with:

RPROMPT='$(python3 ~/.git_zsh_rprompt.py)'

To differentiate it from the other text in my terminal, I've used ANSI escape codes to make it bold and color it based on whether the repo is clean or dirty. For whatever reason, adding these escape codes has caused my RPROMPT to scoot over to the left, and it has also moved my cursor over on top of my left PROMPT. Here's a representation, with the block representing my cursor:

jared@Jareds-Mac⌷ook-Pro:foobar%                    master( +2 ~4 )

Aside from a few non-educated guesses, I'm not entirely sure why this is happening. I'm hoping someone here does and is aware of a solution or workaround that's puts everything back where it's supposed to be. For reference, here is the offending script:

from collections import Counter
from os import devnull
from subprocess import call, check_output, STDOUT

def git_is_repo():
    command = ["git", "branch"]
    return not call(command, stderr = STDOUT, stdout = open(devnull, "w"))

def git_current_branch():
    command = ["git", "branch", "--list"]
    lines = check_output(command).decode("utf-8").strip().splitlines()
    return next((line.split()[1] for line in lines if line.startswith("* ")),
        "unknown branch")

def git_status_counter():
    command = ["git", "status", "--porcelain"]
    lines = check_output(command).decode("utf-8").strip().splitlines()
    return Counter(line.split()[0] for line in lines)

if __name__ == "__main__":
    if git_is_repo():
        counter = git_status_counter()

        # Print bold green if the repo is clean or bold red if it is dirty.
        if counter.elements() == []:
            print("\033[1;32m", end="")
        else:
            print("\033[1;31m", end="")

        print(git_current_branch(), end="")

        # Only print status counters if the repo is dirty.
        if counter.elements() != []:
            print("(", end="")

            if counter["??"] != 0:
                print(" +{}".format(counter["??"]), end="")
            if counter["M"] != 0:
                print(" ~{}".format(counter["M"]), end="")
            if counter["D"] != 0:
                print(" -{}".format(counter["D"]), end="")

            print(" )", end="")

        # Reset text attributes.
        print("\033[0m", end="")

Upvotes: 1

Views: 653

Answers (2)

Eric
Eric

Reputation: 97631

You want to use the %{ %} ZSH escape sequence around your ANSI escapes in the prompt, as documented here

%{...%}

Include a string as a literal escape sequence. The string within the braces should not change the cursor position. Brace pairs can nest.

A positive numeric argument between the % and the { is treated as described for %G below.

Upvotes: 1

abarnert
abarnert

Reputation: 365915

This is the usual problem with zsh, like all other shells, not knowing how long your prompt is except by strlening the string, unless you tag the non-printing characters for it in some way.

There are ways to work around this for every shell,* but in zsh, you really don't have to; that's half the reason it has visual effects commands in its prompt format in the first place.**

So, for example, instead of this:

print("\033[1;32m", end="")

… you can do this:

print("%B%F{green}", end="")

zsh knows how to look up bold and green and insert the appropriate escape sequences for your terminal, and it knows now to count those sequences against your prompt length.

* IIRC, zsh supports the bash-style escaped square brackets, and its own %{…%} tagging.

** The other half is that the visual effects will termcaps to the right escape sequences, or degrade to nothing if that's not possible, instead of just assuming everything is ANSI.

Upvotes: 0

Related Questions