_tmpdirs.py 2.3 KB

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