annotate.py: like 'annotate-output' from ubuntu devscripts but in python

README.md

timestamp-output

Run a command and add timestamps to its standard in and standard out.

This annotates output the same way that annotate-output from the ubuntu devscripts does. The benefit of this over that is primarily in that it does not create a subprocess for every line read (annotate-output users date to get the timestamp).

The string _S_ is replaced in the format with 'O' or 'E' for stdout and stderr respectively.

Examples

• Simple

   $ timestamp-output -- echo hi world
   12:36:47 O: hi world
  

• Sub-second timestamps

   $ timestamp-output --output-prefix="%H:%M:%S.%f _S_: " -- \
       sh -c 'echo hello stdout; echo hello stderr 1>&2'
   12:29:00.169712 E: hello stderr
   12:29:00.169829 O: hello stdout
  

• Delta timestamps

  By providing a output format with '%s' in it, you can get the time in seconds since command start. The '%s' can be anything that is accepted as a float format for printfi ('%.3s' to give 3 decimal places).

   $ timestamp-output --output-prefix="%M:%S.%f %s _S_: " -- \
       sh -c 'echo hello stdout; sleep 1; echo hello stderr 1>&2'
   39:38.121876 0.002955 O: hello stdout
   39:39.123261 1.004340 E: hello stderr
  

timestamp-output

#!/usr/bin/python3
"""timestamp-output

This annotates output the same way that annotate-output from
the ubuntu devscripts does.  The benefit of this over that is
primarily in that it does not create a subprocess for every
line read (annotate-output users `date` to get the timestamp)"""

import argparse
import datetime
import os
import re
import subprocess
import sys
import threading


START_TIME = datetime.datetime.now()
DELTA_FMT_RE = re.compile('(?P<tok>%(s|-?\d*.?\d+s))')

FD_STDOUT = 1
FD_STDERR = 2


def fp_readlines(fp, sizehint=None):
    if sizehint is None:
        sizehint = 1024
    leftover = bytes()
    while True:
        buf = os.read(fp, sizehint)
        consumed = buf == b''

        if leftover:
            buf = leftover + buf
            leftover = bytes()

        while True:
            end = buf.find(b'\n')
            if end != -1:
                yield buf[0:end+1]
                buf = buf[end+1:]
            else:
                leftover = buf
                break

        if consumed:
            break


def ftime(fmt, ts=None):
    if ts is None:
        ts = datetime.datetime.now()
    return datetime.datetime.strftime(ts, fmt)


def uftime(fmt, ts=None, ref=None):
    if ref is None:
        ref = START_TIME
    if ts is None:
        ts = datetime.datetime.now()
    match = DELTA_FMT_RE.search(fmt)
    if match:
        fmtstr = match.group('tok')
        fmtfloat = fmtstr[:-1] + "f"
        fmt = fmt.replace(
            fmtstr, fmtfloat % (ts - ref).total_seconds(), 1)
    return ftime(fmt, ts)


def needs_delta(fmt):
    return bool(DELTA_FMT_RE.search(fmt))


def addtime(fh_in, fh_out, fmt='%H:%M:%S '):
    if needs_delta(fmt):
        fmtfunc = uftime
    else:
        fmtfunc = ftime

    for line in fp_readlines(fh_in):
        os.write(fh_out, fmtfunc(fmt).encode() + line)

    return


def main():
    fmt = "%H:%M:%S"

    parser = argparse.ArgumentParser()
    parser.add_argument('-o', '--stdout-prefix',
                        dest='out_fmt', default=fmt + " O: ",
                        help='the format prefix for stdout')
    parser.add_argument('-e', '--stderr-prefix',
                        dest='err_fmt', default=fmt + " E: ",
                        help='the format prefix for stderr')
    parser.add_argument('-O', '--output-prefix',
                        dest='all_fmt', default=None,
                        help='the format prefix for both out and err')
    parser.add_argument('cmd', nargs='+', default=[])

    args = parser.parse_args()

    if args.all_fmt:
        args.out_fmt = args.all_fmt.replace("_S_", "O")
        args.err_fmt = args.all_fmt.replace("_S_", "E")

    threads = []
    cmd = args.cmd

    out_r, out_w = os.pipe()
    err_r, err_w = os.pipe()

    try:
        t = threading.Thread(
            target=addtime, args=(out_r, FD_STDOUT, args.out_fmt))
        t.start()
        threads.append(t)

        t = threading.Thread(
            target=addtime, args=(err_r, FD_STDERR, args.err_fmt))
        t.start()
        threads.append(t)

        sp = subprocess.Popen(cmd, stdout=out_w, stderr=err_w)
        sp.communicate()
        rc = sp.returncode

        sys.exit(rc)

    finally:
        os.close(out_w)
        os.close(err_w)

        for t in threads:
            t.join()

        os.close(out_r)
        os.close(err_r)


main()