I want to write a decorator for some functions that take file as the first argument. The decorator has to implement the context manager protocol (i.e. turn the wrapped function into a context manager), so I figured I needed to wrap the function with a class.
I'm not really experienced with the decorator pattern and have never implemented a context manager before, but what I wrote works in Python 2.7 and it also works in Python 3.3 if I comment out the wraps line.
from functools import wraps
def _file_reader(func):
"""A decorator implementing the context manager protocol for functions
that read files."""
# @wraps(func)
class CManager:
def __init__(self, source, *args, **kwargs):
self.source = source
self.args = args
self.kwargs = kwargs
self.close = kwargs.get('close', True)
def __enter__(self):
# _file_obj is a little helper that opens the file for reading
self.fsource = _file_obj(self.source, 'r')
return func(self.fsource, *self.args, **self.kwargs)
def __exit__(self, exc_type, exc_value, traceback):
if self.close:
self.fsource.close()
return False
return CManager
The error I get when uncommenting the wraps line occurs inside update_wrapper:
/usr/lib/python3.3/functools.py in update_wrapper(wrapper, wrapped, assigned, updated)
54 setattr(wrapper, attr, value)
55 for attr in updated:
---> 56 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
57 # Return the wrapper so this can be used as a decorator via partial()
58 return wrapper
AttributeError: 'mappingproxy' object has no attribute 'update'
I know the docs don't say that I even can use functools.wraps to wrap a function with a class like this, but then again, it just works in Python 2. Can someone please explain what exactly this traceback is telling me and what I should do to achieve the effects of wraps on both versions of Python?
EDIT: I was mistaken. The code above does not do what I want it to. I want to be able to use the function both with and without with, like the builtin open.
The code above turns the decorated function into a context manager. I want to be able to do:
reader = func('source.txt', arg)
for item in reader:
pass
as well as
with func('source.txt', arg) as reader:
for item in reader:
pass
So my version of the code should probably look approximately as follows:
def _file_reader(func):
"""A decorator implementing the context manager protocol for functions
that read files."""
@wraps(func)
class CManager:
def __init__(self, source, *args, **kwargs):
self.close = kwargs.get('close', True)
self.fsource = _file_obj(source, 'r')
self.reader = func(self.fsource, *args, **kwargs)
def __enter__(self):
return self.reader
def __iter__(self):
return self.reader
def __next__(self):
return next(self.reader)
def __exit__(self, exc_type, exc_value, traceback):
if self.close and not self.fsource.closed:
self.fsource.close()
return False
return CManager
Feel free to comment about anything I have overlooked.
Note: the class version by J.F. Sebastian seems to work then:
I basically removed the wraps from the class and changed return CManager to:
@wraps(func)
def helper(*args, **kwargs):
return CManager(*args, **kwargs)
return helper
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
@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)
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()