選択できるのは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