Using Python Queue with a "with" statement

Question

Is there a standard way to use a Python Queue in a with statement? This is how I would like to be able to use it:

import Queue
myqueue = Queue.Queue()
...
...
...
with myqueue as nextItem:
    doStuff(nextItem)

There are no __enter__ or __exit__ methods for Queue objects, so this does not work. Is there any syntactic sugar to make it look better than this?

import Queue
myqueue = Queue.Queue()
...
...
...
try:
    nextItem = myqueue.get()
    doStuff(nextItem)
finally:
    myqueue.task_done()

Edit: I have two reasons for wanting to use a with statement in this case. First, I think that a with statement would help my code be a little bit cleaner, especially when there is more than just the single call to doStuff. Second, it would be nice to have something simple that I could get in the habit of using every time that would ensure that I don't ever forget to call task_done or have a situation in which an error could cause the call to be skipped.

Answer 1

There are a couple of simple options that will ensure task_done() gets called on your queue no matter what. These are especially useful if the work you need to do with an item from the queue is complicated, or there are multiple code paths for being done with each item.

The first option is specific to queues. It's similar to what Rob suggested, but it doesn't require making a subclass of Queue.

@contextlib.contextmanager
def mark_done(queue):
    try:
        yield queue.get()
    finally:
        queue.task_done()

queue = Queue()
# assume you put stuff in the queue
with mark_done(queue) as item:
    item.do_stuff()

I like a more general form that's useful any time you want to call a function at the end of a code block no matter what, and that function is related to an object at the top of the block:

@contextlib.contextmanager
def defer(func):
    try:
        yield
    finally:
        func()

queue = Queue()
# assume you put stuff in the queue
item = queue.get()
with defer(queue.task_done):
    item.do_stuff()

Using this second form requires one extra line of code, but considering readability, it makes it very clear what the context manager is going to do.

Answer 2

It seems like the answer is no - there is no built-in way to do this.

However as sweeneyrod mentioned, it is possible to subclass Queue and add __enter__ and __exit__ methods. That would look like this:

import Queue

class MyQueue(Queue.Queue):
    def __enter__(self):
        return self.get()

    def __exit__(self, excType, excValue, traceback):
        self.task_done()

Which would allow for it to be used like I showed above, though this has the problem that it treats the queue as if it were a task. To fix that, we can use contextlib to create a method that acts as a context manager.

import Queue
from contextlib import contextmanager

class MyQueue(Queue.Queue):
    @contextmanager
    def task(self):
        try:
            yield self.get()
        finally:
            self.task_done()

You would use this version like this:

with myqueue.task() as next_task:
    doStuff(next_task)

It is also possible to have the task and __exit__ methods do some exception handling, though there are some differences between the two in how that is done.