temp_dir.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. from __future__ import absolute_import
  2. import logging
  3. import os.path
  4. import tempfile
  5. from pip._internal.utils.misc import rmtree
  6. logger = logging.getLogger(__name__)
  7. class TempDirectory(object):
  8. """Helper class that owns and cleans up a temporary directory.
  9. This class can be used as a context manager or as an OO representation of a
  10. temporary directory.
  11. Attributes:
  12. path
  13. Location to the created temporary directory or None
  14. delete
  15. Whether the directory should be deleted when exiting
  16. (when used as a contextmanager)
  17. Methods:
  18. create()
  19. Creates a temporary directory and stores its path in the path
  20. attribute.
  21. cleanup()
  22. Deletes the temporary directory and sets path attribute to None
  23. When used as a context manager, a temporary directory is created on
  24. entering the context and, if the delete attribute is True, on exiting the
  25. context the created directory is deleted.
  26. """
  27. def __init__(self, path=None, delete=None, kind="temp"):
  28. super(TempDirectory, self).__init__()
  29. if path is None and delete is None:
  30. # If we were not given an explicit directory, and we were not given
  31. # an explicit delete option, then we'll default to deleting.
  32. delete = True
  33. self.path = path
  34. self.delete = delete
  35. self.kind = kind
  36. def __repr__(self):
  37. return "<{} {!r}>".format(self.__class__.__name__, self.path)
  38. def __enter__(self):
  39. self.create()
  40. return self
  41. def __exit__(self, exc, value, tb):
  42. if self.delete:
  43. self.cleanup()
  44. def create(self):
  45. """Create a temporary directory and store it's path in self.path
  46. """
  47. if self.path is not None:
  48. logger.debug(
  49. "Skipped creation of temporary directory: {}".format(self.path)
  50. )
  51. return
  52. # We realpath here because some systems have their default tmpdir
  53. # symlinked to another directory. This tends to confuse build
  54. # scripts, so we canonicalize the path by traversing potential
  55. # symlinks here.
  56. self.path = os.path.realpath(
  57. tempfile.mkdtemp(prefix="pip-{}-".format(self.kind))
  58. )
  59. logger.debug("Created temporary directory: {}".format(self.path))
  60. def cleanup(self):
  61. """Remove the temporary directory created and reset state
  62. """
  63. if self.path is not None and os.path.exists(self.path):
  64. rmtree(self.path)
  65. self.path = None