11

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
Lev Levitsky
  • 63,701
  • 20
  • 147
  • 175

2 Answers2

5

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()
jfs
  • 399,953
  • 195
  • 994
  • 1,670
  • Is this equivalent to a full-blown custom context manager class? I looked at `contextlib.contextmanager` but I thought it just patched some stubs onto the function and wasn't really useful. Note that I want the function to work without `with`, too, pretty much like `open`. – Lev Levitsky Dec 30 '12 at 23:26
  • @LevLevitsky: yes, it implements full protocol. The difference (`close` is popped) from your class is intentional to make it more `open`-like. – jfs Dec 30 '12 at 23:44
  • I have a little trouble understanding what you meant by `yield func(f, *args, **kwargs)`. It works in a `with` statement, but gives `TypeError: GeneratorContextManager object is not an iterator` when I use it normally (try to access the iterator it returns). – Lev Levitsky Dec 31 '12 at 09:31
  • @LevLevitsky: file_reader() is a context manager, it is not an iterator. Read contextlib.contextmanager docs. – jfs Dec 31 '12 at 10:01
  • I mean that the function I decorate returns an iterator (like `open`), and I want it to work both ways, like `open` does. But yeah, I see that this construct is mentioned in the docs: `This iterator must yield exactly one value, which will be bound to the targets in the with statement’s as clause, if any.` – Lev Levitsky Dec 31 '12 at 10:05
  • Damn, turns out my code doesn't do it either.. I don't know why I thought it did. That must have confused everyone, I'm sorry. I might have to change the question quite a bit now, but I need to think first. – Lev Levitsky Dec 31 '12 at 12:15
  • Seems like it's not that bad after all. I just had to add a couple of methods to my class. Your class-based solution works fine, thanks. Any comments about the class are most welcome. – Lev Levitsky Dec 31 '12 at 13:18
  • @LevLevitsky: I've added code to allow access without `with`-statement – jfs Jan 03 '13 at 12:38
  • Thanks! It's very useful, but it seems like I wasn't very clear about the kind of iteration `func` does. It doesn't return the file object; rather, it yields data it reads from the file (it's a parser). Accordingly, I don't think I should bind it's return value to `self.file`, rather create some other attribute. Right? The rest is very much to the point, I think, and very well thought-out, thanks again. – Lev Levitsky Jan 04 '13 at 11:08
  • May I ask what you meant by `**kwargs` in `__exit__()` definition? Can it be called with anything other than `sys.exc_info()`? – Lev Levitsky Jan 05 '13 at 21:41
  • @LevLevitsky: I've used `f(*args, **kwargs)` syntax to stress that the call is delegated. `kwargs` is not used by the `with`-statement (though it might be useful to support extensions for the context manager protocol). In addition to `f(*sys.exc_info())` it can be called as `f(None, None, None)`. – jfs Jan 06 '13 at 02:55
  • I wonder if I can somehow call `__exit__` when the iterator is exhausted in a regular use case (without `with`), or create some sort of a hook for this event. I know it kind of decreases (to about zero) the need for supporting `with`, but still... or is this idea too perverted? – Lev Levitsky Jan 06 '13 at 22:09
  • @LevLevitsky: just call `.close()` method after you are done with the object. Putting the call inside `__next__()` (to be called on StopIteration) seems like a bad idea: the outside code should either use `with`-statement or call `.close()` explicitly otherwise there might be a resource leak in case of exceptions – jfs Jan 06 '13 at 22:27
  • So you advise to drop the `close` kwarg altogether then and tell the users to switch to `with`? The things is, `close` is documented in the current version and for now it is processed in the end of each `func`. P.S. Sorry for the ridiculous number of comments... – Lev Levitsky Jan 06 '13 at 22:54
  • @LevLevitsky: I meant that the iterator might never be exhausted in case of exceptions (it doesn't prevent the outside code to call `.close()` explicitly e.g., in `finally` statement). In short, the interface should be identical to `open()` except instead of lines some other tokens are generated. Whether to keep `.close_file` attribute depends on whether there is additional semantics e.g., see `delete` parameter in `tempfile.TemporaryFile()`. `close_file=False` might be useful if `File` does some cleanup on `.close()` but user wants the underlying file to stay open. – jfs Jan 06 '13 at 23:42
  • An unfortunate side effect of this solution is that wrapping with functions instead of classes doesn't let you inherit the classes if necessary, though, so it masks classmethods and staticmethods away for instance :( – monokrome Jun 08 '19 at 21:06
4

Although I don't know what the error you're seeing is from, it looks like you are probably doomed anyway:

>>> import functools
>>> def foo():
...     pass
... 
>>> class bar:
...     pass
... 
>>> functools.wraps(foo)(bar)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib64/python3.2/functools.py", line 48, in update_wrapper
    setattr(wrapper, attr, value)
AttributeError: attribute '__doc__' of 'type' objects is not writable
>>> bar.__doc__
>>> bar.__doc__ = 'Yay'
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: attribute '__doc__' of 'type' objects is not writable

Edit: J.F. Sebastian has the solution, but i'm leaving this bit as exposition on why it must be a function and not a class

Community
  • 1
  • 1
SingleNegationElimination
  • 151,563
  • 33
  • 264
  • 304
  • Thanks for the observation and the workaround, I'll try it out tomorrow (or today, depending on your timezone). What I'd be interested in is the underlying difference between the Python versions. – Lev Levitsky Dec 30 '12 at 23:28