FileChunker.py

#!/usr/bin/env python
# encoding: utf-8
"""
FileChunker.py

Created by Ryan Balfanz on 2009-11-19.
Copyright (c) 2009 Ryan Balfanz. All rights reserved.
"""

import fileinput
import sys
import cStringIO as StringIO

def gfiles():
	"""Generate files from fileinput.input().
	
	Extends the functionality of fileinput.input() to generate complete files.
	Note that line endings are preserved.
	
	Example usage:
		>>> for contents in gfiles():
		>>> 	f = StringIO.StringIO(contents)
		>>> 	for line in f:
		>>> 		sys.stdout.write(line)
	"""
	curFileBuffer = None
	for line in fileinput.input(mode="rb"):
		if fileinput.isfirstline():
			if not curFileBuffer:
				# This is the first line of the first file.
				curFileBuffer = StringIO.StringIO()
				curFileBuffer.write(line)
			else:
				# This is a new file that is not the first file, return the previous file buffer.
				contents = curFileBuffer.getvalue()
				curFileBuffer.close()
				curFileBuffer = None
				# We cannot include the filename becuase of how fileinput.filename() works.
				yield contents
				curFileBuffer = StringIO.StringIO()
				curFileBuffer.write(line)
				# fileinput.nextfile() # This breaks stuff, but I feel like it belongs.
		else:
			# Just some line in the current file, write line to buffer.
			curFileBuffer.write(line)
	contents = curFileBuffer.getvalue()
	curFileBuffer.close()
	curFileBuffer = None
	fileinput.close() # Is this really needed?
	yield contents

if __name__ == '__main__':	
	for n, contents in enumerate(gfiles()):
		f = StringIO.StringIO(contents)
		for i, line in enumerate(f):
			print n, i, line,
		print
		print "*** EOF ***"