add a with_tempdir decorator ensuring all temporary files and dirs are removed

1 files changed, 21 insertions, 0 deletions

diff --git a/testlib.py b/testlib.py
index d95bdd1..b91a297 100644
--- a/testlib.py
+++ b/testlib.py
@@ -33,11 +33,14 @@ import traceback
 import unittest
 import difflib
 import types
+import tempfile
+from shutil import rmtree
 from operator import itemgetter
 from warnings import warn
 from compiler.consts import CO_GENERATOR
 from ConfigParser import ConfigParser
 
+
 # PRINT_ = file('stdout.txt', 'w').write
 
 try:
@@ -65,6 +68,24 @@ DEFAULT_PREFIXES = ('test', 'regrtest', 'smoketest', 'unittest',
 
 ENABLE_DBC = False
 
+
+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"""
+    def proxy(*args, **kargs):
+        
+        old_tmpdir = tempfile.gettempdir()
+        new_tmpdir = tempfile.mkdtemp("-logilab-common-testlib","temp-dir-")
+        tempfile.tempdir = new_tmpdir
+        try:
+            return callable(*args, **kargs)
+        finally:
+            try:
+                rmtree(new_tmpdir, ignore_errors=True)
+            finally:
+                tempfile.tempdir = old_tmpdir
+    return proxy
+
 def main(testdir=None, exitafter=True):
     """Execute a test suite.