#!/usr/bin/python
# Extract archives into directories with sensible names.
# Adam Sampson <ats@offog.org>
#
# The arguments can be files or directories; these will be searched for archive
# files that can be extracted.
#
# foo.zip will be extracted into a directory called foo. If extracting the
# archive produces a single directory, the contents of that directory will be
# moved into foo and the redundant directory removed.

import optparse
import os
import re
import stat
import subprocess
import sys

# A list of supported archive extensions.
# (extension, command to extract into current directory)
archive_types = [
	(".zip", ["unzip"]),
	(".rar", ["unrar", "x"]),
	(".7z", ["7z", "x"]),
	(".tar", ["tar", "xf"]),
	(".tar.bz2", ["tar", "xjf"]),
	(".tar.gz", ["tar", "xzf"]),
	(".tar.xz", ["tar", "xJf"]),
	(".tar.Z", ["tar", "xZf"]),
	]

def is_dir(path):
	return stat.S_ISDIR(os.stat(path).st_mode)

def extract(path, ext, cmd):
	"""Extract an archive.
	Return True if extraction worked, False otherwise."""

	parent = os.path.dirname(path)

	# Chop the extension off the basename.
	destname = os.path.basename(path)[:-len(ext)]
	dest = os.path.join(parent, destname)

	here = os.getcwd()

	if os.access(dest, os.F_OK):
		# Already exists -- assume it's been extracted already.
		sys.stderr.write("Already extracted: " + dest + "\n")
		return True

	# Extract the contents of the archive.
	os.mkdir(dest)
	os.chdir(dest)
	rc = subprocess.call(cmd + [os.path.relpath(path, dest)])

	if rc != 0:
		# Extraction failed -- get rid of the partial result.
		os.chdir(here)
		subprocess.check_call(["rm", "-rf", dest])
		return False

	# Check to see if the contents is all in a single directory.
	entries = os.listdir(".")
	if len(entries) == 1 and is_dir(entries[0]):
		dirname = entries[0]

		# First make sure that the directory name isn't the same as one
		# of the things inside it (e.g. foo/foo).
		subentries = os.listdir(dirname)
		while dirname in subentries:
			dirname += "_"
		if dirname != entries[0]:
			os.rename(entries[0], dirname)

		# Move everything up a level.
		for entry in subentries:
			os.rename(os.path.join(dirname, entry), entry)
		os.rmdir(dirname)

	os.chdir(here)

	return True

def scan(path, failed):
	if is_dir(path):
		# Recurse on the contents.
		for entry in sorted(os.listdir(path)):
			scan(os.path.join(path, entry), failed)
		return

	for (ext, cmd) in archive_types:
		if path.endswith(ext):
			if not extract(path, ext, cmd):
				failed.append(path)
			break

def main():
	parser = optparse.OptionParser()
	(options, args) = parser.parse_args()

	failed = []
	for path in args:
		scan(path, failed)

	if failed != []:
		for path in failed:
			sys.stderr.write("Extraction failed: " + path + "\n")
		sys.exit(1)

if __name__ == "__main__":
	main()