[testlib] check for generators in with_tempdir

If the decorated callable is a generator, we shouldn't reset tempfile.tempdir
until the end of the iteration.  Hopefully nobody uses this with nested/chained
generators, or they'd stomp all over each other's tempdirs.

Closes #117533

3 files changed, 28 insertions, 0 deletions

diff --git a/ChangeLog b/ChangeLog
index 0a2ac84..2e3b430 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,6 +1,8 @@
 ChangeLog for logilab.common
 ============================
 
+--
+    * testlib: check for generators in with_tempdir (closes #117533)
 
 2013-01-21  -- 0.59.0
 
diff --git a/test/unittest_testlib.py b/test/unittest_testlib.py
index 4ea8242..f83e1f2 100644
--- a/test/unittest_testlib.py
+++ b/test/unittest_testlib.py
@@ -628,6 +628,17 @@ class DecoratorTC(TestCase):
         self.assertListEqual(list(os.walk(tempdir)),
             [(tempdir, [], [])])
 
+    def test_tmpdir_generator(self):
+        orig_tempdir = tempfile.gettempdir()
+
+        @with_tempdir
+        def gen():
+            yield tempfile.gettempdir()
+
+        for tempdir in gen():
+            self.assertNotEqual(orig_tempdir, tempdir)
+        self.assertEqual(orig_tempdir, tempfile.gettempdir())
+
     def setUp(self):
         self.pyversion = sys.version_info
 
diff --git a/testlib.py b/testlib.py
index fe519ec..0517032 100644
--- a/testlib.py
+++ b/testlib.py
@@ -124,6 +124,21 @@ __unittest = 1
 def with_tempdir(callable):
     """A decorator ensuring no temporary file left when the function return
     Work only for temporary file create with the tempfile module"""
+    if is_generator(callable):
+        def proxy(*args, **kwargs):
+            old_tmpdir = tempfile.gettempdir()
+            new_tmpdir = tempfile.mkdtemp(prefix="temp-lgc-")
+            tempfile.tempdir = new_tmpdir
+            try:
+                for x in callable(*args, **kwargs):
+                    yield x
+            finally:
+                try:
+                    rmtree(new_tmpdir, ignore_errors=True)
+                finally:
+                    tempfile.tempdir = old_tmpdir
+        return proxy
+
     @wraps(callable)
     def proxy(*args, **kargs):