errors.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. from __future__ import annotations
  2. import contextlib
  3. import dataclasses
  4. import sys
  5. import typing
  6. __all__ = ["ExceptionGroup"]
  7. def __dir__() -> list[str]:
  8. return __all__
  9. if sys.version_info >= (3, 11): # pragma: no cover
  10. from builtins import ExceptionGroup
  11. else: # pragma: no cover
  12. class ExceptionGroup(Exception):
  13. """A minimal implementation of :external:exc:`ExceptionGroup` from Python 3.11.
  14. If :external:exc:`ExceptionGroup` is already defined by Python itself,
  15. that version is used instead.
  16. """
  17. message: str
  18. exceptions: list[Exception]
  19. def __init__(self, message: str, exceptions: list[Exception]) -> None:
  20. self.message = message
  21. self.exceptions = exceptions
  22. def __repr__(self) -> str:
  23. return f"{self.__class__.__name__}({self.message!r}, {self.exceptions!r})"
  24. @dataclasses.dataclass
  25. class _ErrorCollector:
  26. """
  27. Collect errors into ExceptionGroups.
  28. Used like this:
  29. collector = _ErrorCollector()
  30. # Add a single exception
  31. collector.error(ValueError("one"))
  32. # Supports nesting, including combining ExceptionGroups
  33. with collector.collect():
  34. raise ValueError("two")
  35. collector.finalize("Found some errors")
  36. Since making a collector and then calling finalize later is a common pattern,
  37. a convenience method ``on_exit`` is provided.
  38. """
  39. errors: list[Exception] = dataclasses.field(default_factory=list, init=False)
  40. def finalize(self, msg: str) -> None:
  41. """Raise a group exception if there are any errors."""
  42. if self.errors:
  43. raise ExceptionGroup(msg, self.errors)
  44. @contextlib.contextmanager
  45. def on_exit(self, msg: str) -> typing.Generator[_ErrorCollector, None, None]:
  46. """
  47. Calls finalize if no uncollected errors were present.
  48. Uncollected errors are raised normally.
  49. """
  50. yield self
  51. self.finalize(msg)
  52. @contextlib.contextmanager
  53. def collect(self, *err_cls: type[Exception]) -> typing.Generator[None, None, None]:
  54. """
  55. Context manager to collect errors into the error list.
  56. Must be inside loops, as only one error can be collected at a time.
  57. """
  58. error_classes = err_cls or (Exception,)
  59. try:
  60. yield
  61. except ExceptionGroup as error:
  62. self.errors.extend(error.exceptions)
  63. except error_classes as error:
  64. self.errors.append(error)
  65. def error(
  66. self,
  67. error: Exception,
  68. ) -> None:
  69. """Add an error to the list."""
  70. self.errors.append(error)