Make consumer_for a context manager

This patch converts consumer_for into a `contextmanager` decorated
method and adapts `__getattribute__` to support this change. A new
method was not added to avoid clashes with method names in the stages.

Change-Id: I6c0727e86d9ee8d6bd34c312d180040c8baf34ba
This commit is contained in:
Flavio Percoco 2013-11-29 21:30:49 +01:00
parent 11660c2fca
commit 8fa010db98
2 changed files with 13 additions and 3 deletions

View File

@ -30,6 +30,8 @@ At least one of the stages has to implement the calling method. If none of
them do, an AttributeError exception will be raised. them do, an AttributeError exception will be raised.
""" """
import contextlib
import six import six
from marconi.common import decorators from marconi.common import decorators
@ -48,8 +50,10 @@ class Pipeline(object):
@decorators.cached_getattr @decorators.cached_getattr
def __getattr__(self, name): def __getattr__(self, name):
return self.consumer_for(name) with self.consumer_for(name) as consumer:
return consumer
@contextlib.contextmanager
def consumer_for(self, method): def consumer_for(self, method):
"""Creates a closure for `method` """Creates a closure for `method`
@ -106,4 +110,4 @@ class Pipeline(object):
LOG.error(msg) LOG.error(msg)
raise AttributeError(msg) raise AttributeError(msg)
return consumer yield consumer

View File

@ -63,7 +63,7 @@ class TestPipeLine(base.TestBase):
SecondClass()]) SecondClass()])
def test_attribute_error(self): def test_attribute_error(self):
consumer = self.pipeline.consumer_for('does_not_exist') consumer = self.pipeline.does_not_exist
self.assertRaises(AttributeError, consumer) self.assertRaises(AttributeError, consumer)
def test_with_args(self): def test_with_args(self):
@ -87,3 +87,9 @@ class TestPipeLine(base.TestBase):
def test_calls_the_latest(self): def test_calls_the_latest(self):
self.assertTrue(self.pipeline.calls_the_latest()) self.assertTrue(self.pipeline.calls_the_latest())
def test_pipeline_context_manager(self):
ctxt = self.pipeline.consumer_for('does_nothing')
with ctxt as consumer:
self.assertIsNone(consumer())