functools.wraps()
is for wrapper functions:
import contextlib
import functools
def file_reader(func):
@functools.wraps(func)
@contextlib.contextmanager
def wrapper(file, *args, **kwargs):
close = kwargs.pop('close', True) # remove `close` argument if present
f = open(file)
try:
yield func(f, *args, **kwargs)
finally:
if close:
f.close()
return wrapper
Example
@file_reader
def f(file):
print(repr(file.read(10)))
return file
with f('prog.py') as file:
print(repr(file.read(10)))
If you want to use a class-based context manager then a workaround is:
def file_reader(func):
@functools.wraps(func)
def helper(*args, **kwds):
return File(func, *args, **kwds)
return helper
To make it behave identically whether the decorated function is used directly or as a context manager you should return self
in __enter__()
:
import sys
class File(object):
def __init__(self, file, func, *args, **kwargs):
self.close_file = kwargs.pop('close', True)
# accept either filename or file-like object
self.file = file if hasattr(file, 'read') else open(file)
try:
# func is responsible for self.file if it doesn't return it
self.file = func(self.file, *args, **kwargs)
except: # clean up on any error
self.__exit__(*sys.exc_info())
raise
# context manager support
def __enter__(self):
return self
def __exit__(self, *args, **kwargs):
if not self.close_file:
return # do nothing
# clean up
exit = getattr(self.file, '__exit__', None)
if exit is not None:
return exit(*args, **kwargs)
else:
exit = getattr(self.file, 'close', None)
if exit is not None:
exit()
# iterator support
def __iter__(self):
return self
def __next__(self):
return next(self.file)
next = __next__ # Python 2 support
# delegate everything else to file object
def __getattr__(self, attr):
return getattr(self.file, attr)
Example
file = f('prog.py') # use as ordinary function
print(repr(file.read(20)))
file.seek(0)
for line in file:
print(repr(line))
break
file.close()