_utils.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. import asyncio
  2. import functools
  3. import re
  4. import sys
  5. import typing
  6. from contextlib import contextmanager
  7. from starlette.types import Scope
  8. if sys.version_info >= (3, 10): # pragma: no cover
  9. from typing import TypeGuard
  10. else: # pragma: no cover
  11. from typing_extensions import TypeGuard
  12. has_exceptiongroups = True
  13. if sys.version_info < (3, 11): # pragma: no cover
  14. try:
  15. from exceptiongroup import BaseExceptionGroup
  16. except ImportError:
  17. has_exceptiongroups = False
  18. T = typing.TypeVar("T")
  19. AwaitableCallable = typing.Callable[..., typing.Awaitable[T]]
  20. @typing.overload
  21. def is_async_callable(obj: AwaitableCallable[T]) -> TypeGuard[AwaitableCallable[T]]:
  22. ...
  23. @typing.overload
  24. def is_async_callable(obj: typing.Any) -> TypeGuard[AwaitableCallable[typing.Any]]:
  25. ...
  26. def is_async_callable(obj: typing.Any) -> typing.Any:
  27. while isinstance(obj, functools.partial):
  28. obj = obj.func
  29. return asyncio.iscoroutinefunction(obj) or (
  30. callable(obj) and asyncio.iscoroutinefunction(obj.__call__)
  31. )
  32. T_co = typing.TypeVar("T_co", covariant=True)
  33. class AwaitableOrContextManager(
  34. typing.Awaitable[T_co], typing.AsyncContextManager[T_co], typing.Protocol[T_co]
  35. ):
  36. ...
  37. class SupportsAsyncClose(typing.Protocol):
  38. async def close(self) -> None:
  39. ... # pragma: no cover
  40. SupportsAsyncCloseType = typing.TypeVar(
  41. "SupportsAsyncCloseType", bound=SupportsAsyncClose, covariant=False
  42. )
  43. class AwaitableOrContextManagerWrapper(typing.Generic[SupportsAsyncCloseType]):
  44. __slots__ = ("aw", "entered")
  45. def __init__(self, aw: typing.Awaitable[SupportsAsyncCloseType]) -> None:
  46. self.aw = aw
  47. def __await__(self) -> typing.Generator[typing.Any, None, SupportsAsyncCloseType]:
  48. return self.aw.__await__()
  49. async def __aenter__(self) -> SupportsAsyncCloseType:
  50. self.entered = await self.aw
  51. return self.entered
  52. async def __aexit__(self, *args: typing.Any) -> typing.Union[None, bool]:
  53. await self.entered.close()
  54. return None
  55. @contextmanager
  56. def collapse_excgroups() -> typing.Generator[None, None, None]:
  57. try:
  58. yield
  59. except BaseException as exc:
  60. if has_exceptiongroups:
  61. while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
  62. exc = exc.exceptions[0] # pragma: no cover
  63. raise exc
  64. def get_route_path(scope: Scope) -> str:
  65. root_path = scope.get("root_path", "")
  66. route_path = re.sub(r"^" + root_path, "", scope["path"])
  67. return route_path