_fixes.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. # JSONDecodeError was introduced in requests=2.27 released in 2022.
  2. # This allows us to support older requests for users
  3. # More information: https://github.com/psf/requests/pull/5856
  4. try:
  5. from requests import JSONDecodeError # type: ignore # noqa: F401
  6. except ImportError:
  7. try:
  8. from simplejson import JSONDecodeError # type: ignore # noqa: F401
  9. except ImportError:
  10. from json import JSONDecodeError # type: ignore # noqa: F401
  11. import contextlib
  12. import os
  13. import shutil
  14. import stat
  15. import tempfile
  16. import time
  17. from functools import partial
  18. from pathlib import Path
  19. from typing import Callable, Generator, Optional, Union
  20. import yaml
  21. from filelock import BaseFileLock, FileLock, SoftFileLock, Timeout
  22. from .. import constants
  23. from . import logging
  24. logger = logging.get_logger(__name__)
  25. # Wrap `yaml.dump` to set `allow_unicode=True` by default.
  26. #
  27. # Example:
  28. # ```py
  29. # >>> yaml.dump({"emoji": "👀", "some unicode": "日本か"})
  30. # 'emoji: "\\U0001F440"\nsome unicode: "\\u65E5\\u672C\\u304B"\n'
  31. #
  32. # >>> yaml_dump({"emoji": "👀", "some unicode": "日本か"})
  33. # 'emoji: "👀"\nsome unicode: "日本か"\n'
  34. # ```
  35. yaml_dump: Callable[..., str] = partial(yaml.dump, stream=None, allow_unicode=True) # type: ignore
  36. @contextlib.contextmanager
  37. def SoftTemporaryDirectory(
  38. suffix: Optional[str] = None,
  39. prefix: Optional[str] = None,
  40. dir: Optional[Union[Path, str]] = None,
  41. **kwargs,
  42. ) -> Generator[Path, None, None]:
  43. """
  44. Context manager to create a temporary directory and safely delete it.
  45. If tmp directory cannot be deleted normally, we set the WRITE permission and retry.
  46. If cleanup still fails, we give up but don't raise an exception. This is equivalent
  47. to `tempfile.TemporaryDirectory(..., ignore_cleanup_errors=True)` introduced in
  48. Python 3.10.
  49. See https://www.scivision.dev/python-tempfile-permission-error-windows/.
  50. """
  51. tmpdir = tempfile.TemporaryDirectory(prefix=prefix, suffix=suffix, dir=dir, **kwargs)
  52. yield Path(tmpdir.name).resolve()
  53. try:
  54. # First once with normal cleanup
  55. shutil.rmtree(tmpdir.name)
  56. except Exception:
  57. # If failed, try to set write permission and retry
  58. try:
  59. shutil.rmtree(tmpdir.name, onerror=_set_write_permission_and_retry)
  60. except Exception:
  61. pass
  62. # And finally, cleanup the tmpdir.
  63. # If it fails again, give up but do not throw error
  64. try:
  65. tmpdir.cleanup()
  66. except Exception:
  67. pass
  68. def _set_write_permission_and_retry(func, path, excinfo):
  69. os.chmod(path, stat.S_IWRITE)
  70. func(path)
  71. @contextlib.contextmanager
  72. def WeakFileLock(
  73. lock_file: Union[str, Path], *, timeout: Optional[float] = None
  74. ) -> Generator[BaseFileLock, None, None]:
  75. """A filelock with some custom logic.
  76. This filelock is weaker than the default filelock in that:
  77. 1. It won't raise an exception if release fails.
  78. 2. It will default to a SoftFileLock if the filesystem does not support flock.
  79. An INFO log message is emitted every 10 seconds if the lock is not acquired immediately.
  80. If a timeout is provided, a `filelock.Timeout` exception is raised if the lock is not acquired within the timeout.
  81. """
  82. log_interval = constants.FILELOCK_LOG_EVERY_SECONDS
  83. lock = FileLock(lock_file, timeout=log_interval)
  84. start_time = time.time()
  85. while True:
  86. elapsed_time = time.time() - start_time
  87. if timeout is not None and elapsed_time >= timeout:
  88. raise Timeout(str(lock_file))
  89. try:
  90. lock.acquire(timeout=min(log_interval, timeout - elapsed_time) if timeout else log_interval)
  91. except Timeout:
  92. logger.info(
  93. f"Still waiting to acquire lock on {lock_file} (elapsed: {time.time() - start_time:.1f} seconds)"
  94. )
  95. except NotImplementedError as e:
  96. if "use SoftFileLock instead" in str(e):
  97. logger.warning(
  98. "FileSystem does not appear to support flock. Falling back to SoftFileLock for %s", lock_file
  99. )
  100. lock = SoftFileLock(lock_file, timeout=log_interval)
  101. continue
  102. else:
  103. break
  104. try:
  105. yield lock
  106. finally:
  107. try:
  108. lock.release()
  109. except OSError:
  110. try:
  111. Path(lock_file).unlink()
  112. except OSError:
  113. pass