#!/usr/bin/python
# General utilities for Python 2 and 3 programs.
# This is intended to be imported with "from offog import *".
# Adam Sampson <ats@offog.org>

import sys, os, stat, subprocess, fcntl

def warn(*args):
	sys.stderr.write("".join(map(str, args)) + "\n")

def die(*args):
	warn(*args)
	sys.exit(1)

def run_command(*args):
	"""Run a command; die if it exits non-zero."""
	rc = subprocess.call(*args)
	if rc != 0:
		die("Command failed with returncode ", rc, ": ", args)

def ensure_dir(dir):
	"""Like os.makedirs, but don't fail if the directory already exists."""
	try:
		st = os.lstat(dir)
		if stat.S_ISDIR(st.st_mode):
			return
	except OSError:
		pass
	os.makedirs(dir)

class memoised:
	"""Decorator to memoise a function."""
	def __init__(self, function):
		self.memo = {}
		self.function = function
	def __call__(self, *args, **kwargs):
		key = (args, tuple(kwargs.items()))
		if key in self.memo:
			return self.memo[key]
		r = apply(self.function, args, kwargs)
		self.memo[key] = r
		return r

class traced:
	"""Decorator to print the arguments and result for each call to a
	function."""
	def __init__(self, function):
		self.function = function
	def __call__(self, *args, **kwargs):
		arglist = map(repr, args)
		arglist += ["%s=%s" % (n, repr(v))
		            for n, v in sorted(kwargs.items())]
		r = apply(self.function, args, kwargs)
		sys.stderr.write("%s(%s) = %s\n"
		                 % (self.function.__name__,
		                    ", ".join(arglist),
		                    repr(r)))
		return r

class push_dir:
	"""Context guard to change directory."""
	def __init__(self, dir):
		self.dir = dir
	def __enter__(self):
		self.previous = os.getcwd()
		os.chdir(self.dir)
	def __exit__(self, type, value, traceback):
		os.chdir(self.previous)

class temporary_dir:
	"""Context guard to create a temporary directory.
	The directory will be forcibly removed at exit."""
	def __init__(self, dir):
		self.dir = dir
	def __enter__(self):
		ensure_dir(self.dir)
	def __exit__(self, type, value, traceback):
		subprocess.check_call(["rm", "-fr", self.dir])

class atomically_updated:
	"""Context guard to atomically update a file,
	by writing to a temporary name (in the same directory), and then
	renaming over the name specified."""
	def __init__(self, filename, *args):
		self.filename = filename
		self.args = args
	def __enter__(self):
		self.temp_filename = self.filename + (".new-%d" % os.getpid())
		self.f = open(self.temp_filename, *self.args)
		return self.f
	def __exit__(self, type, value, traceback):
		self.f.close()
		if type is None:
			# No exception -- update.
			os.rename(self.temp_filename, self.filename)
		else:
			# Something went wrong.
			os.unlink(self.temp_filename)

class file_locked:
	"""Context guard to claim an fcntl lock."""
	def __init__(self, filename):
		self.filename = filename
	def __enter__(self):
		self.fd = os.open(self.filename, os.O_WRONLY | os.O_CREAT, 0o600)
		fcntl.lockf(self.fd, fcntl.LOCK_EX)
		return self.fd
	def __exit__(self, type, value, traceback):
		os.close(self.fd)