1 #!/usr/bin/env python3 2 3 import os 4 import sys 5 6 try: 7 from tempfile import TemporaryDirectory 8 except ImportError: 9 import shutil 10 import tempfile 11 12 class TemporaryDirectory(object): 13 def __init__(self, suffix='', prefix='tmp', dir=None): 14 self.name = tempfile.mkdtemp(suffix, prefix, dir) 15 16 def __del__(self): 17 self.cleanup() 18 19 def __enter__(self): 20 return self.name 21 22 def __exit__(self, exc, value, tb): 23 self.cleanup() 24 25 def cleanup(self): 26 if self.name: 27 shutil.rmtree(self.name) 28 self.name = None 29 30 if sys.version_info >= (3, 0): 31 from os import makedirs 32 else: 33 def makedirs(path, exist_ok): 34 if exist_ok and os.path.exists(path): 35 return 36 return os.makedirs(path) 37 38 if sys.version_info >= (3, 0): 39 from io import StringIO 40 else: 41 from StringIO import StringIO 42 43 44 try: 45 from unittest.mock import patch 46 except ImportError: 47 import contextlib 48 @contextlib.contextmanager 49 def patch(target, mock): 50 obj, attr = target.rsplit('.') 51 obj = __import__(obj) 52 original_value = getattr(obj, attr) 53 setattr(obj, attr, mock) 54 try: 55 yield 56 finally: 57 setattr(obj, attr, original_value) 58
