#!/usr/bin/env python

PAGERS = ["less", "more", "most", "tail", "head"]

# Usage:
# * Put this file into ~/bin/colorwrap
# * Add these aliases to your bash config
#     alias ls="colorwrap ls"
#     alias grep="colorwrap grep"


import os, os.path, sys, time

class Process(object):
	def __init__(self, pid):
		self.pid = pid
	
	@property
	def stdin(self):
		return os.readlink("/proc/%s/fd/0" % self.pid)
	
	@property
	def stdout(self):
		return os.readlink("/proc/%s/fd/1" % self.pid)
	
	@property
	def cmdline(self):
		return file("/proc/%s/cmdline" % self.pid).read().split("\x00")

def validProcesses():
	for pid in filter(str.isdigit, os.listdir("/proc/")):
		if os.path.exists("/proc/%s/fd/0" % pid):
			yield Process(pid)

def isatty(filename):
	try:
		stdout = file(cur.stdout)
		return os.isatty(stdout.fileno())
	except IOError:
		return False

# Follow from current pid until the end of the chain
count = 0
cur = Process(os.getpid())
while count < 32 and cur.stdout.startswith("pipe:"):
	last = cur
	for cur in validProcesses():
		try:
			if last.stdout == cur.stdin:
				count = 0
				break
		except OSError:
			continue
	else:
		cur = last
		count += 1
		time.sleep(0.001)

# Use colour or not
command = sys.argv[1]
args = sys.argv[1:]
if cur.cmdline in PAGERS or isatty(cur.stdout):
	args.insert(1, "--color=always")

os.execvp(command, args)