Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 
 

146 строки
3.4 KiB

  1. import contextlib
  2. import io
  3. import os
  4. import shutil
  5. import site
  6. import sys
  7. import tempfile
  8. from filelock import FileLock
  9. @contextlib.contextmanager
  10. def tempdir(cd=lambda dir: None, **kwargs):
  11. temp_dir = tempfile.mkdtemp(**kwargs)
  12. orig_dir = os.getcwd()
  13. try:
  14. cd(temp_dir)
  15. yield temp_dir
  16. finally:
  17. cd(orig_dir)
  18. shutil.rmtree(temp_dir)
  19. @contextlib.contextmanager
  20. def environment(**replacements):
  21. """
  22. In a context, patch the environment with replacements. Pass None values
  23. to clear the values.
  24. """
  25. saved = dict((key, os.environ[key]) for key in replacements if key in os.environ)
  26. # remove values that are null
  27. remove = (key for (key, value) in replacements.items() if value is None)
  28. for key in list(remove):
  29. os.environ.pop(key, None)
  30. replacements.pop(key)
  31. os.environ.update(replacements)
  32. try:
  33. yield saved
  34. finally:
  35. for key in replacements:
  36. os.environ.pop(key, None)
  37. os.environ.update(saved)
  38. @contextlib.contextmanager
  39. def quiet():
  40. """
  41. Redirect stdout/stderr to StringIO objects to prevent console output from
  42. distutils commands.
  43. """
  44. old_stdout = sys.stdout
  45. old_stderr = sys.stderr
  46. new_stdout = sys.stdout = io.StringIO()
  47. new_stderr = sys.stderr = io.StringIO()
  48. try:
  49. yield new_stdout, new_stderr
  50. finally:
  51. new_stdout.seek(0)
  52. new_stderr.seek(0)
  53. sys.stdout = old_stdout
  54. sys.stderr = old_stderr
  55. @contextlib.contextmanager
  56. def save_user_site_setting():
  57. saved = site.ENABLE_USER_SITE
  58. try:
  59. yield saved
  60. finally:
  61. site.ENABLE_USER_SITE = saved
  62. @contextlib.contextmanager
  63. def save_pkg_resources_state():
  64. import pkg_resources
  65. pr_state = pkg_resources.__getstate__()
  66. # also save sys.path
  67. sys_path = sys.path[:]
  68. try:
  69. yield pr_state, sys_path
  70. finally:
  71. sys.path[:] = sys_path
  72. pkg_resources.__setstate__(pr_state)
  73. @contextlib.contextmanager
  74. def suppress_exceptions(*excs):
  75. try:
  76. yield
  77. except excs:
  78. pass
  79. def multiproc(request):
  80. """
  81. Return True if running under xdist and multiple
  82. workers are used.
  83. """
  84. try:
  85. worker_id = request.getfixturevalue('worker_id')
  86. except Exception:
  87. return False
  88. return worker_id != 'master'
  89. @contextlib.contextmanager
  90. def session_locked_tmp_dir(request, tmp_path_factory, name):
  91. """Uses a file lock to guarantee only one worker can access a temp dir"""
  92. # get the temp directory shared by all workers
  93. base = tmp_path_factory.getbasetemp()
  94. shared_dir = base.parent if multiproc(request) else base
  95. locked_dir = shared_dir / name
  96. with FileLock(locked_dir.with_suffix(".lock")):
  97. # ^-- prevent multiple workers to access the directory at once
  98. locked_dir.mkdir(exist_ok=True, parents=True)
  99. yield locked_dir
  100. @contextlib.contextmanager
  101. def save_paths():
  102. """Make sure ``sys.path``, ``sys.meta_path`` and ``sys.path_hooks`` are preserved"""
  103. prev = sys.path[:], sys.meta_path[:], sys.path_hooks[:]
  104. try:
  105. yield
  106. finally:
  107. sys.path, sys.meta_path, sys.path_hooks = prev
  108. @contextlib.contextmanager
  109. def save_sys_modules():
  110. """Make sure initial ``sys.modules`` is preserved"""
  111. prev_modules = sys.modules
  112. try:
  113. sys.modules = sys.modules.copy()
  114. yield
  115. finally:
  116. sys.modules = prev_modules