1  """ This module contains classes - NamedFileInTemporaryDirectory, TemporaryWorkingDirectory.
   2  
   3  These classes add extra features such as creating a named file in temporary directory and
   4  creating a context manager for the working directory which is also temporary.
   5  """
   6  
   7  import os as _os
   8  from pathlib import Path
   9  from tempfile import TemporaryDirectory
  10  
  11  
  12  class NamedFileInTemporaryDirectory(object):
  13      def __init__(self, filename, mode="w+b", bufsize=-1, add_to_syspath=False, **kwds):
  14          """
  15          Open a file named `filename` in a temporary directory.
  16  
  17          This context manager is preferred over `NamedTemporaryFile` in
  18          stdlib `tempfile` when one needs to reopen the file.
  19  
  20          Arguments `mode` and `bufsize` are passed to `open`.
  21          Rest of the arguments are passed to `TemporaryDirectory`.
  22  
  23          """
  24          self._tmpdir = TemporaryDirectory(**kwds)
  25          path = Path(self._tmpdir.name) / filename
  26          encoding = None if "b" in mode else "utf-8"
  27          self.file = open(path, mode, bufsize, encoding=encoding)
  28  
  29      def cleanup(self):
  30          self.file.close()
  31          self._tmpdir.cleanup()
  32  
  33      __del__ = cleanup
  34  
  35      def __enter__(self):
  36          return self.file
  37  
  38      def __exit__(self, type, value, traceback):
  39          self.cleanup()
  40  
  41  
  42  class TemporaryWorkingDirectory(TemporaryDirectory):
  43      """
  44      Creates a temporary directory and sets the cwd to that directory.
  45      Automatically reverts to previous cwd upon cleanup.
  46      Usage example:
  47  
  48          with TemporaryWorkingDirectory() as tmpdir:
  49              ...
  50      """
  51  
  52      def __enter__(self):
  53          self.old_wd = Path.cwd()
  54          _os.chdir(self.name)
  55          return super(TemporaryWorkingDirectory, self).__enter__()
  56  
  57      def __exit__(self, exc, value, tb):
  58          _os.chdir(self.old_wd)
  59          return super(TemporaryWorkingDirectory, self).__exit__(exc, value, tb)
  60