_tmpdirs.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. ''' Contexts for *with* statement providing temporary directories
  2. '''
  3. from __future__ import division, print_function, absolute_import
  4. import os
  5. from contextlib import contextmanager
  6. from shutil import rmtree
  7. from tempfile import mkdtemp
  8. @contextmanager
  9. def tempdir():
  10. """Create and return a temporary directory. This has the same
  11. behavior as mkdtemp but can be used as a context manager.
  12. Upon exiting the context, the directory and everything contained
  13. in it are removed.
  14. Examples
  15. --------
  16. >>> import os
  17. >>> with tempdir() as tmpdir:
  18. ... fname = os.path.join(tmpdir, 'example_file.txt')
  19. ... with open(fname, 'wt') as fobj:
  20. ... _ = fobj.write('a string\\n')
  21. >>> os.path.exists(tmpdir)
  22. False
  23. """
  24. d = mkdtemp()
  25. yield d
  26. rmtree(d)
  27. @contextmanager
  28. def in_tempdir():
  29. ''' Create, return, and change directory to a temporary directory
  30. Examples
  31. --------
  32. >>> import os
  33. >>> my_cwd = os.getcwd()
  34. >>> with in_tempdir() as tmpdir:
  35. ... _ = open('test.txt', 'wt').write('some text')
  36. ... assert os.path.isfile('test.txt')
  37. ... assert os.path.isfile(os.path.join(tmpdir, 'test.txt'))
  38. >>> os.path.exists(tmpdir)
  39. False
  40. >>> os.getcwd() == my_cwd
  41. True
  42. '''
  43. pwd = os.getcwd()
  44. d = mkdtemp()
  45. os.chdir(d)
  46. yield d
  47. os.chdir(pwd)
  48. rmtree(d)
  49. @contextmanager
  50. def in_dir(dir=None):
  51. """ Change directory to given directory for duration of ``with`` block
  52. Useful when you want to use `in_tempdir` for the final test, but
  53. you are still debugging. For example, you may want to do this in the end:
  54. >>> with in_tempdir() as tmpdir:
  55. ... # do something complicated which might break
  56. ... pass
  57. But indeed the complicated thing does break, and meanwhile the
  58. ``in_tempdir`` context manager wiped out the directory with the
  59. temporary files that you wanted for debugging. So, while debugging, you
  60. replace with something like:
  61. >>> with in_dir() as tmpdir: # Use working directory by default
  62. ... # do something complicated which might break
  63. ... pass
  64. You can then look at the temporary file outputs to debug what is happening,
  65. fix, and finally replace ``in_dir`` with ``in_tempdir`` again.
  66. """
  67. cwd = os.getcwd()
  68. if dir is None:
  69. yield cwd
  70. return
  71. os.chdir(dir)
  72. yield dir
  73. os.chdir(cwd)