Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Wrapping a file function
class wrap_file_function(object):
"""
Wrap a function which takes a file or a str as it's first argument.
If a str is provided, replace the first argument of the wrapped function
with a file handle, and close the file afterwards
Example:
@wrap_file_function('w')
def write_hi(f):
f.write('hi!\n')
# This will write to already open file handle.
f = open('f1.txt', 'w')
write_hi(f)
f.close()
# This will open file f2.txt with mode 'w', write to it, and close the file.
write_hi('f2.txt')
"""
def __init__(self, *args):
self.modes = args if args else ('r',)
def __call__(self, func):
def wrapped(*args, **kwargs):
close = [] # Files that should be closed
files = [] # File handles that should be passed to func
num_files = len(self.modes)
try:
for i, mode in enumerate(self.modes):
fp = args[i]
if isinstance(fp, str):
fp = open(fp, mode)
close.append(fp)
files.append(fp)
# Replace the files in args when calling func
args = files + list(args[num_files:])
# Make function call and return value
return func(*args, **kwargs)
finally:
for fp in close:
fp.close()
return wrapped
if __name__ == "__main__":
# Demonstration of wrapping a function which writes to a file.
print '-'*50
@wrap_file_function('w')
def write_hi(f):
f.write('hi!\n')
f = open('temp.txt', 'w')
write_hi(f)
write_hi(f)
write_hi(f)
f.close()
write_hi('temp2.txt')
# Demonstration of wrapping a function which reads from a file.
print '-'*50
@wrap_file_function()
def read_file(f):
print f.read()
f = open('temp.txt')
print 'Reading file temp.txt from handle f:'
read_file(f)
print 'Reading file temp2.txt'
read_file('temp2.txt')
# Demonstration of wrapping a function takes multiple files
print '-'*50
@wrap_file_function('r', 'r')
def read_files(f1, f2):
print 'reading f1: '
print f1.read()
print 'reading f2: '
print f2.read()
@wrap_file_function('r', 'w')
def read_write(f1, f2):
f2.write(f1.read())
@wrap_file_function('w', 'r')
def write_to_from(f2, f1):
f2.write(f1.read())
read_files(open('temp.txt'), open('temp2.txt'))
read_files('temp.txt', 'temp2.txt')
read_write('temp.txt', 'temp.copy.txt')
import sys
print 'writing temp.txt to stdout:'
read_write('temp.txt', sys.stdout)
print 'Contents of temp.copy.txt:'
read_file('temp.copy.txt')
@wrap_file_function('w')
def throw_exception(f):
raise RuntimeError('BLAH!')
print '-'*50
print 'Reading a file that does not exist:'
write_to_from('dest.txt', 'file_doesnt_exist.txt')
#print '-'*50
#throw_exception('exception.txt')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment