November 16, 2011

Pretty command line / console output on Unix in Python and Go Lang

Posted in Software at 02:17 by graham

There are lots of easy ways to improve the output of your command line scripts, without going full curses, such as single-line output, using bold text and colors, and even measuring the screen width and height.

The examples are in Python, with a summary example in Go (golang) at the end.

Single line with \r (carriage return)

Instead of printing a \n (which most ‘print’ methods do by default), print a \r. That sends the cursor back to the beginning of the current line (carriage return), without dropping down to a new line (line feed).

import time, sys
total = 10
for i in range(total):
    sys.stdout.write('%d / %d\r' % (i, total))
    sys.stdout.flush()
    time.sleep(0.5)
print('Done     ')

We do sys.stdout.flush() to flush the output, otherwise the operating system buffers it and you don’t see anything until the end. Note also the extra spaces after ‘Done’, to clear previous output. Remove them and you’ll see why they are needed.

ANSI escape sequences

ANSI escape sequences are special non-printable characters you send to the terminal to control it. When you do ls --color, ls is using ANSI escape sequences to pretty things up for you.

When the terminal receives an ESC character (octal 33 – see man ascii) it expects control information to follow, instead of printable data. Most escape sequences start with ESC and [.

Bold

To display something in bold, switch on the bright attribute (1), then reset it afterwards (0).

def bold(msg):
    return u'\033[1m%s\033[0m' % msg

print('This is %s yes it is' % bold('bold'))

Color

There are 8 colors, ANSI codes 30 to 37, which can have the bold modifier, making 16 colors. Colors are surprisingly consistent across terminals, but can usually be changed by the user.

def color(this_color, string):
    return "\033[" + this_color + "m" + string + "\033[0m"

for i in range(30, 38):
    c = str(i)
    print('This is %s' % color(c, 'color ' + c))

    c = '1;' + str(i)
    print('This is %s' % color(c, 'color ' + c))

If that’s not enough for you, you can have 256 colors (thanks lucentbeing), but only if your terminal is set to 256 color mode. You usually get 256 colors by putting one of these lines in your .bashrc:

export TERM='screen-256color'   # Use this is you use tmux or screen. Top choice!
export TERM='xterm-256color'    # Use this otherwise

Here’s lots of colors (copy color function from above):

for i in range(256):
    c = '38;05;%d' % i
    print( color(c, 'color ' + c) )

Clear the screen

A pair of ANSI escape sequences:

import sys

def clear():
    """Clear screen, return cursor to top left"""
    sys.stdout.write('\033[2J')
    sys.stdout.write('\033[H')
    sys.stdout.flush()

clear()

More ANSI power

You can do lots more with ANSI escape sequences, such as change the background color, move the cursor around, and even make it blink. Use your power wisely.

Measure the screen

You can find out how many characters wide or lines high the current terminal is. From the command line, try tput lines or tput cols. To get the data direct from Unix:

import sys
import fcntl
import termios
import struct

lines, cols = struct.unpack('hh',  fcntl.ioctl(sys.stdout, termios.TIOCGWINSZ, '1234'))
print('Terminal is %d lines high by %d chars wide' % (lines, cols))

This is making a system call, hence it’s a bit inscrutable. It’s requesting TIOCGWINSZ information, for file sys.stdout (the terminal), and then interprets it as two packed short integers (hh). The ‘1234’ is just placeholder, which must be four chars long.

A pretty progress bar

Putting some of the above together, here’s a progress bar.

import sys
import fcntl
import termios
import struct
import time

COLS = struct.unpack('hh',  fcntl.ioctl(sys.stdout, termios.TIOCGWINSZ, '1234'))[1]

def bold(msg):
    return u'\033[1m%s\033[0m' % msg

def progress(current, total):
    prefix = '%d / %d' % (current, total)
    bar_start = ' ['
    bar_end = '] '

    bar_size = COLS - len(prefix + bar_start + bar_end)
    amount = int(current / (total / float(bar_size)))
    remain = bar_size - amount

    bar = 'X' * amount + ' ' * remain
    return bold(prefix) + bar_start + bar + bar_end

NUM = 100
for i in range(NUM + 1):
    sys.stdout.write(progress(i, NUM) + '\r')
    sys.stdout.flush()
    time.sleep(0.05)
print('\n')

That’s it! Happy console scripts.

And now the progress bar in Go golang, for fun

package main

import (
    "time"
    "os"
    "strconv"
    "strings"
    "syscall"
    "unsafe"
)

const (
    ONE_MSEC    = 1000 * 1000
    _TIOCGWINSZ = 0x5413    // On OSX use 1074295912. Thanks zeebo
    NUM         = 50
)

func main() {

    var bar string

    cols := TerminalWidth()

    for i := 1; i <= NUM; i++ {
        bar = progress(i, NUM, cols)
        os.Stdout.Write([]byte(bar + "\r"))
        os.Stdout.Sync()
        time.Sleep(ONE_MSEC * 50)
    }
    os.Stdout.Write([]byte("\n"))
}

func Bold(str string) string {
    return "\033[1m" + str + "\033[0m"
}

func TerminalWidth() int {
    sizeobj, _ := GetWinsize()
    return int(sizeobj.Col)
}

func progress(current, total, cols int) string {
    prefix := strconv.Itoa(current) + " / " + strconv.Itoa(total)
    bar_start := " ["
    bar_end := "] "

    bar_size := cols - len(prefix + bar_start + bar_end)
    amount := int(float32(current) / (float32(total) / float32(bar_size)))
    remain := bar_size - amount

    bar := strings.Repeat("X", amount) + strings.Repeat(" ", remain)
    return Bold(prefix) + bar_start + bar + bar_end
}

type winsize struct {
    Row    uint16
    Col    uint16
    Xpixel uint16
    Ypixel uint16
}

func GetWinsize() (*winsize, os.Error) {
    ws := new(winsize)

    r1, _, errno := syscall.Syscall(syscall.SYS_IOCTL,
        uintptr(syscall.Stdin),
        uintptr(_TIOCGWINSZ),
        uintptr(unsafe.Pointer(ws)),
    )

    if int(r1) == -1 {
        return nil, os.NewSyscallError("GetWinsize", int(errno))
    }
    return ws, nil
}

4 Comments »

  1. graham said,

    March 13, 2012 at 06:35

    @David Campbell Looks like exp/terminal is new, it’s in the weekly but not release versions. That would of helped me a bunch with hatcog too, thanks!

  2. David Campbell said,

    March 11, 2012 at 06:43

    The syscall package has a const syscall.TIOCGWINSZ. Also, you may want to check out the exp/terminal GetSize function.

  3. Graham King said,

    November 25, 2011 at 18:23

    @zeebo – thanks! I’ve updated the post.

    Darn it looks like WordPress stripped your includes, presumably because it thought they looked like html tags.

  4. zeebo said,

    November 16, 2011 at 21:48

    After trying your Go program on OS X, it seems like _TIOCGWINSZ was defined incorrectly for it. I ran this program to find the right constant:

        #include
        #include
        int main() {
                printf("%ld\n", TIOCGWINSZ);
                return 0;
        }
    

    Which printed 1074295912, and placing that value for _TIOCGWINSZ worked perfectly.

    Thanks for the great post! :)

Leave a Comment

Note: Your comment will only appear on the site once I approve it manually. This can take a day or two. Thanks for taking the time to comment.