diff --git a/lib/python3.12/site-packages/aiohttp/.hash/_cparser.pxd.hash b/lib/python3.12/site-packages/aiohttp/.hash/_cparser.pxd.hash new file mode 100644 index 0000000000000000000000000000000000000000..3f5cd0e6f720a049ee19149790c808558b492726 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/.hash/_cparser.pxd.hash @@ -0,0 +1 @@ +5276d46021e0e0d7577e0c9155800cbf62932d60a50783fec42aefb63febedec /home/runner/work/aiohttp/aiohttp/aiohttp/_cparser.pxd diff --git a/lib/python3.12/site-packages/aiohttp/.hash/_find_header.pxd.hash b/lib/python3.12/site-packages/aiohttp/.hash/_find_header.pxd.hash new file mode 100644 index 0000000000000000000000000000000000000000..f006c2de5d24a1b5a9c26f83c858127b5e12b07c --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/.hash/_find_header.pxd.hash @@ -0,0 +1 @@ +d067f01423cddb3c442933b5fcc039b18ab651fcec1bc91c577693aafc25cf78 /home/runner/work/aiohttp/aiohttp/aiohttp/_find_header.pxd diff --git a/lib/python3.12/site-packages/aiohttp/.hash/_http_parser.pyx.hash b/lib/python3.12/site-packages/aiohttp/.hash/_http_parser.pyx.hash new file mode 100644 index 0000000000000000000000000000000000000000..1099b5e430125d0a2ced0ddb059110bf97f2109f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/.hash/_http_parser.pyx.hash @@ -0,0 +1 @@ +f9823c608638b8a77fec1c2bd28dc5c043f08c775ec59e7e262dd74b4ebb7384 /home/runner/work/aiohttp/aiohttp/aiohttp/_http_parser.pyx diff --git a/lib/python3.12/site-packages/aiohttp/.hash/_http_writer.pyx.hash b/lib/python3.12/site-packages/aiohttp/.hash/_http_writer.pyx.hash new file mode 100644 index 0000000000000000000000000000000000000000..ad351fe2f96dd4de3a6d83a904ce80802c69908b --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/.hash/_http_writer.pyx.hash @@ -0,0 +1 @@ +56514404ce87a15bfc6b400026d73a270165b2fdbe70313cfa007de29ddd7e14 /home/runner/work/aiohttp/aiohttp/aiohttp/_http_writer.pyx diff --git a/lib/python3.12/site-packages/aiohttp/.hash/hdrs.py.hash b/lib/python3.12/site-packages/aiohttp/.hash/hdrs.py.hash new file mode 100644 index 0000000000000000000000000000000000000000..c8d55240e6a55c305c99244966dee5615565716b --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/.hash/hdrs.py.hash @@ -0,0 +1 @@ +dab8f933203eeb245d60f856e542a45b888d5a110094620e4811f90f816628d1 /home/runner/work/aiohttp/aiohttp/aiohttp/hdrs.py diff --git a/lib/python3.12/site-packages/aiohttp/__init__.py b/lib/python3.12/site-packages/aiohttp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..357baf019de77a08af52b18334cffcc9f078390a --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/__init__.py @@ -0,0 +1,278 @@ +__version__ = "3.13.3" + +from typing import TYPE_CHECKING, Tuple + +from . import hdrs as hdrs +from .client import ( + BaseConnector, + ClientConnectionError, + ClientConnectionResetError, + ClientConnectorCertificateError, + ClientConnectorDNSError, + ClientConnectorError, + ClientConnectorSSLError, + ClientError, + ClientHttpProxyError, + ClientOSError, + ClientPayloadError, + ClientProxyConnectionError, + ClientRequest, + ClientResponse, + ClientResponseError, + ClientSession, + ClientSSLError, + ClientTimeout, + ClientWebSocketResponse, + ClientWSTimeout, + ConnectionTimeoutError, + ContentTypeError, + Fingerprint, + InvalidURL, + InvalidUrlClientError, + InvalidUrlRedirectClientError, + NamedPipeConnector, + NonHttpUrlClientError, + NonHttpUrlRedirectClientError, + RedirectClientError, + RequestInfo, + ServerConnectionError, + ServerDisconnectedError, + ServerFingerprintMismatch, + ServerTimeoutError, + SocketTimeoutError, + TCPConnector, + TooManyRedirects, + UnixConnector, + WSMessageTypeError, + WSServerHandshakeError, + request, +) +from .client_middleware_digest_auth import DigestAuthMiddleware +from .client_middlewares import ClientHandlerType, ClientMiddlewareType +from .compression_utils import set_zlib_backend +from .connector import ( + AddrInfoType as AddrInfoType, + SocketFactoryType as SocketFactoryType, +) +from .cookiejar import CookieJar as CookieJar, DummyCookieJar as DummyCookieJar +from .formdata import FormData as FormData +from .helpers import BasicAuth, ChainMapProxy, ETag +from .http import ( + HttpVersion as HttpVersion, + HttpVersion10 as HttpVersion10, + HttpVersion11 as HttpVersion11, + WebSocketError as WebSocketError, + WSCloseCode as WSCloseCode, + WSMessage as WSMessage, + WSMsgType as WSMsgType, +) +from .multipart import ( + BadContentDispositionHeader as BadContentDispositionHeader, + BadContentDispositionParam as BadContentDispositionParam, + BodyPartReader as BodyPartReader, + MultipartReader as MultipartReader, + MultipartWriter as MultipartWriter, + content_disposition_filename as content_disposition_filename, + parse_content_disposition as parse_content_disposition, +) +from .payload import ( + PAYLOAD_REGISTRY as PAYLOAD_REGISTRY, + AsyncIterablePayload as AsyncIterablePayload, + BufferedReaderPayload as BufferedReaderPayload, + BytesIOPayload as BytesIOPayload, + BytesPayload as BytesPayload, + IOBasePayload as IOBasePayload, + JsonPayload as JsonPayload, + Payload as Payload, + StringIOPayload as StringIOPayload, + StringPayload as StringPayload, + TextIOPayload as TextIOPayload, + get_payload as get_payload, + payload_type as payload_type, +) +from .payload_streamer import streamer as streamer +from .resolver import ( + AsyncResolver as AsyncResolver, + DefaultResolver as DefaultResolver, + ThreadedResolver as ThreadedResolver, +) +from .streams import ( + EMPTY_PAYLOAD as EMPTY_PAYLOAD, + DataQueue as DataQueue, + EofStream as EofStream, + FlowControlDataQueue as FlowControlDataQueue, + StreamReader as StreamReader, +) +from .tracing import ( + TraceConfig as TraceConfig, + TraceConnectionCreateEndParams as TraceConnectionCreateEndParams, + TraceConnectionCreateStartParams as TraceConnectionCreateStartParams, + TraceConnectionQueuedEndParams as TraceConnectionQueuedEndParams, + TraceConnectionQueuedStartParams as TraceConnectionQueuedStartParams, + TraceConnectionReuseconnParams as TraceConnectionReuseconnParams, + TraceDnsCacheHitParams as TraceDnsCacheHitParams, + TraceDnsCacheMissParams as TraceDnsCacheMissParams, + TraceDnsResolveHostEndParams as TraceDnsResolveHostEndParams, + TraceDnsResolveHostStartParams as TraceDnsResolveHostStartParams, + TraceRequestChunkSentParams as TraceRequestChunkSentParams, + TraceRequestEndParams as TraceRequestEndParams, + TraceRequestExceptionParams as TraceRequestExceptionParams, + TraceRequestHeadersSentParams as TraceRequestHeadersSentParams, + TraceRequestRedirectParams as TraceRequestRedirectParams, + TraceRequestStartParams as TraceRequestStartParams, + TraceResponseChunkReceivedParams as TraceResponseChunkReceivedParams, +) + +if TYPE_CHECKING: + # At runtime these are lazy-loaded at the bottom of the file. + from .worker import ( + GunicornUVLoopWebWorker as GunicornUVLoopWebWorker, + GunicornWebWorker as GunicornWebWorker, + ) + +__all__: Tuple[str, ...] = ( + "hdrs", + # client + "AddrInfoType", + "BaseConnector", + "ClientConnectionError", + "ClientConnectionResetError", + "ClientConnectorCertificateError", + "ClientConnectorDNSError", + "ClientConnectorError", + "ClientConnectorSSLError", + "ClientError", + "ClientHttpProxyError", + "ClientOSError", + "ClientPayloadError", + "ClientProxyConnectionError", + "ClientResponse", + "ClientRequest", + "ClientResponseError", + "ClientSSLError", + "ClientSession", + "ClientTimeout", + "ClientWebSocketResponse", + "ClientWSTimeout", + "ConnectionTimeoutError", + "ContentTypeError", + "Fingerprint", + "FlowControlDataQueue", + "InvalidURL", + "InvalidUrlClientError", + "InvalidUrlRedirectClientError", + "NonHttpUrlClientError", + "NonHttpUrlRedirectClientError", + "RedirectClientError", + "RequestInfo", + "ServerConnectionError", + "ServerDisconnectedError", + "ServerFingerprintMismatch", + "ServerTimeoutError", + "SocketFactoryType", + "SocketTimeoutError", + "TCPConnector", + "TooManyRedirects", + "UnixConnector", + "NamedPipeConnector", + "WSServerHandshakeError", + "request", + # client_middleware + "ClientMiddlewareType", + "ClientHandlerType", + # cookiejar + "CookieJar", + "DummyCookieJar", + # formdata + "FormData", + # helpers + "BasicAuth", + "ChainMapProxy", + "DigestAuthMiddleware", + "ETag", + "set_zlib_backend", + # http + "HttpVersion", + "HttpVersion10", + "HttpVersion11", + "WSMsgType", + "WSCloseCode", + "WSMessage", + "WebSocketError", + # multipart + "BadContentDispositionHeader", + "BadContentDispositionParam", + "BodyPartReader", + "MultipartReader", + "MultipartWriter", + "content_disposition_filename", + "parse_content_disposition", + # payload + "AsyncIterablePayload", + "BufferedReaderPayload", + "BytesIOPayload", + "BytesPayload", + "IOBasePayload", + "JsonPayload", + "PAYLOAD_REGISTRY", + "Payload", + "StringIOPayload", + "StringPayload", + "TextIOPayload", + "get_payload", + "payload_type", + # payload_streamer + "streamer", + # resolver + "AsyncResolver", + "DefaultResolver", + "ThreadedResolver", + # streams + "DataQueue", + "EMPTY_PAYLOAD", + "EofStream", + "StreamReader", + # tracing + "TraceConfig", + "TraceConnectionCreateEndParams", + "TraceConnectionCreateStartParams", + "TraceConnectionQueuedEndParams", + "TraceConnectionQueuedStartParams", + "TraceConnectionReuseconnParams", + "TraceDnsCacheHitParams", + "TraceDnsCacheMissParams", + "TraceDnsResolveHostEndParams", + "TraceDnsResolveHostStartParams", + "TraceRequestChunkSentParams", + "TraceRequestEndParams", + "TraceRequestExceptionParams", + "TraceRequestHeadersSentParams", + "TraceRequestRedirectParams", + "TraceRequestStartParams", + "TraceResponseChunkReceivedParams", + # workers (imported lazily with __getattr__) + "GunicornUVLoopWebWorker", + "GunicornWebWorker", + "WSMessageTypeError", +) + + +def __dir__() -> Tuple[str, ...]: + return __all__ + ("__doc__",) + + +def __getattr__(name: str) -> object: + global GunicornUVLoopWebWorker, GunicornWebWorker + + # Importing gunicorn takes a long time (>100ms), so only import if actually needed. + if name in ("GunicornUVLoopWebWorker", "GunicornWebWorker"): + try: + from .worker import GunicornUVLoopWebWorker as guv, GunicornWebWorker as gw + except ImportError: + return None + + GunicornUVLoopWebWorker = guv # type: ignore[misc] + GunicornWebWorker = gw # type: ignore[misc] + return guv if name == "GunicornUVLoopWebWorker" else gw + + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72ddf2b6144d8cd5fdbabcf3b6ad85c15748cd6b Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/_cookie_helpers.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/_cookie_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97be6acc1e45e1445b93acf9a733c40e52c81bb8 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/_cookie_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/abc.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/abc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53b3d0c849e206caef63b7c93f08e8aec60cc3d8 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/abc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/base_protocol.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/base_protocol.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b529d7b38981ee4626a7e77dc18889c78439085a Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/base_protocol.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42f70fb09ba09b1b0eced25640d9208066bcd000 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_exceptions.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92ff62f2439f21e5f10e9e95447544b745d2ca52 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_middleware_digest_auth.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_middleware_digest_auth.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91a641605ac0cbb4d72476d4eef1c431925bc5de Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_middleware_digest_auth.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_middlewares.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_middlewares.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0f424568ad24dbf33db1f7b04cd3e2d4cb11257 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_middlewares.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_proto.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_proto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d8e0d5e3a593cd11fad5b09c00d236c34d25662c Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_proto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_reqrep.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_reqrep.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbd42a49b0f812b7aa0930f2b02949b6e9e5c53a Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_reqrep.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/client_ws.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/client_ws.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bf087039bc877d703656ade3daf7f100b70395b Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/client_ws.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/compression_utils.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/compression_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48ab675ba560c071e7c3195b7dcde871cd8fcd5c Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/compression_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/connector.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/connector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e38d00ae1286b34793232a45c1b805bdc1f7703 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/connector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/cookiejar.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/cookiejar.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b73ee0806287dfb5749aff13d05291dfab95655 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/cookiejar.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/formdata.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/formdata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67c70a1e1a9b2b9bb8f80b0f06cba3794975bf4c Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/formdata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/hdrs.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/hdrs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56bc7374e54384624e61d0a2d3ec7beeba139d35 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/hdrs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/helpers.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..137307f8a32c8b7a82d8d612983b69aad10faf33 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/http.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/http.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..276742159e3fcdf944d8e73584bdcd9504583911 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/http.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/http_exceptions.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/http_exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4162b9b1914f44b4d9ea95097272e4a6ea8a817b Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/http_exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/http_parser.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/http_parser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2bc16aef205a0dfb1acddcb34e0a071493c6f04 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/http_parser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/http_websocket.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/http_websocket.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67c5808822a9ed8ccbded217d7fdda237bf7e641 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/http_websocket.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/http_writer.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/http_writer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c48fa232cf0a23387c6bca4cee58eeea2e4121c Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/http_writer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/log.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/log.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b40eeda2da7c4aca55bbb64fefd2dffc27cd19f8 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/log.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/multipart.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/multipart.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a20061f46e2286bbf7455fc252cc953ceebaf6d Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/multipart.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/payload.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/payload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d003aac3ae8d363c82c1177368d503f9c31bd1ed Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/payload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/payload_streamer.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/payload_streamer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4316993381b8f17d7f8a4cde2aabcbfdf2f50f6 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/payload_streamer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/pytest_plugin.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/pytest_plugin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fb5377d2259e65e6d3ec21fefbf81de31672e00 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/pytest_plugin.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/resolver.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/resolver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..115196b3c24ae7a0a1ecb2d7e0080911acbbfe43 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/resolver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/streams.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/streams.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..915b5e9a562e21c8b617c5e47ab9d5c92d3d6177 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/streams.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/tcp_helpers.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/tcp_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5086c46b864a76fcf9b05b906c4ea57b4f427cf Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/tcp_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/test_utils.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7f241259129ad0130e8636dd8552271221f5134 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/tracing.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/tracing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e6cfdbb5c146d8c007b5bebfb6e4a5f87ef15b8 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/tracing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/typedefs.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/typedefs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..263e2b9dcabf242e46b41d625504ac63155ad899 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/typedefs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05fb1879309776c95c2ef1a3c6e3effb9d626ff2 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_app.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_app.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42ebcd56d875799ab9ac15835e72d3b02cb850a7 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_app.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_exceptions.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..426ea7fc56cbb4d99701aeb8d23f189ae314fd82 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_fileresponse.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_fileresponse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7d798987542ad0d112753b1c45f970080eb4426 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_fileresponse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_log.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_log.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9611a9c0e872901371b9fa68ce1cf60a37f49f4e Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_log.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_middlewares.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_middlewares.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adea57bae8200f84685135357fd071a914d11362 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_middlewares.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_protocol.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_protocol.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8e54e287a622f34cb85cd7ec4c622434f43f566 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_protocol.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_request.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_request.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..546d59a6918825d6fcd6b8bb4ca04d13cc97ad7d Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_request.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_response.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_response.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30eab111bf174f847cef60bfa513c93ee5e6c362 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_response.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_routedef.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_routedef.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a0541f47a19ba386614186d5a9a5b8c1007d451 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_routedef.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_runner.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce734865d88b6cf29139df382b1cf0543ee1bdbf Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_server.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_server.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6785d4e30d71b472b8896ca769433c39146b2977 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_server.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_urldispatcher.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_urldispatcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8676519ec18f45ec8c07e8c9ca6c7a8444d0bd2 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_urldispatcher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/web_ws.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/web_ws.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8d346a99b5558f2a47fcf39e8a3e04a633f7aed Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/web_ws.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/__pycache__/worker.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/__pycache__/worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..405388d9b8aeea1f09d00f4f7ea1202ac43fd850 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/__pycache__/worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_cookie_helpers.py b/lib/python3.12/site-packages/aiohttp/_cookie_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..10e2e0eff9de9d28ecbfe511a1227a1fec1d83cd --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_cookie_helpers.py @@ -0,0 +1,338 @@ +""" +Internal cookie handling helpers. + +This module contains internal utilities for cookie parsing and manipulation. +These are not part of the public API and may change without notice. +""" + +import re +from http.cookies import Morsel +from typing import List, Optional, Sequence, Tuple, cast + +from .log import internal_logger + +__all__ = ( + "parse_set_cookie_headers", + "parse_cookie_header", + "preserve_morsel_with_coded_value", +) + +# Cookie parsing constants +# Allow more characters in cookie names to handle real-world cookies +# that don't strictly follow RFC standards (fixes #2683) +# RFC 6265 defines cookie-name token as per RFC 2616 Section 2.2, +# but many servers send cookies with characters like {} [] () etc. +# This makes the cookie parser more tolerant of real-world cookies +# while still providing some validation to catch obviously malformed names. +_COOKIE_NAME_RE = re.compile(r"^[!#$%&\'()*+\-./0-9:<=>?@A-Z\[\]^_`a-z{|}~]+$") +_COOKIE_KNOWN_ATTRS = frozenset( # AKA Morsel._reserved + ( + "path", + "domain", + "max-age", + "expires", + "secure", + "httponly", + "samesite", + "partitioned", + "version", + "comment", + ) +) +_COOKIE_BOOL_ATTRS = frozenset( # AKA Morsel._flags + ("secure", "httponly", "partitioned") +) + +# SimpleCookie's pattern for parsing cookies with relaxed validation +# Based on http.cookies pattern but extended to allow more characters in cookie names +# to handle real-world cookies (fixes #2683) +_COOKIE_PATTERN = re.compile( + r""" + \s* # Optional whitespace at start of cookie + (?P # Start of group 'key' + # aiohttp has extended to include [] for compatibility with real-world cookies + [\w\d!#%&'~_`><@,:/\$\*\+\-\.\^\|\)\(\?\}\{\[\]]+ # Any word of at least one letter + ) # End of group 'key' + ( # Optional group: there may not be a value. + \s*=\s* # Equal Sign + (?P # Start of group 'val' + "(?:[^\\"]|\\.)*" # Any double-quoted string (properly closed) + | # or + "[^";]* # Unmatched opening quote (differs from SimpleCookie - issue #7993) + | # or + # Special case for "expires" attr - RFC 822, RFC 850, RFC 1036, RFC 1123 + (\w{3,6}day|\w{3}),\s # Day of the week or abbreviated day (with comma) + [\w\d\s-]{9,11}\s[\d:]{8}\s # Date and time in specific format + (GMT|[+-]\d{4}) # Timezone: GMT or RFC 2822 offset like -0000, +0100 + # NOTE: RFC 2822 timezone support is an aiohttp extension + # for issue #4493 - SimpleCookie does NOT support this + | # or + # ANSI C asctime() format: "Wed Jun 9 10:18:14 2021" + # NOTE: This is an aiohttp extension for issue #4327 - SimpleCookie does NOT support this format + \w{3}\s+\w{3}\s+[\s\d]\d\s+\d{2}:\d{2}:\d{2}\s+\d{4} + | # or + [\w\d!#%&'~_`><@,:/\$\*\+\-\.\^\|\)\(\?\}\{\=\[\]]* # Any word or empty string + ) # End of group 'val' + )? # End of optional value group + \s* # Any number of spaces. + (\s+|;|$) # Ending either at space, semicolon, or EOS. + """, + re.VERBOSE | re.ASCII, +) + + +def preserve_morsel_with_coded_value(cookie: Morsel[str]) -> Morsel[str]: + """ + Preserve a Morsel's coded_value exactly as received from the server. + + This function ensures that cookie encoding is preserved exactly as sent by + the server, which is critical for compatibility with old servers that have + strict requirements about cookie formats. + + This addresses the issue described in https://github.com/aio-libs/aiohttp/pull/1453 + where Python's SimpleCookie would re-encode cookies, breaking authentication + with certain servers. + + Args: + cookie: A Morsel object from SimpleCookie + + Returns: + A Morsel object with preserved coded_value + + """ + mrsl_val = cast("Morsel[str]", cookie.get(cookie.key, Morsel())) + # We use __setstate__ instead of the public set() API because it allows us to + # bypass validation and set already validated state. This is more stable than + # setting protected attributes directly and unlikely to change since it would + # break pickling. + mrsl_val.__setstate__( # type: ignore[attr-defined] + {"key": cookie.key, "value": cookie.value, "coded_value": cookie.coded_value} + ) + return mrsl_val + + +_unquote_sub = re.compile(r"\\(?:([0-3][0-7][0-7])|(.))").sub + + +def _unquote_replace(m: re.Match[str]) -> str: + """ + Replace function for _unquote_sub regex substitution. + + Handles escaped characters in cookie values: + - Octal sequences are converted to their character representation + - Other escaped characters are unescaped by removing the backslash + """ + if m[1]: + return chr(int(m[1], 8)) + return m[2] + + +def _unquote(value: str) -> str: + """ + Unquote a cookie value. + + Vendored from http.cookies._unquote to ensure compatibility. + + Note: The original implementation checked for None, but we've removed + that check since all callers already ensure the value is not None. + """ + # If there aren't any doublequotes, + # then there can't be any special characters. See RFC 2109. + if len(value) < 2: + return value + if value[0] != '"' or value[-1] != '"': + return value + + # We have to assume that we must decode this string. + # Down to work. + + # Remove the "s + value = value[1:-1] + + # Check for special sequences. Examples: + # \012 --> \n + # \" --> " + # + return _unquote_sub(_unquote_replace, value) + + +def parse_cookie_header(header: str) -> List[Tuple[str, Morsel[str]]]: + """ + Parse a Cookie header according to RFC 6265 Section 5.4. + + Cookie headers contain only name-value pairs separated by semicolons. + There are no attributes in Cookie headers - even names that match + attribute names (like 'path' or 'secure') should be treated as cookies. + + This parser uses the same regex-based approach as parse_set_cookie_headers + to properly handle quoted values that may contain semicolons. When the + regex fails to match a malformed cookie, it falls back to simple parsing + to ensure subsequent cookies are not lost + https://github.com/aio-libs/aiohttp/issues/11632 + + Args: + header: The Cookie header value to parse + + Returns: + List of (name, Morsel) tuples for compatibility with SimpleCookie.update() + """ + if not header: + return [] + + cookies: List[Tuple[str, Morsel[str]]] = [] + morsel: Morsel[str] + i = 0 + n = len(header) + + invalid_names = [] + while i < n: + # Use the same pattern as parse_set_cookie_headers to find cookies + match = _COOKIE_PATTERN.match(header, i) + if not match: + # Fallback for malformed cookies https://github.com/aio-libs/aiohttp/issues/11632 + # Find next semicolon to skip or attempt simple key=value parsing + next_semi = header.find(";", i) + eq_pos = header.find("=", i) + + # Try to extract key=value if '=' comes before ';' + if eq_pos != -1 and (next_semi == -1 or eq_pos < next_semi): + end_pos = next_semi if next_semi != -1 else n + key = header[i:eq_pos].strip() + value = header[eq_pos + 1 : end_pos].strip() + + # Validate the name (same as regex path) + if not _COOKIE_NAME_RE.match(key): + invalid_names.append(key) + else: + morsel = Morsel() + morsel.__setstate__( # type: ignore[attr-defined] + {"key": key, "value": _unquote(value), "coded_value": value} + ) + cookies.append((key, morsel)) + + # Move to next cookie or end + i = next_semi + 1 if next_semi != -1 else n + continue + + key = match.group("key") + value = match.group("val") or "" + i = match.end(0) + + # Validate the name + if not key or not _COOKIE_NAME_RE.match(key): + invalid_names.append(key) + continue + + # Create new morsel + morsel = Morsel() + # Preserve the original value as coded_value (with quotes if present) + # We use __setstate__ instead of the public set() API because it allows us to + # bypass validation and set already validated state. This is more stable than + # setting protected attributes directly and unlikely to change since it would + # break pickling. + morsel.__setstate__( # type: ignore[attr-defined] + {"key": key, "value": _unquote(value), "coded_value": value} + ) + + cookies.append((key, morsel)) + + if invalid_names: + internal_logger.debug( + "Cannot load cookie. Illegal cookie names: %r", invalid_names + ) + + return cookies + + +def parse_set_cookie_headers(headers: Sequence[str]) -> List[Tuple[str, Morsel[str]]]: + """ + Parse cookie headers using a vendored version of SimpleCookie parsing. + + This implementation is based on SimpleCookie.__parse_string to ensure + compatibility with how SimpleCookie parses cookies, including handling + of malformed cookies with missing semicolons. + + This function is used for both Cookie and Set-Cookie headers in order to be + forgiving. Ideally we would have followed RFC 6265 Section 5.2 (for Cookie + headers) and RFC 6265 Section 4.2.1 (for Set-Cookie headers), but the + real world data makes it impossible since we need to be a bit more forgiving. + + NOTE: This implementation differs from SimpleCookie in handling unmatched quotes. + SimpleCookie will stop parsing when it encounters a cookie value with an unmatched + quote (e.g., 'cookie="value'), causing subsequent cookies to be silently dropped. + This implementation handles unmatched quotes more gracefully to prevent cookie loss. + See https://github.com/aio-libs/aiohttp/issues/7993 + """ + parsed_cookies: List[Tuple[str, Morsel[str]]] = [] + + for header in headers: + if not header: + continue + + # Parse cookie string using SimpleCookie's algorithm + i = 0 + n = len(header) + current_morsel: Optional[Morsel[str]] = None + morsel_seen = False + + while 0 <= i < n: + # Start looking for a cookie + match = _COOKIE_PATTERN.match(header, i) + if not match: + # No more cookies + break + + key, value = match.group("key"), match.group("val") + i = match.end(0) + lower_key = key.lower() + + if key[0] == "$": + if not morsel_seen: + # We ignore attributes which pertain to the cookie + # mechanism as a whole, such as "$Version". + continue + # Process as attribute + if current_morsel is not None: + attr_lower_key = lower_key[1:] + if attr_lower_key in _COOKIE_KNOWN_ATTRS: + current_morsel[attr_lower_key] = value or "" + elif lower_key in _COOKIE_KNOWN_ATTRS: + if not morsel_seen: + # Invalid cookie string - attribute before cookie + break + if lower_key in _COOKIE_BOOL_ATTRS: + # Boolean attribute with any value should be True + if current_morsel is not None and current_morsel.isReservedKey(key): + current_morsel[lower_key] = True + elif value is None: + # Invalid cookie string - non-boolean attribute without value + break + elif current_morsel is not None: + # Regular attribute with value + current_morsel[lower_key] = _unquote(value) + elif value is not None: + # This is a cookie name=value pair + # Validate the name + if key in _COOKIE_KNOWN_ATTRS or not _COOKIE_NAME_RE.match(key): + internal_logger.warning( + "Can not load cookies: Illegal cookie name %r", key + ) + current_morsel = None + else: + # Create new morsel + current_morsel = Morsel() + # Preserve the original value as coded_value (with quotes if present) + # We use __setstate__ instead of the public set() API because it allows us to + # bypass validation and set already validated state. This is more stable than + # setting protected attributes directly and unlikely to change since it would + # break pickling. + current_morsel.__setstate__( # type: ignore[attr-defined] + {"key": key, "value": _unquote(value), "coded_value": value} + ) + parsed_cookies.append((key, current_morsel)) + morsel_seen = True + else: + # Invalid cookie string - no value for non-attribute + break + + return parsed_cookies diff --git a/lib/python3.12/site-packages/aiohttp/_cparser.pxd b/lib/python3.12/site-packages/aiohttp/_cparser.pxd new file mode 100644 index 0000000000000000000000000000000000000000..1b3be6d4efb682bca9397da34f8e727b381bc84f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_cparser.pxd @@ -0,0 +1,158 @@ +from libc.stdint cimport int32_t, uint8_t, uint16_t, uint64_t + + +cdef extern from "llhttp.h": + + struct llhttp__internal_s: + int32_t _index + void* _span_pos0 + void* _span_cb0 + int32_t error + const char* reason + const char* error_pos + void* data + void* _current + uint64_t content_length + uint8_t type + uint8_t method + uint8_t http_major + uint8_t http_minor + uint8_t header_state + uint8_t lenient_flags + uint8_t upgrade + uint8_t finish + uint16_t flags + uint16_t status_code + void* settings + + ctypedef llhttp__internal_s llhttp__internal_t + ctypedef llhttp__internal_t llhttp_t + + ctypedef int (*llhttp_data_cb)(llhttp_t*, const char *at, size_t length) except -1 + ctypedef int (*llhttp_cb)(llhttp_t*) except -1 + + struct llhttp_settings_s: + llhttp_cb on_message_begin + llhttp_data_cb on_url + llhttp_data_cb on_status + llhttp_data_cb on_header_field + llhttp_data_cb on_header_value + llhttp_cb on_headers_complete + llhttp_data_cb on_body + llhttp_cb on_message_complete + llhttp_cb on_chunk_header + llhttp_cb on_chunk_complete + + llhttp_cb on_url_complete + llhttp_cb on_status_complete + llhttp_cb on_header_field_complete + llhttp_cb on_header_value_complete + + ctypedef llhttp_settings_s llhttp_settings_t + + enum llhttp_errno: + HPE_OK, + HPE_INTERNAL, + HPE_STRICT, + HPE_LF_EXPECTED, + HPE_UNEXPECTED_CONTENT_LENGTH, + HPE_CLOSED_CONNECTION, + HPE_INVALID_METHOD, + HPE_INVALID_URL, + HPE_INVALID_CONSTANT, + HPE_INVALID_VERSION, + HPE_INVALID_HEADER_TOKEN, + HPE_INVALID_CONTENT_LENGTH, + HPE_INVALID_CHUNK_SIZE, + HPE_INVALID_STATUS, + HPE_INVALID_EOF_STATE, + HPE_INVALID_TRANSFER_ENCODING, + HPE_CB_MESSAGE_BEGIN, + HPE_CB_HEADERS_COMPLETE, + HPE_CB_MESSAGE_COMPLETE, + HPE_CB_CHUNK_HEADER, + HPE_CB_CHUNK_COMPLETE, + HPE_PAUSED, + HPE_PAUSED_UPGRADE, + HPE_USER + + ctypedef llhttp_errno llhttp_errno_t + + enum llhttp_flags: + F_CHUNKED, + F_CONTENT_LENGTH + + enum llhttp_type: + HTTP_REQUEST, + HTTP_RESPONSE, + HTTP_BOTH + + enum llhttp_method: + HTTP_DELETE, + HTTP_GET, + HTTP_HEAD, + HTTP_POST, + HTTP_PUT, + HTTP_CONNECT, + HTTP_OPTIONS, + HTTP_TRACE, + HTTP_COPY, + HTTP_LOCK, + HTTP_MKCOL, + HTTP_MOVE, + HTTP_PROPFIND, + HTTP_PROPPATCH, + HTTP_SEARCH, + HTTP_UNLOCK, + HTTP_BIND, + HTTP_REBIND, + HTTP_UNBIND, + HTTP_ACL, + HTTP_REPORT, + HTTP_MKACTIVITY, + HTTP_CHECKOUT, + HTTP_MERGE, + HTTP_MSEARCH, + HTTP_NOTIFY, + HTTP_SUBSCRIBE, + HTTP_UNSUBSCRIBE, + HTTP_PATCH, + HTTP_PURGE, + HTTP_MKCALENDAR, + HTTP_LINK, + HTTP_UNLINK, + HTTP_SOURCE, + HTTP_PRI, + HTTP_DESCRIBE, + HTTP_ANNOUNCE, + HTTP_SETUP, + HTTP_PLAY, + HTTP_PAUSE, + HTTP_TEARDOWN, + HTTP_GET_PARAMETER, + HTTP_SET_PARAMETER, + HTTP_REDIRECT, + HTTP_RECORD, + HTTP_FLUSH + + ctypedef llhttp_method llhttp_method_t; + + void llhttp_settings_init(llhttp_settings_t* settings) + void llhttp_init(llhttp_t* parser, llhttp_type type, + const llhttp_settings_t* settings) + + llhttp_errno_t llhttp_execute(llhttp_t* parser, const char* data, size_t len) + + int llhttp_should_keep_alive(const llhttp_t* parser) + + void llhttp_resume_after_upgrade(llhttp_t* parser) + + llhttp_errno_t llhttp_get_errno(const llhttp_t* parser) + const char* llhttp_get_error_reason(const llhttp_t* parser) + const char* llhttp_get_error_pos(const llhttp_t* parser) + + const char* llhttp_method_name(llhttp_method_t method) + + void llhttp_set_lenient_headers(llhttp_t* parser, int enabled) + void llhttp_set_lenient_optional_cr_before_lf(llhttp_t* parser, int enabled) + void llhttp_set_lenient_spaces_after_chunk_size(llhttp_t* parser, int enabled) diff --git a/lib/python3.12/site-packages/aiohttp/_find_header.pxd b/lib/python3.12/site-packages/aiohttp/_find_header.pxd new file mode 100644 index 0000000000000000000000000000000000000000..37a6c37268ee30b182fd77d109688d35d5577c7f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_find_header.pxd @@ -0,0 +1,2 @@ +cdef extern from "_find_header.h": + int find_header(char *, int) diff --git a/lib/python3.12/site-packages/aiohttp/_headers.pxi b/lib/python3.12/site-packages/aiohttp/_headers.pxi new file mode 100644 index 0000000000000000000000000000000000000000..3744721d4786a6c79b90aa349c8d02fa66204ecc --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_headers.pxi @@ -0,0 +1,83 @@ +# The file is autogenerated from aiohttp/hdrs.py +# Run ./tools/gen.py to update it after the origin changing. + +from . import hdrs +cdef tuple headers = ( + hdrs.ACCEPT, + hdrs.ACCEPT_CHARSET, + hdrs.ACCEPT_ENCODING, + hdrs.ACCEPT_LANGUAGE, + hdrs.ACCEPT_RANGES, + hdrs.ACCESS_CONTROL_ALLOW_CREDENTIALS, + hdrs.ACCESS_CONTROL_ALLOW_HEADERS, + hdrs.ACCESS_CONTROL_ALLOW_METHODS, + hdrs.ACCESS_CONTROL_ALLOW_ORIGIN, + hdrs.ACCESS_CONTROL_EXPOSE_HEADERS, + hdrs.ACCESS_CONTROL_MAX_AGE, + hdrs.ACCESS_CONTROL_REQUEST_HEADERS, + hdrs.ACCESS_CONTROL_REQUEST_METHOD, + hdrs.AGE, + hdrs.ALLOW, + hdrs.AUTHORIZATION, + hdrs.CACHE_CONTROL, + hdrs.CONNECTION, + hdrs.CONTENT_DISPOSITION, + hdrs.CONTENT_ENCODING, + hdrs.CONTENT_LANGUAGE, + hdrs.CONTENT_LENGTH, + hdrs.CONTENT_LOCATION, + hdrs.CONTENT_MD5, + hdrs.CONTENT_RANGE, + hdrs.CONTENT_TRANSFER_ENCODING, + hdrs.CONTENT_TYPE, + hdrs.COOKIE, + hdrs.DATE, + hdrs.DESTINATION, + hdrs.DIGEST, + hdrs.ETAG, + hdrs.EXPECT, + hdrs.EXPIRES, + hdrs.FORWARDED, + hdrs.FROM, + hdrs.HOST, + hdrs.IF_MATCH, + hdrs.IF_MODIFIED_SINCE, + hdrs.IF_NONE_MATCH, + hdrs.IF_RANGE, + hdrs.IF_UNMODIFIED_SINCE, + hdrs.KEEP_ALIVE, + hdrs.LAST_EVENT_ID, + hdrs.LAST_MODIFIED, + hdrs.LINK, + hdrs.LOCATION, + hdrs.MAX_FORWARDS, + hdrs.ORIGIN, + hdrs.PRAGMA, + hdrs.PROXY_AUTHENTICATE, + hdrs.PROXY_AUTHORIZATION, + hdrs.RANGE, + hdrs.REFERER, + hdrs.RETRY_AFTER, + hdrs.SEC_WEBSOCKET_ACCEPT, + hdrs.SEC_WEBSOCKET_EXTENSIONS, + hdrs.SEC_WEBSOCKET_KEY, + hdrs.SEC_WEBSOCKET_KEY1, + hdrs.SEC_WEBSOCKET_PROTOCOL, + hdrs.SEC_WEBSOCKET_VERSION, + hdrs.SERVER, + hdrs.SET_COOKIE, + hdrs.TE, + hdrs.TRAILER, + hdrs.TRANSFER_ENCODING, + hdrs.URI, + hdrs.UPGRADE, + hdrs.USER_AGENT, + hdrs.VARY, + hdrs.VIA, + hdrs.WWW_AUTHENTICATE, + hdrs.WANT_DIGEST, + hdrs.WARNING, + hdrs.X_FORWARDED_FOR, + hdrs.X_FORWARDED_HOST, + hdrs.X_FORWARDED_PROTO, +) diff --git a/lib/python3.12/site-packages/aiohttp/_http_parser.pyx b/lib/python3.12/site-packages/aiohttp/_http_parser.pyx new file mode 100644 index 0000000000000000000000000000000000000000..4a7101edbcb1a7f938f8356775bc73ceca8fadf3 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_http_parser.pyx @@ -0,0 +1,835 @@ +# Based on https://github.com/MagicStack/httptools +# + +from cpython cimport ( + Py_buffer, + PyBUF_SIMPLE, + PyBuffer_Release, + PyBytes_AsString, + PyBytes_AsStringAndSize, + PyObject_GetBuffer, +) +from cpython.mem cimport PyMem_Free, PyMem_Malloc +from libc.limits cimport ULLONG_MAX +from libc.string cimport memcpy + +from multidict import CIMultiDict as _CIMultiDict, CIMultiDictProxy as _CIMultiDictProxy +from yarl import URL as _URL + +from aiohttp import hdrs +from aiohttp.helpers import DEBUG, set_exception + +from .http_exceptions import ( + BadHttpMessage, + BadHttpMethod, + BadStatusLine, + ContentLengthError, + InvalidHeader, + InvalidURLError, + LineTooLong, + PayloadEncodingError, + TransferEncodingError, +) +from .http_parser import DeflateBuffer as _DeflateBuffer +from .http_writer import ( + HttpVersion as _HttpVersion, + HttpVersion10 as _HttpVersion10, + HttpVersion11 as _HttpVersion11, +) +from .streams import EMPTY_PAYLOAD as _EMPTY_PAYLOAD, StreamReader as _StreamReader + +cimport cython + +from aiohttp cimport _cparser as cparser + +include "_headers.pxi" + +from aiohttp cimport _find_header + +ALLOWED_UPGRADES = frozenset({"websocket"}) +DEF DEFAULT_FREELIST_SIZE = 250 + +cdef extern from "Python.h": + int PyByteArray_Resize(object, Py_ssize_t) except -1 + Py_ssize_t PyByteArray_Size(object) except -1 + char* PyByteArray_AsString(object) + +__all__ = ('HttpRequestParser', 'HttpResponseParser', + 'RawRequestMessage', 'RawResponseMessage') + +cdef object URL = _URL +cdef object URL_build = URL.build +cdef object CIMultiDict = _CIMultiDict +cdef object CIMultiDictProxy = _CIMultiDictProxy +cdef object HttpVersion = _HttpVersion +cdef object HttpVersion10 = _HttpVersion10 +cdef object HttpVersion11 = _HttpVersion11 +cdef object SEC_WEBSOCKET_KEY1 = hdrs.SEC_WEBSOCKET_KEY1 +cdef object CONTENT_ENCODING = hdrs.CONTENT_ENCODING +cdef object EMPTY_PAYLOAD = _EMPTY_PAYLOAD +cdef object StreamReader = _StreamReader +cdef object DeflateBuffer = _DeflateBuffer +cdef bytes EMPTY_BYTES = b"" + +cdef inline object extend(object buf, const char* at, size_t length): + cdef Py_ssize_t s + cdef char* ptr + s = PyByteArray_Size(buf) + PyByteArray_Resize(buf, s + length) + ptr = PyByteArray_AsString(buf) + memcpy(ptr + s, at, length) + + +DEF METHODS_COUNT = 46; + +cdef list _http_method = [] + +for i in range(METHODS_COUNT): + _http_method.append( + cparser.llhttp_method_name( i).decode('ascii')) + + +cdef inline str http_method_str(int i): + if i < METHODS_COUNT: + return _http_method[i] + else: + return "" + +cdef inline object find_header(bytes raw_header): + cdef Py_ssize_t size + cdef char *buf + cdef int idx + PyBytes_AsStringAndSize(raw_header, &buf, &size) + idx = _find_header.find_header(buf, size) + if idx == -1: + return raw_header.decode('utf-8', 'surrogateescape') + return headers[idx] + + +@cython.freelist(DEFAULT_FREELIST_SIZE) +cdef class RawRequestMessage: + cdef readonly str method + cdef readonly str path + cdef readonly object version # HttpVersion + cdef readonly object headers # CIMultiDict + cdef readonly object raw_headers # tuple + cdef readonly object should_close + cdef readonly object compression + cdef readonly object upgrade + cdef readonly object chunked + cdef readonly object url # yarl.URL + + def __init__(self, method, path, version, headers, raw_headers, + should_close, compression, upgrade, chunked, url): + self.method = method + self.path = path + self.version = version + self.headers = headers + self.raw_headers = raw_headers + self.should_close = should_close + self.compression = compression + self.upgrade = upgrade + self.chunked = chunked + self.url = url + + def __repr__(self): + info = [] + info.append(("method", self.method)) + info.append(("path", self.path)) + info.append(("version", self.version)) + info.append(("headers", self.headers)) + info.append(("raw_headers", self.raw_headers)) + info.append(("should_close", self.should_close)) + info.append(("compression", self.compression)) + info.append(("upgrade", self.upgrade)) + info.append(("chunked", self.chunked)) + info.append(("url", self.url)) + sinfo = ', '.join(name + '=' + repr(val) for name, val in info) + return '' + + def _replace(self, **dct): + cdef RawRequestMessage ret + ret = _new_request_message(self.method, + self.path, + self.version, + self.headers, + self.raw_headers, + self.should_close, + self.compression, + self.upgrade, + self.chunked, + self.url) + if "method" in dct: + ret.method = dct["method"] + if "path" in dct: + ret.path = dct["path"] + if "version" in dct: + ret.version = dct["version"] + if "headers" in dct: + ret.headers = dct["headers"] + if "raw_headers" in dct: + ret.raw_headers = dct["raw_headers"] + if "should_close" in dct: + ret.should_close = dct["should_close"] + if "compression" in dct: + ret.compression = dct["compression"] + if "upgrade" in dct: + ret.upgrade = dct["upgrade"] + if "chunked" in dct: + ret.chunked = dct["chunked"] + if "url" in dct: + ret.url = dct["url"] + return ret + +cdef _new_request_message(str method, + str path, + object version, + object headers, + object raw_headers, + bint should_close, + object compression, + bint upgrade, + bint chunked, + object url): + cdef RawRequestMessage ret + ret = RawRequestMessage.__new__(RawRequestMessage) + ret.method = method + ret.path = path + ret.version = version + ret.headers = headers + ret.raw_headers = raw_headers + ret.should_close = should_close + ret.compression = compression + ret.upgrade = upgrade + ret.chunked = chunked + ret.url = url + return ret + + +@cython.freelist(DEFAULT_FREELIST_SIZE) +cdef class RawResponseMessage: + cdef readonly object version # HttpVersion + cdef readonly int code + cdef readonly str reason + cdef readonly object headers # CIMultiDict + cdef readonly object raw_headers # tuple + cdef readonly object should_close + cdef readonly object compression + cdef readonly object upgrade + cdef readonly object chunked + + def __init__(self, version, code, reason, headers, raw_headers, + should_close, compression, upgrade, chunked): + self.version = version + self.code = code + self.reason = reason + self.headers = headers + self.raw_headers = raw_headers + self.should_close = should_close + self.compression = compression + self.upgrade = upgrade + self.chunked = chunked + + def __repr__(self): + info = [] + info.append(("version", self.version)) + info.append(("code", self.code)) + info.append(("reason", self.reason)) + info.append(("headers", self.headers)) + info.append(("raw_headers", self.raw_headers)) + info.append(("should_close", self.should_close)) + info.append(("compression", self.compression)) + info.append(("upgrade", self.upgrade)) + info.append(("chunked", self.chunked)) + sinfo = ', '.join(name + '=' + repr(val) for name, val in info) + return '' + + +cdef _new_response_message(object version, + int code, + str reason, + object headers, + object raw_headers, + bint should_close, + object compression, + bint upgrade, + bint chunked): + cdef RawResponseMessage ret + ret = RawResponseMessage.__new__(RawResponseMessage) + ret.version = version + ret.code = code + ret.reason = reason + ret.headers = headers + ret.raw_headers = raw_headers + ret.should_close = should_close + ret.compression = compression + ret.upgrade = upgrade + ret.chunked = chunked + return ret + + +@cython.internal +cdef class HttpParser: + + cdef: + cparser.llhttp_t* _cparser + cparser.llhttp_settings_t* _csettings + + bytes _raw_name + object _name + bytes _raw_value + bint _has_value + + object _protocol + object _loop + object _timer + + size_t _max_line_size + size_t _max_field_size + size_t _max_headers + bint _response_with_body + bint _read_until_eof + + bint _started + object _url + bytearray _buf + str _path + str _reason + list _headers + list _raw_headers + bint _upgraded + list _messages + object _payload + bint _payload_error + object _payload_exception + object _last_error + bint _auto_decompress + int _limit + + str _content_encoding + + Py_buffer py_buf + + def __cinit__(self): + self._cparser = \ + PyMem_Malloc(sizeof(cparser.llhttp_t)) + if self._cparser is NULL: + raise MemoryError() + + self._csettings = \ + PyMem_Malloc(sizeof(cparser.llhttp_settings_t)) + if self._csettings is NULL: + raise MemoryError() + + def __dealloc__(self): + PyMem_Free(self._cparser) + PyMem_Free(self._csettings) + + cdef _init( + self, cparser.llhttp_type mode, + object protocol, object loop, int limit, + object timer=None, + size_t max_line_size=8190, size_t max_headers=32768, + size_t max_field_size=8190, payload_exception=None, + bint response_with_body=True, bint read_until_eof=False, + bint auto_decompress=True, + ): + cparser.llhttp_settings_init(self._csettings) + cparser.llhttp_init(self._cparser, mode, self._csettings) + self._cparser.data = self + self._cparser.content_length = 0 + + self._protocol = protocol + self._loop = loop + self._timer = timer + + self._buf = bytearray() + self._payload = None + self._payload_error = 0 + self._payload_exception = payload_exception + self._messages = [] + + self._raw_name = EMPTY_BYTES + self._raw_value = EMPTY_BYTES + self._has_value = False + + self._max_line_size = max_line_size + self._max_headers = max_headers + self._max_field_size = max_field_size + self._response_with_body = response_with_body + self._read_until_eof = read_until_eof + self._upgraded = False + self._auto_decompress = auto_decompress + self._content_encoding = None + + self._csettings.on_url = cb_on_url + self._csettings.on_status = cb_on_status + self._csettings.on_header_field = cb_on_header_field + self._csettings.on_header_value = cb_on_header_value + self._csettings.on_headers_complete = cb_on_headers_complete + self._csettings.on_body = cb_on_body + self._csettings.on_message_begin = cb_on_message_begin + self._csettings.on_message_complete = cb_on_message_complete + self._csettings.on_chunk_header = cb_on_chunk_header + self._csettings.on_chunk_complete = cb_on_chunk_complete + + self._last_error = None + self._limit = limit + + cdef _process_header(self): + cdef str value + if self._raw_name is not EMPTY_BYTES: + name = find_header(self._raw_name) + value = self._raw_value.decode('utf-8', 'surrogateescape') + + self._headers.append((name, value)) + + if name is CONTENT_ENCODING: + self._content_encoding = value + + self._has_value = False + self._raw_headers.append((self._raw_name, self._raw_value)) + self._raw_name = EMPTY_BYTES + self._raw_value = EMPTY_BYTES + + cdef _on_header_field(self, char* at, size_t length): + if self._has_value: + self._process_header() + + if self._raw_name is EMPTY_BYTES: + self._raw_name = at[:length] + else: + self._raw_name += at[:length] + + cdef _on_header_value(self, char* at, size_t length): + if self._raw_value is EMPTY_BYTES: + self._raw_value = at[:length] + else: + self._raw_value += at[:length] + self._has_value = True + + cdef _on_headers_complete(self): + self._process_header() + + should_close = not cparser.llhttp_should_keep_alive(self._cparser) + upgrade = self._cparser.upgrade + chunked = self._cparser.flags & cparser.F_CHUNKED + + raw_headers = tuple(self._raw_headers) + headers = CIMultiDictProxy(CIMultiDict(self._headers)) + + if self._cparser.type == cparser.HTTP_REQUEST: + h_upg = headers.get("upgrade", "") + allowed = upgrade and h_upg.isascii() and h_upg.lower() in ALLOWED_UPGRADES + if allowed or self._cparser.method == cparser.HTTP_CONNECT: + self._upgraded = True + else: + if upgrade and self._cparser.status_code == 101: + self._upgraded = True + + # do not support old websocket spec + if SEC_WEBSOCKET_KEY1 in headers: + raise InvalidHeader(SEC_WEBSOCKET_KEY1) + + encoding = None + enc = self._content_encoding + if enc is not None: + self._content_encoding = None + if enc.isascii() and enc.lower() in {"gzip", "deflate", "br", "zstd"}: + encoding = enc + + if self._cparser.type == cparser.HTTP_REQUEST: + method = http_method_str(self._cparser.method) + msg = _new_request_message( + method, self._path, + self.http_version(), headers, raw_headers, + should_close, encoding, upgrade, chunked, self._url) + else: + msg = _new_response_message( + self.http_version(), self._cparser.status_code, self._reason, + headers, raw_headers, should_close, encoding, + upgrade, chunked) + + if ( + ULLONG_MAX > self._cparser.content_length > 0 or chunked or + self._cparser.method == cparser.HTTP_CONNECT or + (self._cparser.status_code >= 199 and + self._cparser.content_length == 0 and + self._read_until_eof) + ): + payload = StreamReader( + self._protocol, timer=self._timer, loop=self._loop, + limit=self._limit) + else: + payload = EMPTY_PAYLOAD + + self._payload = payload + if encoding is not None and self._auto_decompress: + self._payload = DeflateBuffer(payload, encoding) + + if not self._response_with_body: + payload = EMPTY_PAYLOAD + + self._messages.append((msg, payload)) + + cdef _on_message_complete(self): + self._payload.feed_eof() + self._payload = None + + cdef _on_chunk_header(self): + self._payload.begin_http_chunk_receiving() + + cdef _on_chunk_complete(self): + self._payload.end_http_chunk_receiving() + + cdef object _on_status_complete(self): + pass + + cdef inline http_version(self): + cdef cparser.llhttp_t* parser = self._cparser + + if parser.http_major == 1: + if parser.http_minor == 0: + return HttpVersion10 + elif parser.http_minor == 1: + return HttpVersion11 + + return HttpVersion(parser.http_major, parser.http_minor) + + ### Public API ### + + def feed_eof(self): + cdef bytes desc + + if self._payload is not None: + if self._cparser.flags & cparser.F_CHUNKED: + raise TransferEncodingError( + "Not enough data to satisfy transfer length header.") + elif self._cparser.flags & cparser.F_CONTENT_LENGTH: + raise ContentLengthError( + "Not enough data to satisfy content length header.") + elif cparser.llhttp_get_errno(self._cparser) != cparser.HPE_OK: + desc = cparser.llhttp_get_error_reason(self._cparser) + raise PayloadEncodingError(desc.decode('latin-1')) + else: + self._payload.feed_eof() + elif self._started: + self._on_headers_complete() + if self._messages: + return self._messages[-1][0] + + def feed_data(self, data): + cdef: + size_t data_len + size_t nb + cdef cparser.llhttp_errno_t errno + + PyObject_GetBuffer(data, &self.py_buf, PyBUF_SIMPLE) + data_len = self.py_buf.len + + errno = cparser.llhttp_execute( + self._cparser, + self.py_buf.buf, + data_len) + + if errno is cparser.HPE_PAUSED_UPGRADE: + cparser.llhttp_resume_after_upgrade(self._cparser) + + nb = cparser.llhttp_get_error_pos(self._cparser) - self.py_buf.buf + + PyBuffer_Release(&self.py_buf) + + if errno not in (cparser.HPE_OK, cparser.HPE_PAUSED_UPGRADE): + if self._payload_error == 0: + if self._last_error is not None: + ex = self._last_error + self._last_error = None + else: + after = cparser.llhttp_get_error_pos(self._cparser) + before = data[:after - self.py_buf.buf] + after_b = after.split(b"\r\n", 1)[0] + before = before.rsplit(b"\r\n", 1)[-1] + data = before + after_b + pointer = " " * (len(repr(before))-1) + "^" + ex = parser_error_from_errno(self._cparser, data, pointer) + self._payload = None + raise ex + + if self._messages: + messages = self._messages + self._messages = [] + else: + messages = () + + if self._upgraded: + return messages, True, data[nb:] + else: + return messages, False, b"" + + def set_upgraded(self, val): + self._upgraded = val + + +cdef class HttpRequestParser(HttpParser): + + def __init__( + self, protocol, loop, int limit, timer=None, + size_t max_line_size=8190, size_t max_headers=32768, + size_t max_field_size=8190, payload_exception=None, + bint response_with_body=True, bint read_until_eof=False, + bint auto_decompress=True, + ): + self._init(cparser.HTTP_REQUEST, protocol, loop, limit, timer, + max_line_size, max_headers, max_field_size, + payload_exception, response_with_body, read_until_eof, + auto_decompress) + + cdef object _on_status_complete(self): + cdef int idx1, idx2 + if not self._buf: + return + self._path = self._buf.decode('utf-8', 'surrogateescape') + try: + idx3 = len(self._path) + if self._cparser.method == cparser.HTTP_CONNECT: + # authority-form, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.3 + self._url = URL.build(authority=self._path, encoded=True) + elif idx3 > 1 and self._path[0] == '/': + # origin-form, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.1 + idx1 = self._path.find("?") + if idx1 == -1: + query = "" + idx2 = self._path.find("#") + if idx2 == -1: + path = self._path + fragment = "" + else: + path = self._path[0: idx2] + fragment = self._path[idx2+1:] + + else: + path = self._path[0:idx1] + idx1 += 1 + idx2 = self._path.find("#", idx1+1) + if idx2 == -1: + query = self._path[idx1:] + fragment = "" + else: + query = self._path[idx1: idx2] + fragment = self._path[idx2+1:] + + self._url = URL.build( + path=path, + query_string=query, + fragment=fragment, + encoded=True, + ) + else: + # absolute-form for proxy maybe, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.2 + self._url = URL(self._path, encoded=True) + finally: + PyByteArray_Resize(self._buf, 0) + + +cdef class HttpResponseParser(HttpParser): + + def __init__( + self, protocol, loop, int limit, timer=None, + size_t max_line_size=8190, size_t max_headers=32768, + size_t max_field_size=8190, payload_exception=None, + bint response_with_body=True, bint read_until_eof=False, + bint auto_decompress=True + ): + self._init(cparser.HTTP_RESPONSE, protocol, loop, limit, timer, + max_line_size, max_headers, max_field_size, + payload_exception, response_with_body, read_until_eof, + auto_decompress) + # Use strict parsing on dev mode, so users are warned about broken servers. + if not DEBUG: + cparser.llhttp_set_lenient_headers(self._cparser, 1) + cparser.llhttp_set_lenient_optional_cr_before_lf(self._cparser, 1) + cparser.llhttp_set_lenient_spaces_after_chunk_size(self._cparser, 1) + + cdef object _on_status_complete(self): + if self._buf: + self._reason = self._buf.decode('utf-8', 'surrogateescape') + PyByteArray_Resize(self._buf, 0) + else: + self._reason = self._reason or '' + +cdef int cb_on_message_begin(cparser.llhttp_t* parser) except -1: + cdef HttpParser pyparser = parser.data + + pyparser._started = True + pyparser._headers = [] + pyparser._raw_headers = [] + PyByteArray_Resize(pyparser._buf, 0) + pyparser._path = None + pyparser._reason = None + return 0 + + +cdef int cb_on_url(cparser.llhttp_t* parser, + const char *at, size_t length) except -1: + cdef HttpParser pyparser = parser.data + try: + if length > pyparser._max_line_size: + raise LineTooLong( + 'Status line is too long', pyparser._max_line_size, length) + extend(pyparser._buf, at, length) + except BaseException as ex: + pyparser._last_error = ex + return -1 + else: + return 0 + + +cdef int cb_on_status(cparser.llhttp_t* parser, + const char *at, size_t length) except -1: + cdef HttpParser pyparser = parser.data + cdef str reason + try: + if length > pyparser._max_line_size: + raise LineTooLong( + 'Status line is too long', pyparser._max_line_size, length) + extend(pyparser._buf, at, length) + except BaseException as ex: + pyparser._last_error = ex + return -1 + else: + return 0 + + +cdef int cb_on_header_field(cparser.llhttp_t* parser, + const char *at, size_t length) except -1: + cdef HttpParser pyparser = parser.data + cdef Py_ssize_t size + try: + pyparser._on_status_complete() + size = len(pyparser._raw_name) + length + if size > pyparser._max_field_size: + raise LineTooLong( + 'Header name is too long', pyparser._max_field_size, size) + pyparser._on_header_field(at, length) + except BaseException as ex: + pyparser._last_error = ex + return -1 + else: + return 0 + + +cdef int cb_on_header_value(cparser.llhttp_t* parser, + const char *at, size_t length) except -1: + cdef HttpParser pyparser = parser.data + cdef Py_ssize_t size + try: + size = len(pyparser._raw_value) + length + if size > pyparser._max_field_size: + raise LineTooLong( + 'Header value is too long', pyparser._max_field_size, size) + pyparser._on_header_value(at, length) + except BaseException as ex: + pyparser._last_error = ex + return -1 + else: + return 0 + + +cdef int cb_on_headers_complete(cparser.llhttp_t* parser) except -1: + cdef HttpParser pyparser = parser.data + try: + pyparser._on_status_complete() + pyparser._on_headers_complete() + except BaseException as exc: + pyparser._last_error = exc + return -1 + else: + if pyparser._upgraded or pyparser._cparser.method == cparser.HTTP_CONNECT: + return 2 + else: + return 0 + + +cdef int cb_on_body(cparser.llhttp_t* parser, + const char *at, size_t length) except -1: + cdef HttpParser pyparser = parser.data + cdef bytes body = at[:length] + try: + pyparser._payload.feed_data(body, length) + except BaseException as underlying_exc: + reraised_exc = underlying_exc + if pyparser._payload_exception is not None: + reraised_exc = pyparser._payload_exception(str(underlying_exc)) + + set_exception(pyparser._payload, reraised_exc, underlying_exc) + + pyparser._payload_error = 1 + return -1 + else: + return 0 + + +cdef int cb_on_message_complete(cparser.llhttp_t* parser) except -1: + cdef HttpParser pyparser = parser.data + try: + pyparser._started = False + pyparser._on_message_complete() + except BaseException as exc: + pyparser._last_error = exc + return -1 + else: + return 0 + + +cdef int cb_on_chunk_header(cparser.llhttp_t* parser) except -1: + cdef HttpParser pyparser = parser.data + try: + pyparser._on_chunk_header() + except BaseException as exc: + pyparser._last_error = exc + return -1 + else: + return 0 + + +cdef int cb_on_chunk_complete(cparser.llhttp_t* parser) except -1: + cdef HttpParser pyparser = parser.data + try: + pyparser._on_chunk_complete() + except BaseException as exc: + pyparser._last_error = exc + return -1 + else: + return 0 + + +cdef parser_error_from_errno(cparser.llhttp_t* parser, data, pointer): + cdef cparser.llhttp_errno_t errno = cparser.llhttp_get_errno(parser) + cdef bytes desc = cparser.llhttp_get_error_reason(parser) + + err_msg = "{}:\n\n {!r}\n {}".format(desc.decode("latin-1"), data, pointer) + + if errno in {cparser.HPE_CB_MESSAGE_BEGIN, + cparser.HPE_CB_HEADERS_COMPLETE, + cparser.HPE_CB_MESSAGE_COMPLETE, + cparser.HPE_CB_CHUNK_HEADER, + cparser.HPE_CB_CHUNK_COMPLETE, + cparser.HPE_INVALID_CONSTANT, + cparser.HPE_INVALID_HEADER_TOKEN, + cparser.HPE_INVALID_CONTENT_LENGTH, + cparser.HPE_INVALID_CHUNK_SIZE, + cparser.HPE_INVALID_EOF_STATE, + cparser.HPE_INVALID_TRANSFER_ENCODING}: + return BadHttpMessage(err_msg) + elif errno == cparser.HPE_INVALID_METHOD: + return BadHttpMethod(error=err_msg) + elif errno in {cparser.HPE_INVALID_STATUS, + cparser.HPE_INVALID_VERSION}: + return BadStatusLine(error=err_msg) + elif errno == cparser.HPE_INVALID_URL: + return InvalidURLError(err_msg) + + return BadHttpMessage(err_msg) diff --git a/lib/python3.12/site-packages/aiohttp/_http_writer.pyx b/lib/python3.12/site-packages/aiohttp/_http_writer.pyx new file mode 100644 index 0000000000000000000000000000000000000000..7989c186c89296e7fa9e336bdad447d27e6f3416 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_http_writer.pyx @@ -0,0 +1,162 @@ +from cpython.bytes cimport PyBytes_FromStringAndSize +from cpython.exc cimport PyErr_NoMemory +from cpython.mem cimport PyMem_Free, PyMem_Malloc, PyMem_Realloc +from cpython.object cimport PyObject_Str +from libc.stdint cimport uint8_t, uint64_t +from libc.string cimport memcpy + +from multidict import istr + +DEF BUF_SIZE = 16 * 1024 # 16KiB + +cdef object _istr = istr + + +# ----------------- writer --------------------------- + +cdef struct Writer: + char *buf + Py_ssize_t size + Py_ssize_t pos + bint heap_allocated + +cdef inline void _init_writer(Writer* writer, char *buf): + writer.buf = buf + writer.size = BUF_SIZE + writer.pos = 0 + writer.heap_allocated = 0 + + +cdef inline void _release_writer(Writer* writer): + if writer.heap_allocated: + PyMem_Free(writer.buf) + + +cdef inline int _write_byte(Writer* writer, uint8_t ch): + cdef char * buf + cdef Py_ssize_t size + + if writer.pos == writer.size: + # reallocate + size = writer.size + BUF_SIZE + if not writer.heap_allocated: + buf = PyMem_Malloc(size) + if buf == NULL: + PyErr_NoMemory() + return -1 + memcpy(buf, writer.buf, writer.size) + else: + buf = PyMem_Realloc(writer.buf, size) + if buf == NULL: + PyErr_NoMemory() + return -1 + writer.buf = buf + writer.size = size + writer.heap_allocated = 1 + writer.buf[writer.pos] = ch + writer.pos += 1 + return 0 + + +cdef inline int _write_utf8(Writer* writer, Py_UCS4 symbol): + cdef uint64_t utf = symbol + + if utf < 0x80: + return _write_byte(writer, utf) + elif utf < 0x800: + if _write_byte(writer, (0xc0 | (utf >> 6))) < 0: + return -1 + return _write_byte(writer, (0x80 | (utf & 0x3f))) + elif 0xD800 <= utf <= 0xDFFF: + # surogate pair, ignored + return 0 + elif utf < 0x10000: + if _write_byte(writer, (0xe0 | (utf >> 12))) < 0: + return -1 + if _write_byte(writer, (0x80 | ((utf >> 6) & 0x3f))) < 0: + return -1 + return _write_byte(writer, (0x80 | (utf & 0x3f))) + elif utf > 0x10FFFF: + # symbol is too large + return 0 + else: + if _write_byte(writer, (0xf0 | (utf >> 18))) < 0: + return -1 + if _write_byte(writer, + (0x80 | ((utf >> 12) & 0x3f))) < 0: + return -1 + if _write_byte(writer, + (0x80 | ((utf >> 6) & 0x3f))) < 0: + return -1 + return _write_byte(writer, (0x80 | (utf & 0x3f))) + + +cdef inline int _write_str(Writer* writer, str s): + cdef Py_UCS4 ch + for ch in s: + if _write_utf8(writer, ch) < 0: + return -1 + + +cdef inline int _write_str_raise_on_nlcr(Writer* writer, object s): + cdef Py_UCS4 ch + cdef str out_str + if type(s) is str: + out_str = s + elif type(s) is _istr: + out_str = PyObject_Str(s) + elif not isinstance(s, str): + raise TypeError("Cannot serialize non-str key {!r}".format(s)) + else: + out_str = str(s) + + for ch in out_str: + if ch == 0x0D or ch == 0x0A: + raise ValueError( + "Newline or carriage return detected in headers. " + "Potential header injection attack." + ) + if _write_utf8(writer, ch) < 0: + return -1 + + +# --------------- _serialize_headers ---------------------- + +def _serialize_headers(str status_line, headers): + cdef Writer writer + cdef object key + cdef object val + cdef char buf[BUF_SIZE] + + _init_writer(&writer, buf) + + try: + if _write_str(&writer, status_line) < 0: + raise + if _write_byte(&writer, b'\r') < 0: + raise + if _write_byte(&writer, b'\n') < 0: + raise + + for key, val in headers.items(): + if _write_str_raise_on_nlcr(&writer, key) < 0: + raise + if _write_byte(&writer, b':') < 0: + raise + if _write_byte(&writer, b' ') < 0: + raise + if _write_str_raise_on_nlcr(&writer, val) < 0: + raise + if _write_byte(&writer, b'\r') < 0: + raise + if _write_byte(&writer, b'\n') < 0: + raise + + if _write_byte(&writer, b'\r') < 0: + raise + if _write_byte(&writer, b'\n') < 0: + raise + + return PyBytes_FromStringAndSize(writer.buf, writer.pos) + finally: + _release_writer(&writer) diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pxd.hash b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pxd.hash new file mode 100644 index 0000000000000000000000000000000000000000..eadfed3d48b4781b271296d8864672a9b46b73c8 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pxd.hash @@ -0,0 +1 @@ +b01999d409b29bd916e067bc963d5f2d9ee63cfc9ae0bccb769910131417bf93 /home/runner/work/aiohttp/aiohttp/aiohttp/_websocket/mask.pxd diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pyx.hash b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pyx.hash new file mode 100644 index 0000000000000000000000000000000000000000..5cd7ae67746cbdcc347183a9fdb1aa24a6f89063 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/mask.pyx.hash @@ -0,0 +1 @@ +0478ceb55d0ed30ef1a7da742cd003449bc69a07cf9fdb06789bd2b347cbfffe /home/runner/work/aiohttp/aiohttp/aiohttp/_websocket/mask.pyx diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/.hash/reader_c.pxd.hash b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/reader_c.pxd.hash new file mode 100644 index 0000000000000000000000000000000000000000..f2feb6aad915eb2dfbe1803aec0f0095f67fb5d6 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/.hash/reader_c.pxd.hash @@ -0,0 +1 @@ +9e5fe78ed0ebce5414d2b8e01868d90c1facc20b84d2d5ff6c23e86e44a155ae /home/runner/work/aiohttp/aiohttp/aiohttp/_websocket/reader_c.pxd diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__init__.py b/lib/python3.12/site-packages/aiohttp/_websocket/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..836257cc47aba3e74863c7de0e098d0835bcee1f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/__init__.py @@ -0,0 +1 @@ +"""WebSocket protocol versions 13 and 8.""" diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b392b7d07209e4d44987d7eb2a5ce871ec2f3dc Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/helpers.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..384e449a061e386c9bfd1137ee835cf8a65b9130 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/models.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/models.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a98aa30cedf8d9397c18f661fd1f00830d7dd683 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/models.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a06f4b1cf0ee2d9bff9d9ec730925d21fe6277a1 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_c.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_c.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca3acaf0d79ed39f1348ce9c05f9b67718e12f04 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_c.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_py.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_py.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37c4dac7a979ace895c19405317ef4e7afb2b250 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/reader_py.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/writer.cpython-312.pyc b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/writer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67da637b111a1d402e1b5365227197aaad70b378 Binary files /dev/null and b/lib/python3.12/site-packages/aiohttp/_websocket/__pycache__/writer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/helpers.py b/lib/python3.12/site-packages/aiohttp/_websocket/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..0bb58df9228603fc1eb79a7a2cac8301217a36a1 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/helpers.py @@ -0,0 +1,147 @@ +"""Helpers for WebSocket protocol versions 13 and 8.""" + +import functools +import re +from struct import Struct +from typing import TYPE_CHECKING, Final, List, Optional, Pattern, Tuple + +from ..helpers import NO_EXTENSIONS +from .models import WSHandshakeError + +UNPACK_LEN3 = Struct("!Q").unpack_from +UNPACK_CLOSE_CODE = Struct("!H").unpack +PACK_LEN1 = Struct("!BB").pack +PACK_LEN2 = Struct("!BBH").pack +PACK_LEN3 = Struct("!BBQ").pack +PACK_CLOSE_CODE = Struct("!H").pack +PACK_RANDBITS = Struct("!L").pack +MSG_SIZE: Final[int] = 2**14 +MASK_LEN: Final[int] = 4 + +WS_KEY: Final[bytes] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11" + + +# Used by _websocket_mask_python +@functools.lru_cache +def _xor_table() -> List[bytes]: + return [bytes(a ^ b for a in range(256)) for b in range(256)] + + +def _websocket_mask_python(mask: bytes, data: bytearray) -> None: + """Websocket masking function. + + `mask` is a `bytes` object of length 4; `data` is a `bytearray` + object of any length. The contents of `data` are masked with `mask`, + as specified in section 5.3 of RFC 6455. + + Note that this function mutates the `data` argument. + + This pure-python implementation may be replaced by an optimized + version when available. + + """ + assert isinstance(data, bytearray), data + assert len(mask) == 4, mask + + if data: + _XOR_TABLE = _xor_table() + a, b, c, d = (_XOR_TABLE[n] for n in mask) + data[::4] = data[::4].translate(a) + data[1::4] = data[1::4].translate(b) + data[2::4] = data[2::4].translate(c) + data[3::4] = data[3::4].translate(d) + + +if TYPE_CHECKING or NO_EXTENSIONS: # pragma: no cover + websocket_mask = _websocket_mask_python +else: + try: + from .mask import _websocket_mask_cython # type: ignore[import-not-found] + + websocket_mask = _websocket_mask_cython + except ImportError: # pragma: no cover + websocket_mask = _websocket_mask_python + + +_WS_EXT_RE: Final[Pattern[str]] = re.compile( + r"^(?:;\s*(?:" + r"(server_no_context_takeover)|" + r"(client_no_context_takeover)|" + r"(server_max_window_bits(?:=(\d+))?)|" + r"(client_max_window_bits(?:=(\d+))?)))*$" +) + +_WS_EXT_RE_SPLIT: Final[Pattern[str]] = re.compile(r"permessage-deflate([^,]+)?") + + +def ws_ext_parse(extstr: Optional[str], isserver: bool = False) -> Tuple[int, bool]: + if not extstr: + return 0, False + + compress = 0 + notakeover = False + for ext in _WS_EXT_RE_SPLIT.finditer(extstr): + defext = ext.group(1) + # Return compress = 15 when get `permessage-deflate` + if not defext: + compress = 15 + break + match = _WS_EXT_RE.match(defext) + if match: + compress = 15 + if isserver: + # Server never fail to detect compress handshake. + # Server does not need to send max wbit to client + if match.group(4): + compress = int(match.group(4)) + # Group3 must match if group4 matches + # Compress wbit 8 does not support in zlib + # If compress level not support, + # CONTINUE to next extension + if compress > 15 or compress < 9: + compress = 0 + continue + if match.group(1): + notakeover = True + # Ignore regex group 5 & 6 for client_max_window_bits + break + else: + if match.group(6): + compress = int(match.group(6)) + # Group5 must match if group6 matches + # Compress wbit 8 does not support in zlib + # If compress level not support, + # FAIL the parse progress + if compress > 15 or compress < 9: + raise WSHandshakeError("Invalid window size") + if match.group(2): + notakeover = True + # Ignore regex group 5 & 6 for client_max_window_bits + break + # Return Fail if client side and not match + elif not isserver: + raise WSHandshakeError("Extension for deflate not supported" + ext.group(1)) + + return compress, notakeover + + +def ws_ext_gen( + compress: int = 15, isserver: bool = False, server_notakeover: bool = False +) -> str: + # client_notakeover=False not used for server + # compress wbit 8 does not support in zlib + if compress < 9 or compress > 15: + raise ValueError( + "Compress wbits must between 9 and 15, zlib does not support wbits=8" + ) + enabledext = ["permessage-deflate"] + if not isserver: + enabledext.append("client_max_window_bits") + + if compress < 15: + enabledext.append("server_max_window_bits=" + str(compress)) + if server_notakeover: + enabledext.append("server_no_context_takeover") + # if client_notakeover: + # enabledext.append('client_no_context_takeover') + return "; ".join(enabledext) diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/mask.pxd b/lib/python3.12/site-packages/aiohttp/_websocket/mask.pxd new file mode 100644 index 0000000000000000000000000000000000000000..90983de9ac7e59dfceb639c1e7b656abd5fbb305 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/mask.pxd @@ -0,0 +1,3 @@ +"""Cython declarations for websocket masking.""" + +cpdef void _websocket_mask_cython(bytes mask, bytearray data) diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/mask.pyx b/lib/python3.12/site-packages/aiohttp/_websocket/mask.pyx new file mode 100644 index 0000000000000000000000000000000000000000..2d956c8899644d4c6bce042b928be1f23e51293a --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/mask.pyx @@ -0,0 +1,48 @@ +from cpython cimport PyBytes_AsString + + +#from cpython cimport PyByteArray_AsString # cython still not exports that +cdef extern from "Python.h": + char* PyByteArray_AsString(bytearray ba) except NULL + +from libc.stdint cimport uint32_t, uint64_t, uintmax_t + + +cpdef void _websocket_mask_cython(bytes mask, bytearray data): + """Note, this function mutates its `data` argument + """ + cdef: + Py_ssize_t data_len, i + # bit operations on signed integers are implementation-specific + unsigned char * in_buf + const unsigned char * mask_buf + uint32_t uint32_msk + uint64_t uint64_msk + + assert len(mask) == 4 + + data_len = len(data) + in_buf = PyByteArray_AsString(data) + mask_buf = PyBytes_AsString(mask) + uint32_msk = (mask_buf)[0] + + # TODO: align in_data ptr to achieve even faster speeds + # does it need in python ?! malloc() always aligns to sizeof(long) bytes + + if sizeof(size_t) >= 8: + uint64_msk = uint32_msk + uint64_msk = (uint64_msk << 32) | uint32_msk + + while data_len >= 8: + (in_buf)[0] ^= uint64_msk + in_buf += 8 + data_len -= 8 + + + while data_len >= 4: + (in_buf)[0] ^= uint32_msk + in_buf += 4 + data_len -= 4 + + for i in range(0, data_len): + in_buf[i] ^= mask_buf[i] diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/models.py b/lib/python3.12/site-packages/aiohttp/_websocket/models.py new file mode 100644 index 0000000000000000000000000000000000000000..7e89b9652957e8f4e73916e18048368e8d75911e --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/models.py @@ -0,0 +1,84 @@ +"""Models for WebSocket protocol versions 13 and 8.""" + +import json +from enum import IntEnum +from typing import Any, Callable, Final, NamedTuple, Optional, cast + +WS_DEFLATE_TRAILING: Final[bytes] = bytes([0x00, 0x00, 0xFF, 0xFF]) + + +class WSCloseCode(IntEnum): + OK = 1000 + GOING_AWAY = 1001 + PROTOCOL_ERROR = 1002 + UNSUPPORTED_DATA = 1003 + ABNORMAL_CLOSURE = 1006 + INVALID_TEXT = 1007 + POLICY_VIOLATION = 1008 + MESSAGE_TOO_BIG = 1009 + MANDATORY_EXTENSION = 1010 + INTERNAL_ERROR = 1011 + SERVICE_RESTART = 1012 + TRY_AGAIN_LATER = 1013 + BAD_GATEWAY = 1014 + + +class WSMsgType(IntEnum): + # websocket spec types + CONTINUATION = 0x0 + TEXT = 0x1 + BINARY = 0x2 + PING = 0x9 + PONG = 0xA + CLOSE = 0x8 + + # aiohttp specific types + CLOSING = 0x100 + CLOSED = 0x101 + ERROR = 0x102 + + text = TEXT + binary = BINARY + ping = PING + pong = PONG + close = CLOSE + closing = CLOSING + closed = CLOSED + error = ERROR + + +class WSMessage(NamedTuple): + type: WSMsgType + # To type correctly, this would need some kind of tagged union for each type. + data: Any + extra: Optional[str] + + def json(self, *, loads: Callable[[Any], Any] = json.loads) -> Any: + """Return parsed JSON data. + + .. versionadded:: 0.22 + """ + return loads(self.data) + + +# Constructing the tuple directly to avoid the overhead of +# the lambda and arg processing since NamedTuples are constructed +# with a run time built lambda +# https://github.com/python/cpython/blob/d83fcf8371f2f33c7797bc8f5423a8bca8c46e5c/Lib/collections/__init__.py#L441 +WS_CLOSED_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSED, None, None)) +WS_CLOSING_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSING, None, None)) + + +class WebSocketError(Exception): + """WebSocket protocol parser error.""" + + def __init__(self, code: int, message: str) -> None: + self.code = code + super().__init__(code, message) + + def __str__(self) -> str: + return cast(str, self.args[1]) + + +class WSHandshakeError(Exception): + """WebSocket protocol handshake error.""" diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/reader.py b/lib/python3.12/site-packages/aiohttp/_websocket/reader.py new file mode 100644 index 0000000000000000000000000000000000000000..23f32265cfccbdc8c8fe2f1600accbfb6f816efa --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/reader.py @@ -0,0 +1,31 @@ +"""Reader for WebSocket protocol versions 13 and 8.""" + +from typing import TYPE_CHECKING + +from ..helpers import NO_EXTENSIONS + +if TYPE_CHECKING or NO_EXTENSIONS: # pragma: no cover + from .reader_py import ( + WebSocketDataQueue as WebSocketDataQueuePython, + WebSocketReader as WebSocketReaderPython, + ) + + WebSocketReader = WebSocketReaderPython + WebSocketDataQueue = WebSocketDataQueuePython +else: + try: + from .reader_c import ( # type: ignore[import-not-found] + WebSocketDataQueue as WebSocketDataQueueCython, + WebSocketReader as WebSocketReaderCython, + ) + + WebSocketReader = WebSocketReaderCython + WebSocketDataQueue = WebSocketDataQueueCython + except ImportError: # pragma: no cover + from .reader_py import ( + WebSocketDataQueue as WebSocketDataQueuePython, + WebSocketReader as WebSocketReaderPython, + ) + + WebSocketReader = WebSocketReaderPython + WebSocketDataQueue = WebSocketDataQueuePython diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.pxd b/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.pxd new file mode 100644 index 0000000000000000000000000000000000000000..a7620d8e87fddd189ed9dbb062e6ec34aa8d673e --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.pxd @@ -0,0 +1,110 @@ +import cython + +from .mask cimport _websocket_mask_cython as websocket_mask + + +cdef unsigned int READ_HEADER +cdef unsigned int READ_PAYLOAD_LENGTH +cdef unsigned int READ_PAYLOAD_MASK +cdef unsigned int READ_PAYLOAD + +cdef int OP_CODE_NOT_SET +cdef int OP_CODE_CONTINUATION +cdef int OP_CODE_TEXT +cdef int OP_CODE_BINARY +cdef int OP_CODE_CLOSE +cdef int OP_CODE_PING +cdef int OP_CODE_PONG + +cdef int COMPRESSED_NOT_SET +cdef int COMPRESSED_FALSE +cdef int COMPRESSED_TRUE + +cdef object UNPACK_LEN3 +cdef object UNPACK_CLOSE_CODE +cdef object TUPLE_NEW + +cdef object WSMsgType +cdef object WSMessage + +cdef object WS_MSG_TYPE_TEXT +cdef object WS_MSG_TYPE_BINARY + +cdef set ALLOWED_CLOSE_CODES +cdef set MESSAGE_TYPES_WITH_CONTENT + +cdef tuple EMPTY_FRAME +cdef tuple EMPTY_FRAME_ERROR + +cdef class WebSocketDataQueue: + + cdef unsigned int _size + cdef public object _protocol + cdef unsigned int _limit + cdef object _loop + cdef bint _eof + cdef object _waiter + cdef object _exception + cdef public object _buffer + cdef object _get_buffer + cdef object _put_buffer + + cdef void _release_waiter(self) + + cpdef void feed_data(self, object data, unsigned int size) + + @cython.locals(size="unsigned int") + cdef _read_from_buffer(self) + +cdef class WebSocketReader: + + cdef WebSocketDataQueue queue + cdef unsigned int _max_msg_size + + cdef Exception _exc + cdef bytearray _partial + cdef unsigned int _state + + cdef int _opcode + cdef bint _frame_fin + cdef int _frame_opcode + cdef list _payload_fragments + cdef Py_ssize_t _frame_payload_len + + cdef bytes _tail + cdef bint _has_mask + cdef bytes _frame_mask + cdef Py_ssize_t _payload_bytes_to_read + cdef unsigned int _payload_len_flag + cdef int _compressed + cdef object _decompressobj + cdef bint _compress + + cpdef tuple feed_data(self, object data) + + @cython.locals( + is_continuation=bint, + fin=bint, + has_partial=bint, + payload_merged=bytes, + ) + cpdef void _handle_frame(self, bint fin, int opcode, object payload, int compressed) except * + + @cython.locals( + start_pos=Py_ssize_t, + data_len=Py_ssize_t, + length=Py_ssize_t, + chunk_size=Py_ssize_t, + chunk_len=Py_ssize_t, + data_len=Py_ssize_t, + data_cstr="const unsigned char *", + first_byte="unsigned char", + second_byte="unsigned char", + f_start_pos=Py_ssize_t, + f_end_pos=Py_ssize_t, + has_mask=bint, + fin=bint, + had_fragments=Py_ssize_t, + payload_bytearray=bytearray, + ) + cpdef void _feed_data(self, bytes data) except * diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.py b/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.py new file mode 100644 index 0000000000000000000000000000000000000000..5166d7ec26091a7183db7d9297511d3e62fd639c --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/reader_c.py @@ -0,0 +1,478 @@ +"""Reader for WebSocket protocol versions 13 and 8.""" + +import asyncio +import builtins +from collections import deque +from typing import Deque, Final, Optional, Set, Tuple, Union + +from ..base_protocol import BaseProtocol +from ..compression_utils import ZLibDecompressor +from ..helpers import _EXC_SENTINEL, set_exception +from ..streams import EofStream +from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask +from .models import ( + WS_DEFLATE_TRAILING, + WebSocketError, + WSCloseCode, + WSMessage, + WSMsgType, +) + +ALLOWED_CLOSE_CODES: Final[Set[int]] = {int(i) for i in WSCloseCode} + +# States for the reader, used to parse the WebSocket frame +# integer values are used so they can be cythonized +READ_HEADER = 1 +READ_PAYLOAD_LENGTH = 2 +READ_PAYLOAD_MASK = 3 +READ_PAYLOAD = 4 + +WS_MSG_TYPE_BINARY = WSMsgType.BINARY +WS_MSG_TYPE_TEXT = WSMsgType.TEXT + +# WSMsgType values unpacked so they can by cythonized to ints +OP_CODE_NOT_SET = -1 +OP_CODE_CONTINUATION = WSMsgType.CONTINUATION.value +OP_CODE_TEXT = WSMsgType.TEXT.value +OP_CODE_BINARY = WSMsgType.BINARY.value +OP_CODE_CLOSE = WSMsgType.CLOSE.value +OP_CODE_PING = WSMsgType.PING.value +OP_CODE_PONG = WSMsgType.PONG.value + +EMPTY_FRAME_ERROR = (True, b"") +EMPTY_FRAME = (False, b"") + +COMPRESSED_NOT_SET = -1 +COMPRESSED_FALSE = 0 +COMPRESSED_TRUE = 1 + +TUPLE_NEW = tuple.__new__ + +cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd + + +class WebSocketDataQueue: + """WebSocketDataQueue resumes and pauses an underlying stream. + + It is a destination for WebSocket data. + """ + + def __init__( + self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop + ) -> None: + self._size = 0 + self._protocol = protocol + self._limit = limit * 2 + self._loop = loop + self._eof = False + self._waiter: Optional[asyncio.Future[None]] = None + self._exception: Union[BaseException, None] = None + self._buffer: Deque[Tuple[WSMessage, int]] = deque() + self._get_buffer = self._buffer.popleft + self._put_buffer = self._buffer.append + + def is_eof(self) -> bool: + return self._eof + + def exception(self) -> Optional[BaseException]: + return self._exception + + def set_exception( + self, + exc: BaseException, + exc_cause: builtins.BaseException = _EXC_SENTINEL, + ) -> None: + self._eof = True + self._exception = exc + if (waiter := self._waiter) is not None: + self._waiter = None + set_exception(waiter, exc, exc_cause) + + def _release_waiter(self) -> None: + if (waiter := self._waiter) is None: + return + self._waiter = None + if not waiter.done(): + waiter.set_result(None) + + def feed_eof(self) -> None: + self._eof = True + self._release_waiter() + self._exception = None # Break cyclic references + + def feed_data(self, data: "WSMessage", size: "cython_int") -> None: + self._size += size + self._put_buffer((data, size)) + self._release_waiter() + if self._size > self._limit and not self._protocol._reading_paused: + self._protocol.pause_reading() + + async def read(self) -> WSMessage: + if not self._buffer and not self._eof: + assert not self._waiter + self._waiter = self._loop.create_future() + try: + await self._waiter + except (asyncio.CancelledError, asyncio.TimeoutError): + self._waiter = None + raise + return self._read_from_buffer() + + def _read_from_buffer(self) -> WSMessage: + if self._buffer: + data, size = self._get_buffer() + self._size -= size + if self._size < self._limit and self._protocol._reading_paused: + self._protocol.resume_reading() + return data + if self._exception is not None: + raise self._exception + raise EofStream + + +class WebSocketReader: + def __init__( + self, queue: WebSocketDataQueue, max_msg_size: int, compress: bool = True + ) -> None: + self.queue = queue + self._max_msg_size = max_msg_size + + self._exc: Optional[Exception] = None + self._partial = bytearray() + self._state = READ_HEADER + + self._opcode: int = OP_CODE_NOT_SET + self._frame_fin = False + self._frame_opcode: int = OP_CODE_NOT_SET + self._payload_fragments: list[bytes] = [] + self._frame_payload_len = 0 + + self._tail: bytes = b"" + self._has_mask = False + self._frame_mask: Optional[bytes] = None + self._payload_bytes_to_read = 0 + self._payload_len_flag = 0 + self._compressed: int = COMPRESSED_NOT_SET + self._decompressobj: Optional[ZLibDecompressor] = None + self._compress = compress + + def feed_eof(self) -> None: + self.queue.feed_eof() + + # data can be bytearray on Windows because proactor event loop uses bytearray + # and asyncio types this to Union[bytes, bytearray, memoryview] so we need + # coerce data to bytes if it is not + def feed_data( + self, data: Union[bytes, bytearray, memoryview] + ) -> Tuple[bool, bytes]: + if type(data) is not bytes: + data = bytes(data) + + if self._exc is not None: + return True, data + + try: + self._feed_data(data) + except Exception as exc: + self._exc = exc + set_exception(self.queue, exc) + return EMPTY_FRAME_ERROR + + return EMPTY_FRAME + + def _handle_frame( + self, + fin: bool, + opcode: Union[int, cython_int], # Union intended: Cython pxd uses C int + payload: Union[bytes, bytearray], + compressed: Union[int, cython_int], # Union intended: Cython pxd uses C int + ) -> None: + msg: WSMessage + if opcode in {OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION}: + # Validate continuation frames before processing + if opcode == OP_CODE_CONTINUATION and self._opcode == OP_CODE_NOT_SET: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Continuation frame for non started message", + ) + + # load text/binary + if not fin: + # got partial frame payload + if opcode != OP_CODE_CONTINUATION: + self._opcode = opcode + self._partial += payload + if self._max_msg_size and len(self._partial) >= self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Message size {len(self._partial)} " + f"exceeds limit {self._max_msg_size}", + ) + return + + has_partial = bool(self._partial) + if opcode == OP_CODE_CONTINUATION: + opcode = self._opcode + self._opcode = OP_CODE_NOT_SET + # previous frame was non finished + # we should get continuation opcode + elif has_partial: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "The opcode in non-fin frame is expected " + f"to be zero, got {opcode!r}", + ) + + assembled_payload: Union[bytes, bytearray] + if has_partial: + assembled_payload = self._partial + payload + self._partial.clear() + else: + assembled_payload = payload + + if self._max_msg_size and len(assembled_payload) >= self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Message size {len(assembled_payload)} " + f"exceeds limit {self._max_msg_size}", + ) + + # Decompress process must to be done after all packets + # received. + if compressed: + if not self._decompressobj: + self._decompressobj = ZLibDecompressor(suppress_deflate_header=True) + # XXX: It's possible that the zlib backend (isal is known to + # do this, maybe others too?) will return max_length bytes, + # but internally buffer more data such that the payload is + # >max_length, so we return one extra byte and if we're able + # to do that, then the message is too big. + payload_merged = self._decompressobj.decompress_sync( + assembled_payload + WS_DEFLATE_TRAILING, + ( + self._max_msg_size + 1 + if self._max_msg_size + else self._max_msg_size + ), + ) + if self._max_msg_size and len(payload_merged) > self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Decompressed message exceeds size limit {self._max_msg_size}", + ) + elif type(assembled_payload) is bytes: + payload_merged = assembled_payload + else: + payload_merged = bytes(assembled_payload) + + if opcode == OP_CODE_TEXT: + try: + text = payload_merged.decode("utf-8") + except UnicodeDecodeError as exc: + raise WebSocketError( + WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" + ) from exc + + # XXX: The Text and Binary messages here can be a performance + # bottleneck, so we use tuple.__new__ to improve performance. + # This is not type safe, but many tests should fail in + # test_client_ws_functional.py if this is wrong. + self.queue.feed_data( + TUPLE_NEW(WSMessage, (WS_MSG_TYPE_TEXT, text, "")), + len(payload_merged), + ) + else: + self.queue.feed_data( + TUPLE_NEW(WSMessage, (WS_MSG_TYPE_BINARY, payload_merged, "")), + len(payload_merged), + ) + elif opcode == OP_CODE_CLOSE: + if len(payload) >= 2: + close_code = UNPACK_CLOSE_CODE(payload[:2])[0] + if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + f"Invalid close code: {close_code}", + ) + try: + close_message = payload[2:].decode("utf-8") + except UnicodeDecodeError as exc: + raise WebSocketError( + WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" + ) from exc + msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, close_code, close_message)) + elif payload: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + f"Invalid close frame: {fin} {opcode} {payload!r}", + ) + else: + msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, 0, "")) + + self.queue.feed_data(msg, 0) + elif opcode == OP_CODE_PING: + msg = TUPLE_NEW(WSMessage, (WSMsgType.PING, payload, "")) + self.queue.feed_data(msg, len(payload)) + elif opcode == OP_CODE_PONG: + msg = TUPLE_NEW(WSMessage, (WSMsgType.PONG, payload, "")) + self.queue.feed_data(msg, len(payload)) + else: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}" + ) + + def _feed_data(self, data: bytes) -> None: + """Return the next frame from the socket.""" + if self._tail: + data, self._tail = self._tail + data, b"" + + start_pos: int = 0 + data_len = len(data) + data_cstr = data + + while True: + # read header + if self._state == READ_HEADER: + if data_len - start_pos < 2: + break + first_byte = data_cstr[start_pos] + second_byte = data_cstr[start_pos + 1] + start_pos += 2 + + fin = (first_byte >> 7) & 1 + rsv1 = (first_byte >> 6) & 1 + rsv2 = (first_byte >> 5) & 1 + rsv3 = (first_byte >> 4) & 1 + opcode = first_byte & 0xF + + # frame-fin = %x0 ; more frames of this message follow + # / %x1 ; final frame of this message + # frame-rsv1 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # frame-rsv2 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # frame-rsv3 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # + # Remove rsv1 from this test for deflate development + if rsv2 or rsv3 or (rsv1 and not self._compress): + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received frame with non-zero reserved bits", + ) + + if opcode > 0x7 and fin == 0: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received fragmented control frame", + ) + + has_mask = (second_byte >> 7) & 1 + length = second_byte & 0x7F + + # Control frames MUST have a payload + # length of 125 bytes or less + if opcode > 0x7 and length > 125: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Control frame payload cannot be larger than 125 bytes", + ) + + # Set compress status if last package is FIN + # OR set compress status if this is first fragment + # Raise error if not first fragment with rsv1 = 0x1 + if self._frame_fin or self._compressed == COMPRESSED_NOT_SET: + self._compressed = COMPRESSED_TRUE if rsv1 else COMPRESSED_FALSE + elif rsv1: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received frame with non-zero reserved bits", + ) + + self._frame_fin = bool(fin) + self._frame_opcode = opcode + self._has_mask = bool(has_mask) + self._payload_len_flag = length + self._state = READ_PAYLOAD_LENGTH + + # read payload length + if self._state == READ_PAYLOAD_LENGTH: + len_flag = self._payload_len_flag + if len_flag == 126: + if data_len - start_pos < 2: + break + first_byte = data_cstr[start_pos] + second_byte = data_cstr[start_pos + 1] + start_pos += 2 + self._payload_bytes_to_read = first_byte << 8 | second_byte + elif len_flag > 126: + if data_len - start_pos < 8: + break + self._payload_bytes_to_read = UNPACK_LEN3(data, start_pos)[0] + start_pos += 8 + else: + self._payload_bytes_to_read = len_flag + + self._state = READ_PAYLOAD_MASK if self._has_mask else READ_PAYLOAD + + # read payload mask + if self._state == READ_PAYLOAD_MASK: + if data_len - start_pos < 4: + break + self._frame_mask = data_cstr[start_pos : start_pos + 4] + start_pos += 4 + self._state = READ_PAYLOAD + + if self._state == READ_PAYLOAD: + chunk_len = data_len - start_pos + if self._payload_bytes_to_read >= chunk_len: + f_end_pos = data_len + self._payload_bytes_to_read -= chunk_len + else: + f_end_pos = start_pos + self._payload_bytes_to_read + self._payload_bytes_to_read = 0 + + had_fragments = self._frame_payload_len + self._frame_payload_len += f_end_pos - start_pos + f_start_pos = start_pos + start_pos = f_end_pos + + if self._payload_bytes_to_read != 0: + # If we don't have a complete frame, we need to save the + # data for the next call to feed_data. + self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) + break + + payload: Union[bytes, bytearray] + if had_fragments: + # We have to join the payload fragments get the payload + self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) + if self._has_mask: + assert self._frame_mask is not None + payload_bytearray = bytearray(b"".join(self._payload_fragments)) + websocket_mask(self._frame_mask, payload_bytearray) + payload = payload_bytearray + else: + payload = b"".join(self._payload_fragments) + self._payload_fragments.clear() + elif self._has_mask: + assert self._frame_mask is not None + payload_bytearray = data_cstr[f_start_pos:f_end_pos] # type: ignore[assignment] + if type(payload_bytearray) is not bytearray: # pragma: no branch + # Cython will do the conversion for us + # but we need to do it for Python and we + # will always get here in Python + payload_bytearray = bytearray(payload_bytearray) + websocket_mask(self._frame_mask, payload_bytearray) + payload = payload_bytearray + else: + payload = data_cstr[f_start_pos:f_end_pos] + + self._handle_frame( + self._frame_fin, self._frame_opcode, payload, self._compressed + ) + self._frame_payload_len = 0 + self._state = READ_HEADER + + # XXX: Cython needs slices to be bounded, so we can't omit the slice end here. + self._tail = data_cstr[start_pos:data_len] if start_pos < data_len else b"" diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/reader_py.py b/lib/python3.12/site-packages/aiohttp/_websocket/reader_py.py new file mode 100644 index 0000000000000000000000000000000000000000..5166d7ec26091a7183db7d9297511d3e62fd639c --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/reader_py.py @@ -0,0 +1,478 @@ +"""Reader for WebSocket protocol versions 13 and 8.""" + +import asyncio +import builtins +from collections import deque +from typing import Deque, Final, Optional, Set, Tuple, Union + +from ..base_protocol import BaseProtocol +from ..compression_utils import ZLibDecompressor +from ..helpers import _EXC_SENTINEL, set_exception +from ..streams import EofStream +from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask +from .models import ( + WS_DEFLATE_TRAILING, + WebSocketError, + WSCloseCode, + WSMessage, + WSMsgType, +) + +ALLOWED_CLOSE_CODES: Final[Set[int]] = {int(i) for i in WSCloseCode} + +# States for the reader, used to parse the WebSocket frame +# integer values are used so they can be cythonized +READ_HEADER = 1 +READ_PAYLOAD_LENGTH = 2 +READ_PAYLOAD_MASK = 3 +READ_PAYLOAD = 4 + +WS_MSG_TYPE_BINARY = WSMsgType.BINARY +WS_MSG_TYPE_TEXT = WSMsgType.TEXT + +# WSMsgType values unpacked so they can by cythonized to ints +OP_CODE_NOT_SET = -1 +OP_CODE_CONTINUATION = WSMsgType.CONTINUATION.value +OP_CODE_TEXT = WSMsgType.TEXT.value +OP_CODE_BINARY = WSMsgType.BINARY.value +OP_CODE_CLOSE = WSMsgType.CLOSE.value +OP_CODE_PING = WSMsgType.PING.value +OP_CODE_PONG = WSMsgType.PONG.value + +EMPTY_FRAME_ERROR = (True, b"") +EMPTY_FRAME = (False, b"") + +COMPRESSED_NOT_SET = -1 +COMPRESSED_FALSE = 0 +COMPRESSED_TRUE = 1 + +TUPLE_NEW = tuple.__new__ + +cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd + + +class WebSocketDataQueue: + """WebSocketDataQueue resumes and pauses an underlying stream. + + It is a destination for WebSocket data. + """ + + def __init__( + self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop + ) -> None: + self._size = 0 + self._protocol = protocol + self._limit = limit * 2 + self._loop = loop + self._eof = False + self._waiter: Optional[asyncio.Future[None]] = None + self._exception: Union[BaseException, None] = None + self._buffer: Deque[Tuple[WSMessage, int]] = deque() + self._get_buffer = self._buffer.popleft + self._put_buffer = self._buffer.append + + def is_eof(self) -> bool: + return self._eof + + def exception(self) -> Optional[BaseException]: + return self._exception + + def set_exception( + self, + exc: BaseException, + exc_cause: builtins.BaseException = _EXC_SENTINEL, + ) -> None: + self._eof = True + self._exception = exc + if (waiter := self._waiter) is not None: + self._waiter = None + set_exception(waiter, exc, exc_cause) + + def _release_waiter(self) -> None: + if (waiter := self._waiter) is None: + return + self._waiter = None + if not waiter.done(): + waiter.set_result(None) + + def feed_eof(self) -> None: + self._eof = True + self._release_waiter() + self._exception = None # Break cyclic references + + def feed_data(self, data: "WSMessage", size: "cython_int") -> None: + self._size += size + self._put_buffer((data, size)) + self._release_waiter() + if self._size > self._limit and not self._protocol._reading_paused: + self._protocol.pause_reading() + + async def read(self) -> WSMessage: + if not self._buffer and not self._eof: + assert not self._waiter + self._waiter = self._loop.create_future() + try: + await self._waiter + except (asyncio.CancelledError, asyncio.TimeoutError): + self._waiter = None + raise + return self._read_from_buffer() + + def _read_from_buffer(self) -> WSMessage: + if self._buffer: + data, size = self._get_buffer() + self._size -= size + if self._size < self._limit and self._protocol._reading_paused: + self._protocol.resume_reading() + return data + if self._exception is not None: + raise self._exception + raise EofStream + + +class WebSocketReader: + def __init__( + self, queue: WebSocketDataQueue, max_msg_size: int, compress: bool = True + ) -> None: + self.queue = queue + self._max_msg_size = max_msg_size + + self._exc: Optional[Exception] = None + self._partial = bytearray() + self._state = READ_HEADER + + self._opcode: int = OP_CODE_NOT_SET + self._frame_fin = False + self._frame_opcode: int = OP_CODE_NOT_SET + self._payload_fragments: list[bytes] = [] + self._frame_payload_len = 0 + + self._tail: bytes = b"" + self._has_mask = False + self._frame_mask: Optional[bytes] = None + self._payload_bytes_to_read = 0 + self._payload_len_flag = 0 + self._compressed: int = COMPRESSED_NOT_SET + self._decompressobj: Optional[ZLibDecompressor] = None + self._compress = compress + + def feed_eof(self) -> None: + self.queue.feed_eof() + + # data can be bytearray on Windows because proactor event loop uses bytearray + # and asyncio types this to Union[bytes, bytearray, memoryview] so we need + # coerce data to bytes if it is not + def feed_data( + self, data: Union[bytes, bytearray, memoryview] + ) -> Tuple[bool, bytes]: + if type(data) is not bytes: + data = bytes(data) + + if self._exc is not None: + return True, data + + try: + self._feed_data(data) + except Exception as exc: + self._exc = exc + set_exception(self.queue, exc) + return EMPTY_FRAME_ERROR + + return EMPTY_FRAME + + def _handle_frame( + self, + fin: bool, + opcode: Union[int, cython_int], # Union intended: Cython pxd uses C int + payload: Union[bytes, bytearray], + compressed: Union[int, cython_int], # Union intended: Cython pxd uses C int + ) -> None: + msg: WSMessage + if opcode in {OP_CODE_TEXT, OP_CODE_BINARY, OP_CODE_CONTINUATION}: + # Validate continuation frames before processing + if opcode == OP_CODE_CONTINUATION and self._opcode == OP_CODE_NOT_SET: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Continuation frame for non started message", + ) + + # load text/binary + if not fin: + # got partial frame payload + if opcode != OP_CODE_CONTINUATION: + self._opcode = opcode + self._partial += payload + if self._max_msg_size and len(self._partial) >= self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Message size {len(self._partial)} " + f"exceeds limit {self._max_msg_size}", + ) + return + + has_partial = bool(self._partial) + if opcode == OP_CODE_CONTINUATION: + opcode = self._opcode + self._opcode = OP_CODE_NOT_SET + # previous frame was non finished + # we should get continuation opcode + elif has_partial: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "The opcode in non-fin frame is expected " + f"to be zero, got {opcode!r}", + ) + + assembled_payload: Union[bytes, bytearray] + if has_partial: + assembled_payload = self._partial + payload + self._partial.clear() + else: + assembled_payload = payload + + if self._max_msg_size and len(assembled_payload) >= self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Message size {len(assembled_payload)} " + f"exceeds limit {self._max_msg_size}", + ) + + # Decompress process must to be done after all packets + # received. + if compressed: + if not self._decompressobj: + self._decompressobj = ZLibDecompressor(suppress_deflate_header=True) + # XXX: It's possible that the zlib backend (isal is known to + # do this, maybe others too?) will return max_length bytes, + # but internally buffer more data such that the payload is + # >max_length, so we return one extra byte and if we're able + # to do that, then the message is too big. + payload_merged = self._decompressobj.decompress_sync( + assembled_payload + WS_DEFLATE_TRAILING, + ( + self._max_msg_size + 1 + if self._max_msg_size + else self._max_msg_size + ), + ) + if self._max_msg_size and len(payload_merged) > self._max_msg_size: + raise WebSocketError( + WSCloseCode.MESSAGE_TOO_BIG, + f"Decompressed message exceeds size limit {self._max_msg_size}", + ) + elif type(assembled_payload) is bytes: + payload_merged = assembled_payload + else: + payload_merged = bytes(assembled_payload) + + if opcode == OP_CODE_TEXT: + try: + text = payload_merged.decode("utf-8") + except UnicodeDecodeError as exc: + raise WebSocketError( + WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" + ) from exc + + # XXX: The Text and Binary messages here can be a performance + # bottleneck, so we use tuple.__new__ to improve performance. + # This is not type safe, but many tests should fail in + # test_client_ws_functional.py if this is wrong. + self.queue.feed_data( + TUPLE_NEW(WSMessage, (WS_MSG_TYPE_TEXT, text, "")), + len(payload_merged), + ) + else: + self.queue.feed_data( + TUPLE_NEW(WSMessage, (WS_MSG_TYPE_BINARY, payload_merged, "")), + len(payload_merged), + ) + elif opcode == OP_CODE_CLOSE: + if len(payload) >= 2: + close_code = UNPACK_CLOSE_CODE(payload[:2])[0] + if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + f"Invalid close code: {close_code}", + ) + try: + close_message = payload[2:].decode("utf-8") + except UnicodeDecodeError as exc: + raise WebSocketError( + WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message" + ) from exc + msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, close_code, close_message)) + elif payload: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + f"Invalid close frame: {fin} {opcode} {payload!r}", + ) + else: + msg = TUPLE_NEW(WSMessage, (WSMsgType.CLOSE, 0, "")) + + self.queue.feed_data(msg, 0) + elif opcode == OP_CODE_PING: + msg = TUPLE_NEW(WSMessage, (WSMsgType.PING, payload, "")) + self.queue.feed_data(msg, len(payload)) + elif opcode == OP_CODE_PONG: + msg = TUPLE_NEW(WSMessage, (WSMsgType.PONG, payload, "")) + self.queue.feed_data(msg, len(payload)) + else: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}" + ) + + def _feed_data(self, data: bytes) -> None: + """Return the next frame from the socket.""" + if self._tail: + data, self._tail = self._tail + data, b"" + + start_pos: int = 0 + data_len = len(data) + data_cstr = data + + while True: + # read header + if self._state == READ_HEADER: + if data_len - start_pos < 2: + break + first_byte = data_cstr[start_pos] + second_byte = data_cstr[start_pos + 1] + start_pos += 2 + + fin = (first_byte >> 7) & 1 + rsv1 = (first_byte >> 6) & 1 + rsv2 = (first_byte >> 5) & 1 + rsv3 = (first_byte >> 4) & 1 + opcode = first_byte & 0xF + + # frame-fin = %x0 ; more frames of this message follow + # / %x1 ; final frame of this message + # frame-rsv1 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # frame-rsv2 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # frame-rsv3 = %x0 ; + # 1 bit, MUST be 0 unless negotiated otherwise + # + # Remove rsv1 from this test for deflate development + if rsv2 or rsv3 or (rsv1 and not self._compress): + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received frame with non-zero reserved bits", + ) + + if opcode > 0x7 and fin == 0: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received fragmented control frame", + ) + + has_mask = (second_byte >> 7) & 1 + length = second_byte & 0x7F + + # Control frames MUST have a payload + # length of 125 bytes or less + if opcode > 0x7 and length > 125: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Control frame payload cannot be larger than 125 bytes", + ) + + # Set compress status if last package is FIN + # OR set compress status if this is first fragment + # Raise error if not first fragment with rsv1 = 0x1 + if self._frame_fin or self._compressed == COMPRESSED_NOT_SET: + self._compressed = COMPRESSED_TRUE if rsv1 else COMPRESSED_FALSE + elif rsv1: + raise WebSocketError( + WSCloseCode.PROTOCOL_ERROR, + "Received frame with non-zero reserved bits", + ) + + self._frame_fin = bool(fin) + self._frame_opcode = opcode + self._has_mask = bool(has_mask) + self._payload_len_flag = length + self._state = READ_PAYLOAD_LENGTH + + # read payload length + if self._state == READ_PAYLOAD_LENGTH: + len_flag = self._payload_len_flag + if len_flag == 126: + if data_len - start_pos < 2: + break + first_byte = data_cstr[start_pos] + second_byte = data_cstr[start_pos + 1] + start_pos += 2 + self._payload_bytes_to_read = first_byte << 8 | second_byte + elif len_flag > 126: + if data_len - start_pos < 8: + break + self._payload_bytes_to_read = UNPACK_LEN3(data, start_pos)[0] + start_pos += 8 + else: + self._payload_bytes_to_read = len_flag + + self._state = READ_PAYLOAD_MASK if self._has_mask else READ_PAYLOAD + + # read payload mask + if self._state == READ_PAYLOAD_MASK: + if data_len - start_pos < 4: + break + self._frame_mask = data_cstr[start_pos : start_pos + 4] + start_pos += 4 + self._state = READ_PAYLOAD + + if self._state == READ_PAYLOAD: + chunk_len = data_len - start_pos + if self._payload_bytes_to_read >= chunk_len: + f_end_pos = data_len + self._payload_bytes_to_read -= chunk_len + else: + f_end_pos = start_pos + self._payload_bytes_to_read + self._payload_bytes_to_read = 0 + + had_fragments = self._frame_payload_len + self._frame_payload_len += f_end_pos - start_pos + f_start_pos = start_pos + start_pos = f_end_pos + + if self._payload_bytes_to_read != 0: + # If we don't have a complete frame, we need to save the + # data for the next call to feed_data. + self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) + break + + payload: Union[bytes, bytearray] + if had_fragments: + # We have to join the payload fragments get the payload + self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos]) + if self._has_mask: + assert self._frame_mask is not None + payload_bytearray = bytearray(b"".join(self._payload_fragments)) + websocket_mask(self._frame_mask, payload_bytearray) + payload = payload_bytearray + else: + payload = b"".join(self._payload_fragments) + self._payload_fragments.clear() + elif self._has_mask: + assert self._frame_mask is not None + payload_bytearray = data_cstr[f_start_pos:f_end_pos] # type: ignore[assignment] + if type(payload_bytearray) is not bytearray: # pragma: no branch + # Cython will do the conversion for us + # but we need to do it for Python and we + # will always get here in Python + payload_bytearray = bytearray(payload_bytearray) + websocket_mask(self._frame_mask, payload_bytearray) + payload = payload_bytearray + else: + payload = data_cstr[f_start_pos:f_end_pos] + + self._handle_frame( + self._frame_fin, self._frame_opcode, payload, self._compressed + ) + self._frame_payload_len = 0 + self._state = READ_HEADER + + # XXX: Cython needs slices to be bounded, so we can't omit the slice end here. + self._tail = data_cstr[start_pos:data_len] if start_pos < data_len else b"" diff --git a/lib/python3.12/site-packages/aiohttp/_websocket/writer.py b/lib/python3.12/site-packages/aiohttp/_websocket/writer.py new file mode 100644 index 0000000000000000000000000000000000000000..9604202357cddd5714242144e30f4a05c4f638fa --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/_websocket/writer.py @@ -0,0 +1,262 @@ +"""WebSocket protocol versions 13 and 8.""" + +import asyncio +import random +import sys +from functools import partial +from typing import Final, Optional, Set, Union + +from ..base_protocol import BaseProtocol +from ..client_exceptions import ClientConnectionResetError +from ..compression_utils import ZLibBackend, ZLibCompressor +from .helpers import ( + MASK_LEN, + MSG_SIZE, + PACK_CLOSE_CODE, + PACK_LEN1, + PACK_LEN2, + PACK_LEN3, + PACK_RANDBITS, + websocket_mask, +) +from .models import WS_DEFLATE_TRAILING, WSMsgType + +DEFAULT_LIMIT: Final[int] = 2**16 + +# WebSocket opcode boundary: opcodes 0-7 are data frames, 8-15 are control frames +# Control frames (ping, pong, close) are never compressed +WS_CONTROL_FRAME_OPCODE: Final[int] = 8 + +# For websockets, keeping latency low is extremely important as implementations +# generally expect to be able to send and receive messages quickly. We use a +# larger chunk size to reduce the number of executor calls and avoid task +# creation overhead, since both are significant sources of latency when chunks +# are small. A size of 16KiB was chosen as a balance between avoiding task +# overhead and not blocking the event loop too long with synchronous compression. + +WEBSOCKET_MAX_SYNC_CHUNK_SIZE = 16 * 1024 + + +class WebSocketWriter: + """WebSocket writer. + + The writer is responsible for sending messages to the client. It is + created by the protocol when a connection is established. The writer + should avoid implementing any application logic and should only be + concerned with the low-level details of the WebSocket protocol. + """ + + def __init__( + self, + protocol: BaseProtocol, + transport: asyncio.Transport, + *, + use_mask: bool = False, + limit: int = DEFAULT_LIMIT, + random: random.Random = random.Random(), + compress: int = 0, + notakeover: bool = False, + ) -> None: + """Initialize a WebSocket writer.""" + self.protocol = protocol + self.transport = transport + self.use_mask = use_mask + self.get_random_bits = partial(random.getrandbits, 32) + self.compress = compress + self.notakeover = notakeover + self._closing = False + self._limit = limit + self._output_size = 0 + self._compressobj: Optional[ZLibCompressor] = None + self._send_lock = asyncio.Lock() + self._background_tasks: Set[asyncio.Task[None]] = set() + + async def send_frame( + self, message: bytes, opcode: int, compress: Optional[int] = None + ) -> None: + """Send a frame over the websocket with message as its payload.""" + if self._closing and not (opcode & WSMsgType.CLOSE): + raise ClientConnectionResetError("Cannot write to closing transport") + + if not (compress or self.compress) or opcode >= WS_CONTROL_FRAME_OPCODE: + # Non-compressed frames don't need lock or shield + self._write_websocket_frame(message, opcode, 0) + elif len(message) <= WEBSOCKET_MAX_SYNC_CHUNK_SIZE: + # Small compressed payloads - compress synchronously in event loop + # We need the lock even though sync compression has no await points. + # This prevents small frames from interleaving with large frames that + # compress in the executor, avoiding compressor state corruption. + async with self._send_lock: + self._send_compressed_frame_sync(message, opcode, compress) + else: + # Large compressed frames need shield to prevent corruption + # For large compressed frames, the entire compress+send + # operation must be atomic. If cancelled after compression but + # before send, the compressor state would be advanced but data + # not sent, corrupting subsequent frames. + # Create a task to shield from cancellation + # The lock is acquired inside the shielded task so the entire + # operation (lock + compress + send) completes atomically. + # Use eager_start on Python 3.12+ to avoid scheduling overhead + loop = asyncio.get_running_loop() + coro = self._send_compressed_frame_async_locked(message, opcode, compress) + if sys.version_info >= (3, 12): + send_task = asyncio.Task(coro, loop=loop, eager_start=True) + else: + send_task = loop.create_task(coro) + # Keep a strong reference to prevent garbage collection + self._background_tasks.add(send_task) + send_task.add_done_callback(self._background_tasks.discard) + await asyncio.shield(send_task) + + # It is safe to return control to the event loop when using compression + # after this point as we have already sent or buffered all the data. + # Once we have written output_size up to the limit, we call the + # drain helper which waits for the transport to be ready to accept + # more data. This is a flow control mechanism to prevent the buffer + # from growing too large. The drain helper will return right away + # if the writer is not paused. + if self._output_size > self._limit: + self._output_size = 0 + if self.protocol._paused: + await self.protocol._drain_helper() + + def _write_websocket_frame(self, message: bytes, opcode: int, rsv: int) -> None: + """ + Write a websocket frame to the transport. + + This method handles frame header construction, masking, and writing to transport. + It does not handle compression or flow control - those are the responsibility + of the caller. + """ + msg_length = len(message) + + use_mask = self.use_mask + mask_bit = 0x80 if use_mask else 0 + + # Depending on the message length, the header is assembled differently. + # The first byte is reserved for the opcode and the RSV bits. + first_byte = 0x80 | rsv | opcode + if msg_length < 126: + header = PACK_LEN1(first_byte, msg_length | mask_bit) + header_len = 2 + elif msg_length < 65536: + header = PACK_LEN2(first_byte, 126 | mask_bit, msg_length) + header_len = 4 + else: + header = PACK_LEN3(first_byte, 127 | mask_bit, msg_length) + header_len = 10 + + if self.transport.is_closing(): + raise ClientConnectionResetError("Cannot write to closing transport") + + # https://datatracker.ietf.org/doc/html/rfc6455#section-5.3 + # If we are using a mask, we need to generate it randomly + # and apply it to the message before sending it. A mask is + # a 32-bit value that is applied to the message using a + # bitwise XOR operation. It is used to prevent certain types + # of attacks on the websocket protocol. The mask is only used + # when aiohttp is acting as a client. Servers do not use a mask. + if use_mask: + mask = PACK_RANDBITS(self.get_random_bits()) + message = bytearray(message) + websocket_mask(mask, message) + self.transport.write(header + mask + message) + self._output_size += MASK_LEN + elif msg_length > MSG_SIZE: + self.transport.write(header) + self.transport.write(message) + else: + self.transport.write(header + message) + + self._output_size += header_len + msg_length + + def _get_compressor(self, compress: Optional[int]) -> ZLibCompressor: + """Get or create a compressor object for the given compression level.""" + if compress: + # Do not set self._compress if compressing is for this frame + return ZLibCompressor( + level=ZLibBackend.Z_BEST_SPEED, + wbits=-compress, + max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE, + ) + if not self._compressobj: + self._compressobj = ZLibCompressor( + level=ZLibBackend.Z_BEST_SPEED, + wbits=-self.compress, + max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE, + ) + return self._compressobj + + def _send_compressed_frame_sync( + self, message: bytes, opcode: int, compress: Optional[int] + ) -> None: + """ + Synchronous send for small compressed frames. + + This is used for small compressed payloads that compress synchronously in the event loop. + Since there are no await points, this is inherently cancellation-safe. + """ + # RSV are the reserved bits in the frame header. They are used to + # indicate that the frame is using an extension. + # https://datatracker.ietf.org/doc/html/rfc6455#section-5.2 + compressobj = self._get_compressor(compress) + # (0x40) RSV1 is set for compressed frames + # https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1 + self._write_websocket_frame( + ( + compressobj.compress_sync(message) + + compressobj.flush( + ZLibBackend.Z_FULL_FLUSH + if self.notakeover + else ZLibBackend.Z_SYNC_FLUSH + ) + ).removesuffix(WS_DEFLATE_TRAILING), + opcode, + 0x40, + ) + + async def _send_compressed_frame_async_locked( + self, message: bytes, opcode: int, compress: Optional[int] + ) -> None: + """ + Async send for large compressed frames with lock. + + Acquires the lock and compresses large payloads asynchronously in + the executor. The lock is held for the entire operation to ensure + the compressor state is not corrupted by concurrent sends. + + MUST be run shielded from cancellation. If cancelled after + compression but before sending, the compressor state would be + advanced but data not sent, corrupting subsequent frames. + """ + async with self._send_lock: + # RSV are the reserved bits in the frame header. They are used to + # indicate that the frame is using an extension. + # https://datatracker.ietf.org/doc/html/rfc6455#section-5.2 + compressobj = self._get_compressor(compress) + # (0x40) RSV1 is set for compressed frames + # https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1 + self._write_websocket_frame( + ( + await compressobj.compress(message) + + compressobj.flush( + ZLibBackend.Z_FULL_FLUSH + if self.notakeover + else ZLibBackend.Z_SYNC_FLUSH + ) + ).removesuffix(WS_DEFLATE_TRAILING), + opcode, + 0x40, + ) + + async def close(self, code: int = 1000, message: Union[bytes, str] = b"") -> None: + """Close the websocket, sending the specified code and message.""" + if isinstance(message, str): + message = message.encode("utf-8") + try: + await self.send_frame( + PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE + ) + finally: + self._closing = True diff --git a/lib/python3.12/site-packages/aiohttp/abc.py b/lib/python3.12/site-packages/aiohttp/abc.py new file mode 100644 index 0000000000000000000000000000000000000000..faf09575afbb2d6842ad639dada0fdebc2e28945 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/abc.py @@ -0,0 +1,268 @@ +import asyncio +import logging +import socket +from abc import ABC, abstractmethod +from collections.abc import Sized +from http.cookies import BaseCookie, Morsel +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Dict, + Generator, + Iterable, + List, + Optional, + Sequence, + Tuple, + TypedDict, + Union, +) + +from multidict import CIMultiDict +from yarl import URL + +from ._cookie_helpers import parse_set_cookie_headers +from .typedefs import LooseCookies + +if TYPE_CHECKING: + from .web_app import Application + from .web_exceptions import HTTPException + from .web_request import BaseRequest, Request + from .web_response import StreamResponse +else: + BaseRequest = Request = Application = StreamResponse = None + HTTPException = None + + +class AbstractRouter(ABC): + def __init__(self) -> None: + self._frozen = False + + def post_init(self, app: Application) -> None: + """Post init stage. + + Not an abstract method for sake of backward compatibility, + but if the router wants to be aware of the application + it can override this. + """ + + @property + def frozen(self) -> bool: + return self._frozen + + def freeze(self) -> None: + """Freeze router.""" + self._frozen = True + + @abstractmethod + async def resolve(self, request: Request) -> "AbstractMatchInfo": + """Return MATCH_INFO for given request""" + + +class AbstractMatchInfo(ABC): + + __slots__ = () + + @property # pragma: no branch + @abstractmethod + def handler(self) -> Callable[[Request], Awaitable[StreamResponse]]: + """Execute matched request handler""" + + @property + @abstractmethod + def expect_handler( + self, + ) -> Callable[[Request], Awaitable[Optional[StreamResponse]]]: + """Expect handler for 100-continue processing""" + + @property # pragma: no branch + @abstractmethod + def http_exception(self) -> Optional[HTTPException]: + """HTTPException instance raised on router's resolving, or None""" + + @abstractmethod # pragma: no branch + def get_info(self) -> Dict[str, Any]: + """Return a dict with additional info useful for introspection""" + + @property # pragma: no branch + @abstractmethod + def apps(self) -> Tuple[Application, ...]: + """Stack of nested applications. + + Top level application is left-most element. + + """ + + @abstractmethod + def add_app(self, app: Application) -> None: + """Add application to the nested apps stack.""" + + @abstractmethod + def freeze(self) -> None: + """Freeze the match info. + + The method is called after route resolution. + + After the call .add_app() is forbidden. + + """ + + +class AbstractView(ABC): + """Abstract class based view.""" + + def __init__(self, request: Request) -> None: + self._request = request + + @property + def request(self) -> Request: + """Request instance.""" + return self._request + + @abstractmethod + def __await__(self) -> Generator[None, None, StreamResponse]: + """Execute the view handler.""" + + +class ResolveResult(TypedDict): + """Resolve result. + + This is the result returned from an AbstractResolver's + resolve method. + + :param hostname: The hostname that was provided. + :param host: The IP address that was resolved. + :param port: The port that was resolved. + :param family: The address family that was resolved. + :param proto: The protocol that was resolved. + :param flags: The flags that were resolved. + """ + + hostname: str + host: str + port: int + family: int + proto: int + flags: int + + +class AbstractResolver(ABC): + """Abstract DNS resolver.""" + + @abstractmethod + async def resolve( + self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET + ) -> List[ResolveResult]: + """Return IP address for given hostname""" + + @abstractmethod + async def close(self) -> None: + """Release resolver""" + + +if TYPE_CHECKING: + IterableBase = Iterable[Morsel[str]] +else: + IterableBase = Iterable + + +ClearCookiePredicate = Callable[["Morsel[str]"], bool] + + +class AbstractCookieJar(Sized, IterableBase): + """Abstract Cookie Jar.""" + + def __init__(self, *, loop: Optional[asyncio.AbstractEventLoop] = None) -> None: + self._loop = loop or asyncio.get_running_loop() + + @property + @abstractmethod + def quote_cookie(self) -> bool: + """Return True if cookies should be quoted.""" + + @abstractmethod + def clear(self, predicate: Optional[ClearCookiePredicate] = None) -> None: + """Clear all cookies if no predicate is passed.""" + + @abstractmethod + def clear_domain(self, domain: str) -> None: + """Clear all cookies for domain and all subdomains.""" + + @abstractmethod + def update_cookies(self, cookies: LooseCookies, response_url: URL = URL()) -> None: + """Update cookies.""" + + def update_cookies_from_headers( + self, headers: Sequence[str], response_url: URL + ) -> None: + """Update cookies from raw Set-Cookie headers.""" + if headers and (cookies_to_update := parse_set_cookie_headers(headers)): + self.update_cookies(cookies_to_update, response_url) + + @abstractmethod + def filter_cookies(self, request_url: URL) -> "BaseCookie[str]": + """Return the jar's cookies filtered by their attributes.""" + + +class AbstractStreamWriter(ABC): + """Abstract stream writer.""" + + buffer_size: int = 0 + output_size: int = 0 + length: Optional[int] = 0 + + @abstractmethod + async def write(self, chunk: Union[bytes, bytearray, memoryview]) -> None: + """Write chunk into stream.""" + + @abstractmethod + async def write_eof(self, chunk: bytes = b"") -> None: + """Write last chunk.""" + + @abstractmethod + async def drain(self) -> None: + """Flush the write buffer.""" + + @abstractmethod + def enable_compression( + self, encoding: str = "deflate", strategy: Optional[int] = None + ) -> None: + """Enable HTTP body compression""" + + @abstractmethod + def enable_chunking(self) -> None: + """Enable HTTP chunked mode""" + + @abstractmethod + async def write_headers( + self, status_line: str, headers: "CIMultiDict[str]" + ) -> None: + """Write HTTP headers""" + + def send_headers(self) -> None: + """Force sending buffered headers if not already sent. + + Required only if write_headers() buffers headers instead of sending immediately. + For backwards compatibility, this method does nothing by default. + """ + + +class AbstractAccessLogger(ABC): + """Abstract writer to access log.""" + + __slots__ = ("logger", "log_format") + + def __init__(self, logger: logging.Logger, log_format: str) -> None: + self.logger = logger + self.log_format = log_format + + @abstractmethod + def log(self, request: BaseRequest, response: StreamResponse, time: float) -> None: + """Emit log to logger.""" + + @property + def enabled(self) -> bool: + """Check if logger is enabled.""" + return True diff --git a/lib/python3.12/site-packages/aiohttp/base_protocol.py b/lib/python3.12/site-packages/aiohttp/base_protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..b0a67ed6ff68ca5bc48be9ac472ee755369b2720 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/base_protocol.py @@ -0,0 +1,100 @@ +import asyncio +from typing import Optional, cast + +from .client_exceptions import ClientConnectionResetError +from .helpers import set_exception +from .tcp_helpers import tcp_nodelay + + +class BaseProtocol(asyncio.Protocol): + __slots__ = ( + "_loop", + "_paused", + "_drain_waiter", + "_connection_lost", + "_reading_paused", + "transport", + ) + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop: asyncio.AbstractEventLoop = loop + self._paused = False + self._drain_waiter: Optional[asyncio.Future[None]] = None + self._reading_paused = False + + self.transport: Optional[asyncio.Transport] = None + + @property + def connected(self) -> bool: + """Return True if the connection is open.""" + return self.transport is not None + + @property + def writing_paused(self) -> bool: + return self._paused + + def pause_writing(self) -> None: + assert not self._paused + self._paused = True + + def resume_writing(self) -> None: + assert self._paused + self._paused = False + + waiter = self._drain_waiter + if waiter is not None: + self._drain_waiter = None + if not waiter.done(): + waiter.set_result(None) + + def pause_reading(self) -> None: + if not self._reading_paused and self.transport is not None: + try: + self.transport.pause_reading() + except (AttributeError, NotImplementedError, RuntimeError): + pass + self._reading_paused = True + + def resume_reading(self) -> None: + if self._reading_paused and self.transport is not None: + try: + self.transport.resume_reading() + except (AttributeError, NotImplementedError, RuntimeError): + pass + self._reading_paused = False + + def connection_made(self, transport: asyncio.BaseTransport) -> None: + tr = cast(asyncio.Transport, transport) + tcp_nodelay(tr, True) + self.transport = tr + + def connection_lost(self, exc: Optional[BaseException]) -> None: + # Wake up the writer if currently paused. + self.transport = None + if not self._paused: + return + waiter = self._drain_waiter + if waiter is None: + return + self._drain_waiter = None + if waiter.done(): + return + if exc is None: + waiter.set_result(None) + else: + set_exception( + waiter, + ConnectionError("Connection lost"), + exc, + ) + + async def _drain_helper(self) -> None: + if self.transport is None: + raise ClientConnectionResetError("Connection lost") + if not self._paused: + return + waiter = self._drain_waiter + if waiter is None: + waiter = self._loop.create_future() + self._drain_waiter = waiter + await asyncio.shield(waiter) diff --git a/lib/python3.12/site-packages/aiohttp/client.py b/lib/python3.12/site-packages/aiohttp/client.py new file mode 100644 index 0000000000000000000000000000000000000000..bc4ee17caf088e6148b72a5daa5e854a269732eb --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client.py @@ -0,0 +1,1635 @@ +"""HTTP Client for asyncio.""" + +import asyncio +import base64 +import hashlib +import json +import os +import sys +import traceback +import warnings +from contextlib import suppress +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Coroutine, + Final, + FrozenSet, + Generator, + Generic, + Iterable, + List, + Mapping, + Optional, + Sequence, + Set, + Tuple, + Type, + TypedDict, + TypeVar, + Union, +) + +import attr +from multidict import CIMultiDict, MultiDict, MultiDictProxy, istr +from yarl import URL + +from . import hdrs, http, payload +from ._websocket.reader import WebSocketDataQueue +from .abc import AbstractCookieJar +from .client_exceptions import ( + ClientConnectionError, + ClientConnectionResetError, + ClientConnectorCertificateError, + ClientConnectorDNSError, + ClientConnectorError, + ClientConnectorSSLError, + ClientError, + ClientHttpProxyError, + ClientOSError, + ClientPayloadError, + ClientProxyConnectionError, + ClientResponseError, + ClientSSLError, + ConnectionTimeoutError, + ContentTypeError, + InvalidURL, + InvalidUrlClientError, + InvalidUrlRedirectClientError, + NonHttpUrlClientError, + NonHttpUrlRedirectClientError, + RedirectClientError, + ServerConnectionError, + ServerDisconnectedError, + ServerFingerprintMismatch, + ServerTimeoutError, + SocketTimeoutError, + TooManyRedirects, + WSMessageTypeError, + WSServerHandshakeError, +) +from .client_middlewares import ClientMiddlewareType, build_client_middlewares +from .client_reqrep import ( + ClientRequest as ClientRequest, + ClientResponse as ClientResponse, + Fingerprint as Fingerprint, + RequestInfo as RequestInfo, + _merge_ssl_params, +) +from .client_ws import ( + DEFAULT_WS_CLIENT_TIMEOUT, + ClientWebSocketResponse as ClientWebSocketResponse, + ClientWSTimeout as ClientWSTimeout, +) +from .connector import ( + HTTP_AND_EMPTY_SCHEMA_SET, + BaseConnector as BaseConnector, + NamedPipeConnector as NamedPipeConnector, + TCPConnector as TCPConnector, + UnixConnector as UnixConnector, +) +from .cookiejar import CookieJar +from .helpers import ( + _SENTINEL, + DEBUG, + EMPTY_BODY_METHODS, + BasicAuth, + TimeoutHandle, + basicauth_from_netrc, + get_env_proxy_for_url, + netrc_from_env, + sentinel, + strip_auth_from_url, +) +from .http import WS_KEY, HttpVersion, WebSocketReader, WebSocketWriter +from .http_websocket import WSHandshakeError, ws_ext_gen, ws_ext_parse +from .tracing import Trace, TraceConfig +from .typedefs import JSONEncoder, LooseCookies, LooseHeaders, Query, StrOrURL + +__all__ = ( + # client_exceptions + "ClientConnectionError", + "ClientConnectionResetError", + "ClientConnectorCertificateError", + "ClientConnectorDNSError", + "ClientConnectorError", + "ClientConnectorSSLError", + "ClientError", + "ClientHttpProxyError", + "ClientOSError", + "ClientPayloadError", + "ClientProxyConnectionError", + "ClientResponseError", + "ClientSSLError", + "ConnectionTimeoutError", + "ContentTypeError", + "InvalidURL", + "InvalidUrlClientError", + "RedirectClientError", + "NonHttpUrlClientError", + "InvalidUrlRedirectClientError", + "NonHttpUrlRedirectClientError", + "ServerConnectionError", + "ServerDisconnectedError", + "ServerFingerprintMismatch", + "ServerTimeoutError", + "SocketTimeoutError", + "TooManyRedirects", + "WSServerHandshakeError", + # client_reqrep + "ClientRequest", + "ClientResponse", + "Fingerprint", + "RequestInfo", + # connector + "BaseConnector", + "TCPConnector", + "UnixConnector", + "NamedPipeConnector", + # client_ws + "ClientWebSocketResponse", + # client + "ClientSession", + "ClientTimeout", + "ClientWSTimeout", + "request", + "WSMessageTypeError", +) + + +if TYPE_CHECKING: + from ssl import SSLContext +else: + SSLContext = None + +if sys.version_info >= (3, 11) and TYPE_CHECKING: + from typing import Unpack + + +class _RequestOptions(TypedDict, total=False): + params: Query + data: Any + json: Any + cookies: Union[LooseCookies, None] + headers: Union[LooseHeaders, None] + skip_auto_headers: Union[Iterable[str], None] + auth: Union[BasicAuth, None] + allow_redirects: bool + max_redirects: int + compress: Union[str, bool, None] + chunked: Union[bool, None] + expect100: bool + raise_for_status: Union[None, bool, Callable[[ClientResponse], Awaitable[None]]] + read_until_eof: bool + proxy: Union[StrOrURL, None] + proxy_auth: Union[BasicAuth, None] + timeout: "Union[ClientTimeout, _SENTINEL, None]" + ssl: Union[SSLContext, bool, Fingerprint] + server_hostname: Union[str, None] + proxy_headers: Union[LooseHeaders, None] + trace_request_ctx: Union[Mapping[str, Any], None] + read_bufsize: Union[int, None] + auto_decompress: Union[bool, None] + max_line_size: Union[int, None] + max_field_size: Union[int, None] + middlewares: Optional[Sequence[ClientMiddlewareType]] + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class ClientTimeout: + total: Optional[float] = None + connect: Optional[float] = None + sock_read: Optional[float] = None + sock_connect: Optional[float] = None + ceil_threshold: float = 5 + + # pool_queue_timeout: Optional[float] = None + # dns_resolution_timeout: Optional[float] = None + # socket_connect_timeout: Optional[float] = None + # connection_acquiring_timeout: Optional[float] = None + # new_connection_timeout: Optional[float] = None + # http_header_timeout: Optional[float] = None + # response_body_timeout: Optional[float] = None + + # to create a timeout specific for a single request, either + # - create a completely new one to overwrite the default + # - or use http://www.attrs.org/en/stable/api.html#attr.evolve + # to overwrite the defaults + + +# 5 Minute default read timeout +DEFAULT_TIMEOUT: Final[ClientTimeout] = ClientTimeout(total=5 * 60, sock_connect=30) + +# https://www.rfc-editor.org/rfc/rfc9110#section-9.2.2 +IDEMPOTENT_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE", "PUT", "DELETE"}) + +_RetType = TypeVar("_RetType", ClientResponse, ClientWebSocketResponse) +_CharsetResolver = Callable[[ClientResponse, bytes], str] + + +class ClientSession: + """First-class interface for making HTTP requests.""" + + ATTRS = frozenset( + [ + "_base_url", + "_base_url_origin", + "_source_traceback", + "_connector", + "_loop", + "_cookie_jar", + "_connector_owner", + "_default_auth", + "_version", + "_json_serialize", + "_requote_redirect_url", + "_timeout", + "_raise_for_status", + "_auto_decompress", + "_trust_env", + "_default_headers", + "_skip_auto_headers", + "_request_class", + "_response_class", + "_ws_response_class", + "_trace_configs", + "_read_bufsize", + "_max_line_size", + "_max_field_size", + "_resolve_charset", + "_default_proxy", + "_default_proxy_auth", + "_retry_connection", + "_middlewares", + "requote_redirect_url", + ] + ) + + _source_traceback: Optional[traceback.StackSummary] = None + _connector: Optional[BaseConnector] = None + + def __init__( + self, + base_url: Optional[StrOrURL] = None, + *, + connector: Optional[BaseConnector] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + cookies: Optional[LooseCookies] = None, + headers: Optional[LooseHeaders] = None, + proxy: Optional[StrOrURL] = None, + proxy_auth: Optional[BasicAuth] = None, + skip_auto_headers: Optional[Iterable[str]] = None, + auth: Optional[BasicAuth] = None, + json_serialize: JSONEncoder = json.dumps, + request_class: Type[ClientRequest] = ClientRequest, + response_class: Type[ClientResponse] = ClientResponse, + ws_response_class: Type[ClientWebSocketResponse] = ClientWebSocketResponse, + version: HttpVersion = http.HttpVersion11, + cookie_jar: Optional[AbstractCookieJar] = None, + connector_owner: bool = True, + raise_for_status: Union[ + bool, Callable[[ClientResponse], Awaitable[None]] + ] = False, + read_timeout: Union[float, _SENTINEL] = sentinel, + conn_timeout: Optional[float] = None, + timeout: Union[object, ClientTimeout] = sentinel, + auto_decompress: bool = True, + trust_env: bool = False, + requote_redirect_url: bool = True, + trace_configs: Optional[List[TraceConfig]] = None, + read_bufsize: int = 2**16, + max_line_size: int = 8190, + max_field_size: int = 8190, + fallback_charset_resolver: _CharsetResolver = lambda r, b: "utf-8", + middlewares: Sequence[ClientMiddlewareType] = (), + ssl_shutdown_timeout: Union[_SENTINEL, None, float] = sentinel, + ) -> None: + # We initialise _connector to None immediately, as it's referenced in __del__() + # and could cause issues if an exception occurs during initialisation. + self._connector: Optional[BaseConnector] = None + + if loop is None: + if connector is not None: + loop = connector._loop + + loop = loop or asyncio.get_running_loop() + + if base_url is None or isinstance(base_url, URL): + self._base_url: Optional[URL] = base_url + self._base_url_origin = None if base_url is None else base_url.origin() + else: + self._base_url = URL(base_url) + self._base_url_origin = self._base_url.origin() + assert self._base_url.absolute, "Only absolute URLs are supported" + if self._base_url is not None and not self._base_url.path.endswith("/"): + raise ValueError("base_url must have a trailing '/'") + + if timeout is sentinel or timeout is None: + self._timeout = DEFAULT_TIMEOUT + if read_timeout is not sentinel: + warnings.warn( + "read_timeout is deprecated, use timeout argument instead", + DeprecationWarning, + stacklevel=2, + ) + self._timeout = attr.evolve(self._timeout, total=read_timeout) + if conn_timeout is not None: + self._timeout = attr.evolve(self._timeout, connect=conn_timeout) + warnings.warn( + "conn_timeout is deprecated, use timeout argument instead", + DeprecationWarning, + stacklevel=2, + ) + else: + if not isinstance(timeout, ClientTimeout): + raise ValueError( + f"timeout parameter cannot be of {type(timeout)} type, " + "please use 'timeout=ClientTimeout(...)'", + ) + self._timeout = timeout + if read_timeout is not sentinel: + raise ValueError( + "read_timeout and timeout parameters " + "conflict, please setup " + "timeout.read" + ) + if conn_timeout is not None: + raise ValueError( + "conn_timeout and timeout parameters " + "conflict, please setup " + "timeout.connect" + ) + + if ssl_shutdown_timeout is not sentinel: + warnings.warn( + "The ssl_shutdown_timeout parameter is deprecated and will be removed in aiohttp 4.0", + DeprecationWarning, + stacklevel=2, + ) + + if connector is None: + connector = TCPConnector( + loop=loop, ssl_shutdown_timeout=ssl_shutdown_timeout + ) + + if connector._loop is not loop: + raise RuntimeError("Session and connector has to use same event loop") + + self._loop = loop + + if loop.get_debug(): + self._source_traceback = traceback.extract_stack(sys._getframe(1)) + + if cookie_jar is None: + cookie_jar = CookieJar(loop=loop) + self._cookie_jar = cookie_jar + + if cookies: + self._cookie_jar.update_cookies(cookies) + + self._connector = connector + self._connector_owner = connector_owner + self._default_auth = auth + self._version = version + self._json_serialize = json_serialize + self._raise_for_status = raise_for_status + self._auto_decompress = auto_decompress + self._trust_env = trust_env + self._requote_redirect_url = requote_redirect_url + self._read_bufsize = read_bufsize + self._max_line_size = max_line_size + self._max_field_size = max_field_size + + # Convert to list of tuples + if headers: + real_headers: CIMultiDict[str] = CIMultiDict(headers) + else: + real_headers = CIMultiDict() + self._default_headers: CIMultiDict[str] = real_headers + if skip_auto_headers is not None: + self._skip_auto_headers = frozenset(istr(i) for i in skip_auto_headers) + else: + self._skip_auto_headers = frozenset() + + self._request_class = request_class + self._response_class = response_class + self._ws_response_class = ws_response_class + + self._trace_configs = trace_configs or [] + for trace_config in self._trace_configs: + trace_config.freeze() + + self._resolve_charset = fallback_charset_resolver + + self._default_proxy = proxy + self._default_proxy_auth = proxy_auth + self._retry_connection: bool = True + self._middlewares = middlewares + + def __init_subclass__(cls: Type["ClientSession"]) -> None: + warnings.warn( + "Inheritance class {} from ClientSession " + "is discouraged".format(cls.__name__), + DeprecationWarning, + stacklevel=2, + ) + + if DEBUG: + + def __setattr__(self, name: str, val: Any) -> None: + if name not in self.ATTRS: + warnings.warn( + "Setting custom ClientSession.{} attribute " + "is discouraged".format(name), + DeprecationWarning, + stacklevel=2, + ) + super().__setattr__(name, val) + + def __del__(self, _warnings: Any = warnings) -> None: + if not self.closed: + kwargs = {"source": self} + _warnings.warn( + f"Unclosed client session {self!r}", ResourceWarning, **kwargs + ) + context = {"client_session": self, "message": "Unclosed client session"} + if self._source_traceback is not None: + context["source_traceback"] = self._source_traceback + self._loop.call_exception_handler(context) + + if sys.version_info >= (3, 11) and TYPE_CHECKING: + + def request( + self, + method: str, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + else: + + def request( + self, method: str, url: StrOrURL, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP request.""" + return _RequestContextManager(self._request(method, url, **kwargs)) + + def _build_url(self, str_or_url: StrOrURL) -> URL: + url = URL(str_or_url) + if self._base_url and not url.absolute: + return self._base_url.join(url) + return url + + async def _request( + self, + method: str, + str_or_url: StrOrURL, + *, + params: Query = None, + data: Any = None, + json: Any = None, + cookies: Optional[LooseCookies] = None, + headers: Optional[LooseHeaders] = None, + skip_auto_headers: Optional[Iterable[str]] = None, + auth: Optional[BasicAuth] = None, + allow_redirects: bool = True, + max_redirects: int = 10, + compress: Union[str, bool, None] = None, + chunked: Optional[bool] = None, + expect100: bool = False, + raise_for_status: Union[ + None, bool, Callable[[ClientResponse], Awaitable[None]] + ] = None, + read_until_eof: bool = True, + proxy: Optional[StrOrURL] = None, + proxy_auth: Optional[BasicAuth] = None, + timeout: Union[ClientTimeout, _SENTINEL] = sentinel, + verify_ssl: Optional[bool] = None, + fingerprint: Optional[bytes] = None, + ssl_context: Optional[SSLContext] = None, + ssl: Union[SSLContext, bool, Fingerprint] = True, + server_hostname: Optional[str] = None, + proxy_headers: Optional[LooseHeaders] = None, + trace_request_ctx: Optional[Mapping[str, Any]] = None, + read_bufsize: Optional[int] = None, + auto_decompress: Optional[bool] = None, + max_line_size: Optional[int] = None, + max_field_size: Optional[int] = None, + middlewares: Optional[Sequence[ClientMiddlewareType]] = None, + ) -> ClientResponse: + + # NOTE: timeout clamps existing connect and read timeouts. We cannot + # set the default to None because we need to detect if the user wants + # to use the existing timeouts by setting timeout to None. + + if self.closed: + raise RuntimeError("Session is closed") + + ssl = _merge_ssl_params(ssl, verify_ssl, ssl_context, fingerprint) + + if data is not None and json is not None: + raise ValueError( + "data and json parameters can not be used at the same time" + ) + elif json is not None: + data = payload.JsonPayload(json, dumps=self._json_serialize) + + if not isinstance(chunked, bool) and chunked is not None: + warnings.warn("Chunk size is deprecated #1615", DeprecationWarning) + + redirects = 0 + history: List[ClientResponse] = [] + version = self._version + params = params or {} + + # Merge with default headers and transform to CIMultiDict + headers = self._prepare_headers(headers) + + try: + url = self._build_url(str_or_url) + except ValueError as e: + raise InvalidUrlClientError(str_or_url) from e + + assert self._connector is not None + if url.scheme not in self._connector.allowed_protocol_schema_set: + raise NonHttpUrlClientError(url) + + skip_headers: Optional[Iterable[istr]] + if skip_auto_headers is not None: + skip_headers = { + istr(i) for i in skip_auto_headers + } | self._skip_auto_headers + elif self._skip_auto_headers: + skip_headers = self._skip_auto_headers + else: + skip_headers = None + + if proxy is None: + proxy = self._default_proxy + if proxy_auth is None: + proxy_auth = self._default_proxy_auth + + if proxy is None: + proxy_headers = None + else: + proxy_headers = self._prepare_headers(proxy_headers) + try: + proxy = URL(proxy) + except ValueError as e: + raise InvalidURL(proxy) from e + + if timeout is sentinel: + real_timeout: ClientTimeout = self._timeout + else: + if not isinstance(timeout, ClientTimeout): + real_timeout = ClientTimeout(total=timeout) + else: + real_timeout = timeout + # timeout is cumulative for all request operations + # (request, redirects, responses, data consuming) + tm = TimeoutHandle( + self._loop, real_timeout.total, ceil_threshold=real_timeout.ceil_threshold + ) + handle = tm.start() + + if read_bufsize is None: + read_bufsize = self._read_bufsize + + if auto_decompress is None: + auto_decompress = self._auto_decompress + + if max_line_size is None: + max_line_size = self._max_line_size + + if max_field_size is None: + max_field_size = self._max_field_size + + traces = [ + Trace( + self, + trace_config, + trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx), + ) + for trace_config in self._trace_configs + ] + + for trace in traces: + await trace.send_request_start(method, url.update_query(params), headers) + + timer = tm.timer() + try: + with timer: + # https://www.rfc-editor.org/rfc/rfc9112.html#name-retrying-requests + retry_persistent_connection = ( + self._retry_connection and method in IDEMPOTENT_METHODS + ) + while True: + url, auth_from_url = strip_auth_from_url(url) + if not url.raw_host: + # NOTE: Bail early, otherwise, causes `InvalidURL` through + # NOTE: `self._request_class()` below. + err_exc_cls = ( + InvalidUrlRedirectClientError + if redirects + else InvalidUrlClientError + ) + raise err_exc_cls(url) + # If `auth` was passed for an already authenticated URL, + # disallow only if this is the initial URL; this is to avoid issues + # with sketchy redirects that are not the caller's responsibility + if not history and (auth and auth_from_url): + raise ValueError( + "Cannot combine AUTH argument with " + "credentials encoded in URL" + ) + + # Override the auth with the one from the URL only if we + # have no auth, or if we got an auth from a redirect URL + if auth is None or (history and auth_from_url is not None): + auth = auth_from_url + + if ( + auth is None + and self._default_auth + and ( + not self._base_url or self._base_url_origin == url.origin() + ) + ): + auth = self._default_auth + + # Try netrc if auth is still None and trust_env is enabled. + if auth is None and self._trust_env and url.host is not None: + auth = await self._loop.run_in_executor( + None, self._get_netrc_auth, url.host + ) + + # It would be confusing if we support explicit + # Authorization header with auth argument + if ( + headers is not None + and auth is not None + and hdrs.AUTHORIZATION in headers + ): + raise ValueError( + "Cannot combine AUTHORIZATION header " + "with AUTH argument or credentials " + "encoded in URL" + ) + + all_cookies = self._cookie_jar.filter_cookies(url) + + if cookies is not None: + tmp_cookie_jar = CookieJar( + quote_cookie=self._cookie_jar.quote_cookie + ) + tmp_cookie_jar.update_cookies(cookies) + req_cookies = tmp_cookie_jar.filter_cookies(url) + if req_cookies: + all_cookies.load(req_cookies) + + proxy_: Optional[URL] = None + if proxy is not None: + proxy_ = URL(proxy) + elif self._trust_env: + with suppress(LookupError): + proxy_, proxy_auth = await asyncio.to_thread( + get_env_proxy_for_url, url + ) + + req = self._request_class( + method, + url, + params=params, + headers=headers, + skip_auto_headers=skip_headers, + data=data, + cookies=all_cookies, + auth=auth, + version=version, + compress=compress, + chunked=chunked, + expect100=expect100, + loop=self._loop, + response_class=self._response_class, + proxy=proxy_, + proxy_auth=proxy_auth, + timer=timer, + session=self, + ssl=ssl if ssl is not None else True, + server_hostname=server_hostname, + proxy_headers=proxy_headers, + traces=traces, + trust_env=self.trust_env, + ) + + async def _connect_and_send_request( + req: ClientRequest, + ) -> ClientResponse: + # connection timeout + assert self._connector is not None + try: + conn = await self._connector.connect( + req, traces=traces, timeout=real_timeout + ) + except asyncio.TimeoutError as exc: + raise ConnectionTimeoutError( + f"Connection timeout to host {req.url}" + ) from exc + + assert conn.protocol is not None + conn.protocol.set_response_params( + timer=timer, + skip_payload=req.method in EMPTY_BODY_METHODS, + read_until_eof=read_until_eof, + auto_decompress=auto_decompress, + read_timeout=real_timeout.sock_read, + read_bufsize=read_bufsize, + timeout_ceil_threshold=self._connector._timeout_ceil_threshold, + max_line_size=max_line_size, + max_field_size=max_field_size, + ) + try: + resp = await req.send(conn) + try: + await resp.start(conn) + except BaseException: + resp.close() + raise + except BaseException: + conn.close() + raise + return resp + + # Apply middleware (if any) - per-request middleware overrides session middleware + effective_middlewares = ( + self._middlewares if middlewares is None else middlewares + ) + + if effective_middlewares: + handler = build_client_middlewares( + _connect_and_send_request, effective_middlewares + ) + else: + handler = _connect_and_send_request + + try: + resp = await handler(req) + # Client connector errors should not be retried + except ( + ConnectionTimeoutError, + ClientConnectorError, + ClientConnectorCertificateError, + ClientConnectorSSLError, + ): + raise + except (ClientOSError, ServerDisconnectedError): + if retry_persistent_connection: + retry_persistent_connection = False + continue + raise + except ClientError: + raise + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise ClientOSError(*exc.args) from exc + + # Update cookies from raw headers to preserve duplicates + if resp._raw_cookie_headers: + self._cookie_jar.update_cookies_from_headers( + resp._raw_cookie_headers, resp.url + ) + + # redirects + if resp.status in (301, 302, 303, 307, 308) and allow_redirects: + + for trace in traces: + await trace.send_request_redirect( + method, url.update_query(params), headers, resp + ) + + redirects += 1 + history.append(resp) + if max_redirects and redirects >= max_redirects: + if req._body is not None: + await req._body.close() + resp.close() + raise TooManyRedirects( + history[0].request_info, tuple(history) + ) + + # For 301 and 302, mimic IE, now changed in RFC + # https://github.com/kennethreitz/requests/pull/269 + if (resp.status == 303 and resp.method != hdrs.METH_HEAD) or ( + resp.status in (301, 302) and resp.method == hdrs.METH_POST + ): + method = hdrs.METH_GET + data = None + if headers.get(hdrs.CONTENT_LENGTH): + headers.pop(hdrs.CONTENT_LENGTH) + else: + # For 307/308, always preserve the request body + # For 301/302 with non-POST methods, preserve the request body + # https://www.rfc-editor.org/rfc/rfc9110#section-15.4.3-3.1 + # Use the existing payload to avoid recreating it from a potentially consumed file + data = req._body + + r_url = resp.headers.get(hdrs.LOCATION) or resp.headers.get( + hdrs.URI + ) + if r_url is None: + # see github.com/aio-libs/aiohttp/issues/2022 + break + else: + # reading from correct redirection + # response is forbidden + resp.release() + + try: + parsed_redirect_url = URL( + r_url, encoded=not self._requote_redirect_url + ) + except ValueError as e: + if req._body is not None: + await req._body.close() + resp.close() + raise InvalidUrlRedirectClientError( + r_url, + "Server attempted redirecting to a location that does not look like a URL", + ) from e + + scheme = parsed_redirect_url.scheme + if scheme not in HTTP_AND_EMPTY_SCHEMA_SET: + if req._body is not None: + await req._body.close() + resp.close() + raise NonHttpUrlRedirectClientError(r_url) + elif not scheme: + parsed_redirect_url = url.join(parsed_redirect_url) + + try: + redirect_origin = parsed_redirect_url.origin() + except ValueError as origin_val_err: + if req._body is not None: + await req._body.close() + resp.close() + raise InvalidUrlRedirectClientError( + parsed_redirect_url, + "Invalid redirect URL origin", + ) from origin_val_err + + if url.origin() != redirect_origin: + auth = None + headers.pop(hdrs.AUTHORIZATION, None) + + url = parsed_redirect_url + params = {} + resp.release() + continue + + break + + if req._body is not None: + await req._body.close() + # check response status + if raise_for_status is None: + raise_for_status = self._raise_for_status + + if raise_for_status is None: + pass + elif callable(raise_for_status): + await raise_for_status(resp) + elif raise_for_status: + resp.raise_for_status() + + # register connection + if handle is not None: + if resp.connection is not None: + resp.connection.add_callback(handle.cancel) + else: + handle.cancel() + + resp._history = tuple(history) + + for trace in traces: + await trace.send_request_end( + method, url.update_query(params), headers, resp + ) + return resp + + except BaseException as e: + # cleanup timer + tm.close() + if handle: + handle.cancel() + handle = None + + for trace in traces: + await trace.send_request_exception( + method, url.update_query(params), headers, e + ) + raise + + def ws_connect( + self, + url: StrOrURL, + *, + method: str = hdrs.METH_GET, + protocols: Iterable[str] = (), + timeout: Union[ClientWSTimeout, _SENTINEL] = sentinel, + receive_timeout: Optional[float] = None, + autoclose: bool = True, + autoping: bool = True, + heartbeat: Optional[float] = None, + auth: Optional[BasicAuth] = None, + origin: Optional[str] = None, + params: Query = None, + headers: Optional[LooseHeaders] = None, + proxy: Optional[StrOrURL] = None, + proxy_auth: Optional[BasicAuth] = None, + ssl: Union[SSLContext, bool, Fingerprint] = True, + verify_ssl: Optional[bool] = None, + fingerprint: Optional[bytes] = None, + ssl_context: Optional[SSLContext] = None, + server_hostname: Optional[str] = None, + proxy_headers: Optional[LooseHeaders] = None, + compress: int = 0, + max_msg_size: int = 4 * 1024 * 1024, + ) -> "_WSRequestContextManager": + """Initiate websocket connection.""" + return _WSRequestContextManager( + self._ws_connect( + url, + method=method, + protocols=protocols, + timeout=timeout, + receive_timeout=receive_timeout, + autoclose=autoclose, + autoping=autoping, + heartbeat=heartbeat, + auth=auth, + origin=origin, + params=params, + headers=headers, + proxy=proxy, + proxy_auth=proxy_auth, + ssl=ssl, + verify_ssl=verify_ssl, + fingerprint=fingerprint, + ssl_context=ssl_context, + server_hostname=server_hostname, + proxy_headers=proxy_headers, + compress=compress, + max_msg_size=max_msg_size, + ) + ) + + async def _ws_connect( + self, + url: StrOrURL, + *, + method: str = hdrs.METH_GET, + protocols: Iterable[str] = (), + timeout: Union[ClientWSTimeout, _SENTINEL] = sentinel, + receive_timeout: Optional[float] = None, + autoclose: bool = True, + autoping: bool = True, + heartbeat: Optional[float] = None, + auth: Optional[BasicAuth] = None, + origin: Optional[str] = None, + params: Query = None, + headers: Optional[LooseHeaders] = None, + proxy: Optional[StrOrURL] = None, + proxy_auth: Optional[BasicAuth] = None, + ssl: Union[SSLContext, bool, Fingerprint] = True, + verify_ssl: Optional[bool] = None, + fingerprint: Optional[bytes] = None, + ssl_context: Optional[SSLContext] = None, + server_hostname: Optional[str] = None, + proxy_headers: Optional[LooseHeaders] = None, + compress: int = 0, + max_msg_size: int = 4 * 1024 * 1024, + ) -> ClientWebSocketResponse: + if timeout is not sentinel: + if isinstance(timeout, ClientWSTimeout): + ws_timeout = timeout + else: + warnings.warn( + "parameter 'timeout' of type 'float' " + "is deprecated, please use " + "'timeout=ClientWSTimeout(ws_close=...)'", + DeprecationWarning, + stacklevel=2, + ) + ws_timeout = ClientWSTimeout(ws_close=timeout) + else: + ws_timeout = DEFAULT_WS_CLIENT_TIMEOUT + if receive_timeout is not None: + warnings.warn( + "float parameter 'receive_timeout' " + "is deprecated, please use parameter " + "'timeout=ClientWSTimeout(ws_receive=...)'", + DeprecationWarning, + stacklevel=2, + ) + ws_timeout = attr.evolve(ws_timeout, ws_receive=receive_timeout) + + if headers is None: + real_headers: CIMultiDict[str] = CIMultiDict() + else: + real_headers = CIMultiDict(headers) + + default_headers = { + hdrs.UPGRADE: "websocket", + hdrs.CONNECTION: "Upgrade", + hdrs.SEC_WEBSOCKET_VERSION: "13", + } + + for key, value in default_headers.items(): + real_headers.setdefault(key, value) + + sec_key = base64.b64encode(os.urandom(16)) + real_headers[hdrs.SEC_WEBSOCKET_KEY] = sec_key.decode() + + if protocols: + real_headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ",".join(protocols) + if origin is not None: + real_headers[hdrs.ORIGIN] = origin + if compress: + extstr = ws_ext_gen(compress=compress) + real_headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = extstr + + # For the sake of backward compatibility, if user passes in None, convert it to True + if ssl is None: + warnings.warn( + "ssl=None is deprecated, please use ssl=True", + DeprecationWarning, + stacklevel=2, + ) + ssl = True + ssl = _merge_ssl_params(ssl, verify_ssl, ssl_context, fingerprint) + + # send request + resp = await self.request( + method, + url, + params=params, + headers=real_headers, + read_until_eof=False, + auth=auth, + proxy=proxy, + proxy_auth=proxy_auth, + ssl=ssl, + server_hostname=server_hostname, + proxy_headers=proxy_headers, + ) + + try: + # check handshake + if resp.status != 101: + raise WSServerHandshakeError( + resp.request_info, + resp.history, + message="Invalid response status", + status=resp.status, + headers=resp.headers, + ) + + if resp.headers.get(hdrs.UPGRADE, "").lower() != "websocket": + raise WSServerHandshakeError( + resp.request_info, + resp.history, + message="Invalid upgrade header", + status=resp.status, + headers=resp.headers, + ) + + if resp.headers.get(hdrs.CONNECTION, "").lower() != "upgrade": + raise WSServerHandshakeError( + resp.request_info, + resp.history, + message="Invalid connection header", + status=resp.status, + headers=resp.headers, + ) + + # key calculation + r_key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, "") + match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()).decode() + if r_key != match: + raise WSServerHandshakeError( + resp.request_info, + resp.history, + message="Invalid challenge response", + status=resp.status, + headers=resp.headers, + ) + + # websocket protocol + protocol = None + if protocols and hdrs.SEC_WEBSOCKET_PROTOCOL in resp.headers: + resp_protocols = [ + proto.strip() + for proto in resp.headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(",") + ] + + for proto in resp_protocols: + if proto in protocols: + protocol = proto + break + + # websocket compress + notakeover = False + if compress: + compress_hdrs = resp.headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS) + if compress_hdrs: + try: + compress, notakeover = ws_ext_parse(compress_hdrs) + except WSHandshakeError as exc: + raise WSServerHandshakeError( + resp.request_info, + resp.history, + message=exc.args[0], + status=resp.status, + headers=resp.headers, + ) from exc + else: + compress = 0 + notakeover = False + + conn = resp.connection + assert conn is not None + conn_proto = conn.protocol + assert conn_proto is not None + + # For WS connection the read_timeout must be either receive_timeout or greater + # None == no timeout, i.e. infinite timeout, so None is the max timeout possible + if ws_timeout.ws_receive is None: + # Reset regardless + conn_proto.read_timeout = None + elif conn_proto.read_timeout is not None: + conn_proto.read_timeout = max( + ws_timeout.ws_receive, conn_proto.read_timeout + ) + + transport = conn.transport + assert transport is not None + reader = WebSocketDataQueue(conn_proto, 2**16, loop=self._loop) + conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader) + writer = WebSocketWriter( + conn_proto, + transport, + use_mask=True, + compress=compress, + notakeover=notakeover, + ) + except BaseException: + resp.close() + raise + else: + return self._ws_response_class( + reader, + writer, + protocol, + resp, + ws_timeout, + autoclose, + autoping, + self._loop, + heartbeat=heartbeat, + compress=compress, + client_notakeover=notakeover, + ) + + def _prepare_headers(self, headers: Optional[LooseHeaders]) -> "CIMultiDict[str]": + """Add default headers and transform it to CIMultiDict""" + # Convert headers to MultiDict + result = CIMultiDict(self._default_headers) + if headers: + if not isinstance(headers, (MultiDictProxy, MultiDict)): + headers = CIMultiDict(headers) + added_names: Set[str] = set() + for key, value in headers.items(): + if key in added_names: + result.add(key, value) + else: + result[key] = value + added_names.add(key) + return result + + def _get_netrc_auth(self, host: str) -> Optional[BasicAuth]: + """ + Get auth from netrc for the given host. + + This method is designed to be called in an executor to avoid + blocking I/O in the event loop. + """ + netrc_obj = netrc_from_env() + try: + return basicauth_from_netrc(netrc_obj, host) + except LookupError: + return None + + if sys.version_info >= (3, 11) and TYPE_CHECKING: + + def get( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def options( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def head( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def post( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def put( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def patch( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + def delete( + self, + url: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> "_RequestContextManager": ... + + else: + + def get( + self, url: StrOrURL, *, allow_redirects: bool = True, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP GET request.""" + return _RequestContextManager( + self._request( + hdrs.METH_GET, url, allow_redirects=allow_redirects, **kwargs + ) + ) + + def options( + self, url: StrOrURL, *, allow_redirects: bool = True, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP OPTIONS request.""" + return _RequestContextManager( + self._request( + hdrs.METH_OPTIONS, url, allow_redirects=allow_redirects, **kwargs + ) + ) + + def head( + self, url: StrOrURL, *, allow_redirects: bool = False, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP HEAD request.""" + return _RequestContextManager( + self._request( + hdrs.METH_HEAD, url, allow_redirects=allow_redirects, **kwargs + ) + ) + + def post( + self, url: StrOrURL, *, data: Any = None, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP POST request.""" + return _RequestContextManager( + self._request(hdrs.METH_POST, url, data=data, **kwargs) + ) + + def put( + self, url: StrOrURL, *, data: Any = None, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP PUT request.""" + return _RequestContextManager( + self._request(hdrs.METH_PUT, url, data=data, **kwargs) + ) + + def patch( + self, url: StrOrURL, *, data: Any = None, **kwargs: Any + ) -> "_RequestContextManager": + """Perform HTTP PATCH request.""" + return _RequestContextManager( + self._request(hdrs.METH_PATCH, url, data=data, **kwargs) + ) + + def delete(self, url: StrOrURL, **kwargs: Any) -> "_RequestContextManager": + """Perform HTTP DELETE request.""" + return _RequestContextManager( + self._request(hdrs.METH_DELETE, url, **kwargs) + ) + + async def close(self) -> None: + """Close underlying connector. + + Release all acquired resources. + """ + if not self.closed: + if self._connector is not None and self._connector_owner: + await self._connector.close() + self._connector = None + + @property + def closed(self) -> bool: + """Is client session closed. + + A readonly property. + """ + return self._connector is None or self._connector.closed + + @property + def connector(self) -> Optional[BaseConnector]: + """Connector instance used for the session.""" + return self._connector + + @property + def cookie_jar(self) -> AbstractCookieJar: + """The session cookies.""" + return self._cookie_jar + + @property + def version(self) -> Tuple[int, int]: + """The session HTTP protocol version.""" + return self._version + + @property + def requote_redirect_url(self) -> bool: + """Do URL requoting on redirection handling.""" + return self._requote_redirect_url + + @requote_redirect_url.setter + def requote_redirect_url(self, val: bool) -> None: + """Do URL requoting on redirection handling.""" + warnings.warn( + "session.requote_redirect_url modification is deprecated #2778", + DeprecationWarning, + stacklevel=2, + ) + self._requote_redirect_url = val + + @property + def loop(self) -> asyncio.AbstractEventLoop: + """Session's loop.""" + warnings.warn( + "client.loop property is deprecated", DeprecationWarning, stacklevel=2 + ) + return self._loop + + @property + def timeout(self) -> ClientTimeout: + """Timeout for the session.""" + return self._timeout + + @property + def headers(self) -> "CIMultiDict[str]": + """The default headers of the client session.""" + return self._default_headers + + @property + def skip_auto_headers(self) -> FrozenSet[istr]: + """Headers for which autogeneration should be skipped""" + return self._skip_auto_headers + + @property + def auth(self) -> Optional[BasicAuth]: + """An object that represents HTTP Basic Authorization""" + return self._default_auth + + @property + def json_serialize(self) -> JSONEncoder: + """Json serializer callable""" + return self._json_serialize + + @property + def connector_owner(self) -> bool: + """Should connector be closed on session closing""" + return self._connector_owner + + @property + def raise_for_status( + self, + ) -> Union[bool, Callable[[ClientResponse], Awaitable[None]]]: + """Should `ClientResponse.raise_for_status()` be called for each response.""" + return self._raise_for_status + + @property + def auto_decompress(self) -> bool: + """Should the body response be automatically decompressed.""" + return self._auto_decompress + + @property + def trust_env(self) -> bool: + """ + Should proxies information from environment or netrc be trusted. + + Information is from HTTP_PROXY / HTTPS_PROXY environment variables + or ~/.netrc file if present. + """ + return self._trust_env + + @property + def trace_configs(self) -> List[TraceConfig]: + """A list of TraceConfig instances used for client tracing""" + return self._trace_configs + + def detach(self) -> None: + """Detach connector from session without closing the former. + + Session is switched to closed state anyway. + """ + self._connector = None + + def __enter__(self) -> None: + raise TypeError("Use async with instead") + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + # __exit__ should exist in pair with __enter__ but never executed + pass # pragma: no cover + + async def __aenter__(self) -> "ClientSession": + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + await self.close() + + +class _BaseRequestContextManager(Coroutine[Any, Any, _RetType], Generic[_RetType]): + + __slots__ = ("_coro", "_resp") + + def __init__(self, coro: Coroutine["asyncio.Future[Any]", None, _RetType]) -> None: + self._coro: Coroutine["asyncio.Future[Any]", None, _RetType] = coro + + def send(self, arg: None) -> "asyncio.Future[Any]": + return self._coro.send(arg) + + def throw(self, *args: Any, **kwargs: Any) -> "asyncio.Future[Any]": + return self._coro.throw(*args, **kwargs) + + def close(self) -> None: + return self._coro.close() + + def __await__(self) -> Generator[Any, None, _RetType]: + ret = self._coro.__await__() + return ret + + def __iter__(self) -> Generator[Any, None, _RetType]: + return self.__await__() + + async def __aenter__(self) -> _RetType: + self._resp: _RetType = await self._coro + return await self._resp.__aenter__() + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> None: + await self._resp.__aexit__(exc_type, exc, tb) + + +_RequestContextManager = _BaseRequestContextManager[ClientResponse] +_WSRequestContextManager = _BaseRequestContextManager[ClientWebSocketResponse] + + +class _SessionRequestContextManager: + + __slots__ = ("_coro", "_resp", "_session") + + def __init__( + self, + coro: Coroutine["asyncio.Future[Any]", None, ClientResponse], + session: ClientSession, + ) -> None: + self._coro = coro + self._resp: Optional[ClientResponse] = None + self._session = session + + async def __aenter__(self) -> ClientResponse: + try: + self._resp = await self._coro + except BaseException: + await self._session.close() + raise + else: + return self._resp + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> None: + assert self._resp is not None + self._resp.close() + await self._session.close() + + +if sys.version_info >= (3, 11) and TYPE_CHECKING: + + def request( + method: str, + url: StrOrURL, + *, + version: HttpVersion = http.HttpVersion11, + connector: Optional[BaseConnector] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + **kwargs: Unpack[_RequestOptions], + ) -> _SessionRequestContextManager: ... + +else: + + def request( + method: str, + url: StrOrURL, + *, + version: HttpVersion = http.HttpVersion11, + connector: Optional[BaseConnector] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + **kwargs: Any, + ) -> _SessionRequestContextManager: + """Constructs and sends a request. + + Returns response object. + method - HTTP method + url - request url + params - (optional) Dictionary or bytes to be sent in the query + string of the new request + data - (optional) Dictionary, bytes, or file-like object to + send in the body of the request + json - (optional) Any json compatible python object + headers - (optional) Dictionary of HTTP Headers to send with + the request + cookies - (optional) Dict object to send with the request + auth - (optional) BasicAuth named tuple represent HTTP Basic Auth + auth - aiohttp.helpers.BasicAuth + allow_redirects - (optional) If set to False, do not follow + redirects + version - Request HTTP version. + compress - Set to True if request has to be compressed + with deflate encoding. + chunked - Set to chunk size for chunked transfer encoding. + expect100 - Expect 100-continue response from server. + connector - BaseConnector sub-class instance to support + connection pooling. + read_until_eof - Read response until eof if response + does not have Content-Length header. + loop - Optional event loop. + timeout - Optional ClientTimeout settings structure, 5min + total timeout by default. + Usage:: + >>> import aiohttp + >>> async with aiohttp.request('GET', 'http://python.org/') as resp: + ... print(resp) + ... data = await resp.read() + + """ + connector_owner = False + if connector is None: + connector_owner = True + connector = TCPConnector(loop=loop, force_close=True) + + session = ClientSession( + loop=loop, + cookies=kwargs.pop("cookies", None), + version=version, + timeout=kwargs.pop("timeout", sentinel), + connector=connector, + connector_owner=connector_owner, + ) + + return _SessionRequestContextManager( + session._request(method, url, **kwargs), + session, + ) diff --git a/lib/python3.12/site-packages/aiohttp/client_exceptions.py b/lib/python3.12/site-packages/aiohttp/client_exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..1d298e9a8cf663cdc8a85d3b7d1f9264ff5e03c9 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_exceptions.py @@ -0,0 +1,421 @@ +"""HTTP related errors.""" + +import asyncio +import warnings +from typing import TYPE_CHECKING, Optional, Tuple, Union + +from multidict import MultiMapping + +from .typedefs import StrOrURL + +if TYPE_CHECKING: + import ssl + + SSLContext = ssl.SSLContext +else: + try: + import ssl + + SSLContext = ssl.SSLContext + except ImportError: # pragma: no cover + ssl = SSLContext = None # type: ignore[assignment] + +if TYPE_CHECKING: + from .client_reqrep import ClientResponse, ConnectionKey, Fingerprint, RequestInfo + from .http_parser import RawResponseMessage +else: + RequestInfo = ClientResponse = ConnectionKey = RawResponseMessage = None + +__all__ = ( + "ClientError", + "ClientConnectionError", + "ClientConnectionResetError", + "ClientOSError", + "ClientConnectorError", + "ClientProxyConnectionError", + "ClientSSLError", + "ClientConnectorDNSError", + "ClientConnectorSSLError", + "ClientConnectorCertificateError", + "ConnectionTimeoutError", + "SocketTimeoutError", + "ServerConnectionError", + "ServerTimeoutError", + "ServerDisconnectedError", + "ServerFingerprintMismatch", + "ClientResponseError", + "ClientHttpProxyError", + "WSServerHandshakeError", + "ContentTypeError", + "ClientPayloadError", + "InvalidURL", + "InvalidUrlClientError", + "RedirectClientError", + "NonHttpUrlClientError", + "InvalidUrlRedirectClientError", + "NonHttpUrlRedirectClientError", + "WSMessageTypeError", +) + + +class ClientError(Exception): + """Base class for client connection errors.""" + + +class ClientResponseError(ClientError): + """Base class for exceptions that occur after getting a response. + + request_info: An instance of RequestInfo. + history: A sequence of responses, if redirects occurred. + status: HTTP status code. + message: Error message. + headers: Response headers. + """ + + def __init__( + self, + request_info: RequestInfo, + history: Tuple[ClientResponse, ...], + *, + code: Optional[int] = None, + status: Optional[int] = None, + message: str = "", + headers: Optional[MultiMapping[str]] = None, + ) -> None: + self.request_info = request_info + if code is not None: + if status is not None: + raise ValueError( + "Both code and status arguments are provided; " + "code is deprecated, use status instead" + ) + warnings.warn( + "code argument is deprecated, use status instead", + DeprecationWarning, + stacklevel=2, + ) + if status is not None: + self.status = status + elif code is not None: + self.status = code + else: + self.status = 0 + self.message = message + self.headers = headers + self.history = history + self.args = (request_info, history) + + def __str__(self) -> str: + return "{}, message={!r}, url={!r}".format( + self.status, + self.message, + str(self.request_info.real_url), + ) + + def __repr__(self) -> str: + args = f"{self.request_info!r}, {self.history!r}" + if self.status != 0: + args += f", status={self.status!r}" + if self.message != "": + args += f", message={self.message!r}" + if self.headers is not None: + args += f", headers={self.headers!r}" + return f"{type(self).__name__}({args})" + + @property + def code(self) -> int: + warnings.warn( + "code property is deprecated, use status instead", + DeprecationWarning, + stacklevel=2, + ) + return self.status + + @code.setter + def code(self, value: int) -> None: + warnings.warn( + "code property is deprecated, use status instead", + DeprecationWarning, + stacklevel=2, + ) + self.status = value + + +class ContentTypeError(ClientResponseError): + """ContentType found is not valid.""" + + +class WSServerHandshakeError(ClientResponseError): + """websocket server handshake error.""" + + +class ClientHttpProxyError(ClientResponseError): + """HTTP proxy error. + + Raised in :class:`aiohttp.connector.TCPConnector` if + proxy responds with status other than ``200 OK`` + on ``CONNECT`` request. + """ + + +class TooManyRedirects(ClientResponseError): + """Client was redirected too many times.""" + + +class ClientConnectionError(ClientError): + """Base class for client socket errors.""" + + +class ClientConnectionResetError(ClientConnectionError, ConnectionResetError): + """ConnectionResetError""" + + +class ClientOSError(ClientConnectionError, OSError): + """OSError error.""" + + +class ClientConnectorError(ClientOSError): + """Client connector error. + + Raised in :class:`aiohttp.connector.TCPConnector` if + a connection can not be established. + """ + + def __init__(self, connection_key: ConnectionKey, os_error: OSError) -> None: + self._conn_key = connection_key + self._os_error = os_error + super().__init__(os_error.errno, os_error.strerror) + self.args = (connection_key, os_error) + + @property + def os_error(self) -> OSError: + return self._os_error + + @property + def host(self) -> str: + return self._conn_key.host + + @property + def port(self) -> Optional[int]: + return self._conn_key.port + + @property + def ssl(self) -> Union[SSLContext, bool, "Fingerprint"]: + return self._conn_key.ssl + + def __str__(self) -> str: + return "Cannot connect to host {0.host}:{0.port} ssl:{1} [{2}]".format( + self, "default" if self.ssl is True else self.ssl, self.strerror + ) + + # OSError.__reduce__ does too much black magick + __reduce__ = BaseException.__reduce__ + + +class ClientConnectorDNSError(ClientConnectorError): + """DNS resolution failed during client connection. + + Raised in :class:`aiohttp.connector.TCPConnector` if + DNS resolution fails. + """ + + +class ClientProxyConnectionError(ClientConnectorError): + """Proxy connection error. + + Raised in :class:`aiohttp.connector.TCPConnector` if + connection to proxy can not be established. + """ + + +class UnixClientConnectorError(ClientConnectorError): + """Unix connector error. + + Raised in :py:class:`aiohttp.connector.UnixConnector` + if connection to unix socket can not be established. + """ + + def __init__( + self, path: str, connection_key: ConnectionKey, os_error: OSError + ) -> None: + self._path = path + super().__init__(connection_key, os_error) + + @property + def path(self) -> str: + return self._path + + def __str__(self) -> str: + return "Cannot connect to unix socket {0.path} ssl:{1} [{2}]".format( + self, "default" if self.ssl is True else self.ssl, self.strerror + ) + + +class ServerConnectionError(ClientConnectionError): + """Server connection errors.""" + + +class ServerDisconnectedError(ServerConnectionError): + """Server disconnected.""" + + def __init__(self, message: Union[RawResponseMessage, str, None] = None) -> None: + if message is None: + message = "Server disconnected" + + self.args = (message,) + self.message = message + + +class ServerTimeoutError(ServerConnectionError, asyncio.TimeoutError): + """Server timeout error.""" + + +class ConnectionTimeoutError(ServerTimeoutError): + """Connection timeout error.""" + + +class SocketTimeoutError(ServerTimeoutError): + """Socket timeout error.""" + + +class ServerFingerprintMismatch(ServerConnectionError): + """SSL certificate does not match expected fingerprint.""" + + def __init__(self, expected: bytes, got: bytes, host: str, port: int) -> None: + self.expected = expected + self.got = got + self.host = host + self.port = port + self.args = (expected, got, host, port) + + def __repr__(self) -> str: + return "<{} expected={!r} got={!r} host={!r} port={!r}>".format( + self.__class__.__name__, self.expected, self.got, self.host, self.port + ) + + +class ClientPayloadError(ClientError): + """Response payload error.""" + + +class InvalidURL(ClientError, ValueError): + """Invalid URL. + + URL used for fetching is malformed, e.g. it doesn't contains host + part. + """ + + # Derive from ValueError for backward compatibility + + def __init__(self, url: StrOrURL, description: Union[str, None] = None) -> None: + # The type of url is not yarl.URL because the exception can be raised + # on URL(url) call + self._url = url + self._description = description + + if description: + super().__init__(url, description) + else: + super().__init__(url) + + @property + def url(self) -> StrOrURL: + return self._url + + @property + def description(self) -> "str | None": + return self._description + + def __repr__(self) -> str: + return f"<{self.__class__.__name__} {self}>" + + def __str__(self) -> str: + if self._description: + return f"{self._url} - {self._description}" + return str(self._url) + + +class InvalidUrlClientError(InvalidURL): + """Invalid URL client error.""" + + +class RedirectClientError(ClientError): + """Client redirect error.""" + + +class NonHttpUrlClientError(ClientError): + """Non http URL client error.""" + + +class InvalidUrlRedirectClientError(InvalidUrlClientError, RedirectClientError): + """Invalid URL redirect client error.""" + + +class NonHttpUrlRedirectClientError(NonHttpUrlClientError, RedirectClientError): + """Non http URL redirect client error.""" + + +class ClientSSLError(ClientConnectorError): + """Base error for ssl.*Errors.""" + + +if ssl is not None: + cert_errors = (ssl.CertificateError,) + cert_errors_bases = ( + ClientSSLError, + ssl.CertificateError, + ) + + ssl_errors = (ssl.SSLError,) + ssl_error_bases = (ClientSSLError, ssl.SSLError) +else: # pragma: no cover + cert_errors = tuple() + cert_errors_bases = ( + ClientSSLError, + ValueError, + ) + + ssl_errors = tuple() + ssl_error_bases = (ClientSSLError,) + + +class ClientConnectorSSLError(*ssl_error_bases): # type: ignore[misc] + """Response ssl error.""" + + +class ClientConnectorCertificateError(*cert_errors_bases): # type: ignore[misc] + """Response certificate error.""" + + def __init__( + self, connection_key: ConnectionKey, certificate_error: Exception + ) -> None: + self._conn_key = connection_key + self._certificate_error = certificate_error + self.args = (connection_key, certificate_error) + + @property + def certificate_error(self) -> Exception: + return self._certificate_error + + @property + def host(self) -> str: + return self._conn_key.host + + @property + def port(self) -> Optional[int]: + return self._conn_key.port + + @property + def ssl(self) -> bool: + return self._conn_key.is_ssl + + def __str__(self) -> str: + return ( + "Cannot connect to host {0.host}:{0.port} ssl:{0.ssl} " + "[{0.certificate_error.__class__.__name__}: " + "{0.certificate_error.args}]".format(self) + ) + + +class WSMessageTypeError(TypeError): + """WebSocket message type is not valid.""" diff --git a/lib/python3.12/site-packages/aiohttp/client_middleware_digest_auth.py b/lib/python3.12/site-packages/aiohttp/client_middleware_digest_auth.py new file mode 100644 index 0000000000000000000000000000000000000000..5aab5acb85acbda5136fe0bef84df7ff4f492405 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_middleware_digest_auth.py @@ -0,0 +1,480 @@ +""" +Digest authentication middleware for aiohttp client. + +This middleware implements HTTP Digest Authentication according to RFC 7616, +providing a more secure alternative to Basic Authentication. It supports all +standard hash algorithms including MD5, SHA, SHA-256, SHA-512 and their session +variants, as well as both 'auth' and 'auth-int' quality of protection (qop) options. +""" + +import hashlib +import os +import re +import sys +import time +from typing import ( + Callable, + Dict, + Final, + FrozenSet, + List, + Literal, + Tuple, + TypedDict, + Union, +) + +from yarl import URL + +from . import hdrs +from .client_exceptions import ClientError +from .client_middlewares import ClientHandlerType +from .client_reqrep import ClientRequest, ClientResponse +from .payload import Payload + + +class DigestAuthChallenge(TypedDict, total=False): + realm: str + nonce: str + qop: str + algorithm: str + opaque: str + domain: str + stale: str + + +DigestFunctions: Dict[str, Callable[[bytes], "hashlib._Hash"]] = { + "MD5": hashlib.md5, + "MD5-SESS": hashlib.md5, + "SHA": hashlib.sha1, + "SHA-SESS": hashlib.sha1, + "SHA256": hashlib.sha256, + "SHA256-SESS": hashlib.sha256, + "SHA-256": hashlib.sha256, + "SHA-256-SESS": hashlib.sha256, + "SHA512": hashlib.sha512, + "SHA512-SESS": hashlib.sha512, + "SHA-512": hashlib.sha512, + "SHA-512-SESS": hashlib.sha512, +} + + +# Compile the regex pattern once at module level for performance +_HEADER_PAIRS_PATTERN = re.compile( + r'(?:^|\s|,\s*)(\w+)\s*=\s*(?:"((?:[^"\\]|\\.)*)"|([^\s,]+))' + if sys.version_info < (3, 11) + else r'(?:^|\s|,\s*)((?>\w+))\s*=\s*(?:"((?:[^"\\]|\\.)*)"|([^\s,]+))' + # +------------|--------|--|-|-|--|----|------|----|--||-----|-> Match valid start/sep + # +--------|--|-|-|--|----|------|----|--||-----|-> alphanumeric key (atomic + # | | | | | | | | || | group reduces backtracking) + # +--|-|-|--|----|------|----|--||-----|-> maybe whitespace + # | | | | | | | || | + # +-|-|--|----|------|----|--||-----|-> = (delimiter) + # +-|--|----|------|----|--||-----|-> maybe whitespace + # | | | | | || | + # +--|----|------|----|--||-----|-> group quoted or unquoted + # | | | | || | + # +----|------|----|--||-----|-> if quoted... + # +------|----|--||-----|-> anything but " or \ + # +----|--||-----|-> escaped characters allowed + # +--||-----|-> or can be empty string + # || | + # +|-----|-> if unquoted... + # +-----|-> anything but , or + # +-> at least one char req'd +) + + +# RFC 7616: Challenge parameters to extract +CHALLENGE_FIELDS: Final[ + Tuple[ + Literal["realm", "nonce", "qop", "algorithm", "opaque", "domain", "stale"], ... + ] +] = ( + "realm", + "nonce", + "qop", + "algorithm", + "opaque", + "domain", + "stale", +) + +# Supported digest authentication algorithms +# Use a tuple of sorted keys for predictable documentation and error messages +SUPPORTED_ALGORITHMS: Final[Tuple[str, ...]] = tuple(sorted(DigestFunctions.keys())) + +# RFC 7616: Fields that require quoting in the Digest auth header +# These fields must be enclosed in double quotes in the Authorization header. +# Algorithm, qop, and nc are never quoted per RFC specifications. +# This frozen set is used by the template-based header construction to +# automatically determine which fields need quotes. +QUOTED_AUTH_FIELDS: Final[FrozenSet[str]] = frozenset( + {"username", "realm", "nonce", "uri", "response", "opaque", "cnonce"} +) + + +def escape_quotes(value: str) -> str: + """Escape double quotes for HTTP header values.""" + return value.replace('"', '\\"') + + +def unescape_quotes(value: str) -> str: + """Unescape double quotes in HTTP header values.""" + return value.replace('\\"', '"') + + +def parse_header_pairs(header: str) -> Dict[str, str]: + """ + Parse key-value pairs from WWW-Authenticate or similar HTTP headers. + + This function handles the complex format of WWW-Authenticate header values, + supporting both quoted and unquoted values, proper handling of commas in + quoted values, and whitespace variations per RFC 7616. + + Examples of supported formats: + - key1="value1", key2=value2 + - key1 = "value1" , key2="value, with, commas" + - key1=value1,key2="value2" + - realm="example.com", nonce="12345", qop="auth" + + Args: + header: The header value string to parse + + Returns: + Dictionary mapping parameter names to their values + """ + return { + stripped_key: unescape_quotes(quoted_val) if quoted_val else unquoted_val + for key, quoted_val, unquoted_val in _HEADER_PAIRS_PATTERN.findall(header) + if (stripped_key := key.strip()) + } + + +class DigestAuthMiddleware: + """ + HTTP digest authentication middleware for aiohttp client. + + This middleware intercepts 401 Unauthorized responses containing a Digest + authentication challenge, calculates the appropriate digest credentials, + and automatically retries the request with the proper Authorization header. + + Features: + - Handles all aspects of Digest authentication handshake automatically + - Supports all standard hash algorithms: + - MD5, MD5-SESS + - SHA, SHA-SESS + - SHA256, SHA256-SESS, SHA-256, SHA-256-SESS + - SHA512, SHA512-SESS, SHA-512, SHA-512-SESS + - Supports 'auth' and 'auth-int' quality of protection modes + - Properly handles quoted strings and parameter parsing + - Includes replay attack protection with client nonce count tracking + - Supports preemptive authentication per RFC 7616 Section 3.6 + + Standards compliance: + - RFC 7616: HTTP Digest Access Authentication (primary reference) + - RFC 2617: HTTP Authentication (deprecated by RFC 7616) + - RFC 1945: Section 11.1 (username restrictions) + + Implementation notes: + The core digest calculation is inspired by the implementation in + https://github.com/requests/requests/blob/v2.18.4/requests/auth.py + with added support for modern digest auth features and error handling. + """ + + def __init__( + self, + login: str, + password: str, + preemptive: bool = True, + ) -> None: + if login is None: + raise ValueError("None is not allowed as login value") + + if password is None: + raise ValueError("None is not allowed as password value") + + if ":" in login: + raise ValueError('A ":" is not allowed in username (RFC 1945#section-11.1)') + + self._login_str: Final[str] = login + self._login_bytes: Final[bytes] = login.encode("utf-8") + self._password_bytes: Final[bytes] = password.encode("utf-8") + + self._last_nonce_bytes = b"" + self._nonce_count = 0 + self._challenge: DigestAuthChallenge = {} + self._preemptive: bool = preemptive + # Set of URLs defining the protection space + self._protection_space: List[str] = [] + + async def _encode( + self, method: str, url: URL, body: Union[Payload, Literal[b""]] + ) -> str: + """ + Build digest authorization header for the current challenge. + + Args: + method: The HTTP method (GET, POST, etc.) + url: The request URL + body: The request body (used for qop=auth-int) + + Returns: + A fully formatted Digest authorization header string + + Raises: + ClientError: If the challenge is missing required parameters or + contains unsupported values + + """ + challenge = self._challenge + if "realm" not in challenge: + raise ClientError( + "Malformed Digest auth challenge: Missing 'realm' parameter" + ) + + if "nonce" not in challenge: + raise ClientError( + "Malformed Digest auth challenge: Missing 'nonce' parameter" + ) + + # Empty realm values are allowed per RFC 7616 (SHOULD, not MUST, contain host name) + realm = challenge["realm"] + nonce = challenge["nonce"] + + # Empty nonce values are not allowed as they are security-critical for replay protection + if not nonce: + raise ClientError( + "Security issue: Digest auth challenge contains empty 'nonce' value" + ) + + qop_raw = challenge.get("qop", "") + # Preserve original algorithm case for response while using uppercase for processing + algorithm_original = challenge.get("algorithm", "MD5") + algorithm = algorithm_original.upper() + opaque = challenge.get("opaque", "") + + # Convert string values to bytes once + nonce_bytes = nonce.encode("utf-8") + realm_bytes = realm.encode("utf-8") + path = URL(url).path_qs + + # Process QoP + qop = "" + qop_bytes = b"" + if qop_raw: + valid_qops = {"auth", "auth-int"}.intersection( + {q.strip() for q in qop_raw.split(",") if q.strip()} + ) + if not valid_qops: + raise ClientError( + f"Digest auth error: Unsupported Quality of Protection (qop) value(s): {qop_raw}" + ) + + qop = "auth-int" if "auth-int" in valid_qops else "auth" + qop_bytes = qop.encode("utf-8") + + if algorithm not in DigestFunctions: + raise ClientError( + f"Digest auth error: Unsupported hash algorithm: {algorithm}. " + f"Supported algorithms: {', '.join(SUPPORTED_ALGORITHMS)}" + ) + hash_fn: Final = DigestFunctions[algorithm] + + def H(x: bytes) -> bytes: + """RFC 7616 Section 3: Hash function H(data) = hex(hash(data)).""" + return hash_fn(x).hexdigest().encode() + + def KD(s: bytes, d: bytes) -> bytes: + """RFC 7616 Section 3: KD(secret, data) = H(concat(secret, ":", data)).""" + return H(b":".join((s, d))) + + # Calculate A1 and A2 + A1 = b":".join((self._login_bytes, realm_bytes, self._password_bytes)) + A2 = f"{method.upper()}:{path}".encode() + if qop == "auth-int": + if isinstance(body, Payload): # will always be empty bytes unless Payload + entity_bytes = await body.as_bytes() # Get bytes from Payload + else: + entity_bytes = body + entity_hash = H(entity_bytes) + A2 = b":".join((A2, entity_hash)) + + HA1 = H(A1) + HA2 = H(A2) + + # Nonce count handling + if nonce_bytes == self._last_nonce_bytes: + self._nonce_count += 1 + else: + self._nonce_count = 1 + + self._last_nonce_bytes = nonce_bytes + ncvalue = f"{self._nonce_count:08x}" + ncvalue_bytes = ncvalue.encode("utf-8") + + # Generate client nonce + cnonce = hashlib.sha1( + b"".join( + [ + str(self._nonce_count).encode("utf-8"), + nonce_bytes, + time.ctime().encode("utf-8"), + os.urandom(8), + ] + ) + ).hexdigest()[:16] + cnonce_bytes = cnonce.encode("utf-8") + + # Special handling for session-based algorithms + if algorithm.upper().endswith("-SESS"): + HA1 = H(b":".join((HA1, nonce_bytes, cnonce_bytes))) + + # Calculate the response digest + if qop: + noncebit = b":".join( + (nonce_bytes, ncvalue_bytes, cnonce_bytes, qop_bytes, HA2) + ) + response_digest = KD(HA1, noncebit) + else: + response_digest = KD(HA1, b":".join((nonce_bytes, HA2))) + + # Define a dict mapping of header fields to their values + # Group fields into always-present, optional, and qop-dependent + header_fields = { + # Always present fields + "username": escape_quotes(self._login_str), + "realm": escape_quotes(realm), + "nonce": escape_quotes(nonce), + "uri": path, + "response": response_digest.decode(), + "algorithm": algorithm_original, + } + + # Optional fields + if opaque: + header_fields["opaque"] = escape_quotes(opaque) + + # QoP-dependent fields + if qop: + header_fields["qop"] = qop + header_fields["nc"] = ncvalue + header_fields["cnonce"] = cnonce + + # Build header using templates for each field type + pairs: List[str] = [] + for field, value in header_fields.items(): + if field in QUOTED_AUTH_FIELDS: + pairs.append(f'{field}="{value}"') + else: + pairs.append(f"{field}={value}") + + return f"Digest {', '.join(pairs)}" + + def _in_protection_space(self, url: URL) -> bool: + """ + Check if the given URL is within the current protection space. + + According to RFC 7616, a URI is in the protection space if any URI + in the protection space is a prefix of it (after both have been made absolute). + """ + request_str = str(url) + for space_str in self._protection_space: + # Check if request starts with space URL + if not request_str.startswith(space_str): + continue + # Exact match or space ends with / (proper directory prefix) + if len(request_str) == len(space_str) or space_str[-1] == "/": + return True + # Check next char is / to ensure proper path boundary + if request_str[len(space_str)] == "/": + return True + return False + + def _authenticate(self, response: ClientResponse) -> bool: + """ + Takes the given response and tries digest-auth, if needed. + + Returns true if the original request must be resent. + """ + if response.status != 401: + return False + + auth_header = response.headers.get("www-authenticate", "") + if not auth_header: + return False # No authentication header present + + method, sep, headers = auth_header.partition(" ") + if not sep: + # No space found in www-authenticate header + return False # Malformed auth header, missing scheme separator + + if method.lower() != "digest": + # Not a digest auth challenge (could be Basic, Bearer, etc.) + return False + + if not headers: + # We have a digest scheme but no parameters + return False # Malformed digest header, missing parameters + + # We have a digest auth header with content + if not (header_pairs := parse_header_pairs(headers)): + # Failed to parse any key-value pairs + return False # Malformed digest header, no valid parameters + + # Extract challenge parameters + self._challenge = {} + for field in CHALLENGE_FIELDS: + if value := header_pairs.get(field): + self._challenge[field] = value + + # Update protection space based on domain parameter or default to origin + origin = response.url.origin() + + if domain := self._challenge.get("domain"): + # Parse space-separated list of URIs + self._protection_space = [] + for uri in domain.split(): + # Remove quotes if present + uri = uri.strip('"') + if uri.startswith("/"): + # Path-absolute, relative to origin + self._protection_space.append(str(origin.join(URL(uri)))) + else: + # Absolute URI + self._protection_space.append(str(URL(uri))) + else: + # No domain specified, protection space is entire origin + self._protection_space = [str(origin)] + + # Return True only if we found at least one challenge parameter + return bool(self._challenge) + + async def __call__( + self, request: ClientRequest, handler: ClientHandlerType + ) -> ClientResponse: + """Run the digest auth middleware.""" + response = None + for retry_count in range(2): + # Apply authorization header if: + # 1. This is a retry after 401 (retry_count > 0), OR + # 2. Preemptive auth is enabled AND we have a challenge AND the URL is in protection space + if retry_count > 0 or ( + self._preemptive + and self._challenge + and self._in_protection_space(request.url) + ): + request.headers[hdrs.AUTHORIZATION] = await self._encode( + request.method, request.url, request.body + ) + + # Send the request + response = await handler(request) + + # Check if we need to authenticate + if not self._authenticate(response): + break + + # At this point, response is guaranteed to be defined + assert response is not None + return response diff --git a/lib/python3.12/site-packages/aiohttp/client_middlewares.py b/lib/python3.12/site-packages/aiohttp/client_middlewares.py new file mode 100644 index 0000000000000000000000000000000000000000..3ca2cb202ad93963369f2a10fd1d118a194c4405 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_middlewares.py @@ -0,0 +1,55 @@ +"""Client middleware support.""" + +from collections.abc import Awaitable, Callable, Sequence + +from .client_reqrep import ClientRequest, ClientResponse + +__all__ = ("ClientMiddlewareType", "ClientHandlerType", "build_client_middlewares") + +# Type alias for client request handlers - functions that process requests and return responses +ClientHandlerType = Callable[[ClientRequest], Awaitable[ClientResponse]] + +# Type for client middleware - similar to server but uses ClientRequest/ClientResponse +ClientMiddlewareType = Callable[ + [ClientRequest, ClientHandlerType], Awaitable[ClientResponse] +] + + +def build_client_middlewares( + handler: ClientHandlerType, + middlewares: Sequence[ClientMiddlewareType], +) -> ClientHandlerType: + """ + Apply middlewares to request handler. + + The middlewares are applied in reverse order, so the first middleware + in the list wraps all subsequent middlewares and the handler. + + This implementation avoids using partial/update_wrapper to minimize overhead + and doesn't cache to avoid holding references to stateful middleware. + """ + # Optimize for single middleware case + if len(middlewares) == 1: + middleware = middlewares[0] + + async def single_middleware_handler(req: ClientRequest) -> ClientResponse: + return await middleware(req, handler) + + return single_middleware_handler + + # Build the chain for multiple middlewares + current_handler = handler + + for middleware in reversed(middlewares): + # Create a new closure that captures the current state + def make_wrapper( + mw: ClientMiddlewareType, next_h: ClientHandlerType + ) -> ClientHandlerType: + async def wrapped(req: ClientRequest) -> ClientResponse: + return await mw(req, next_h) + + return wrapped + + current_handler = make_wrapper(middleware, current_handler) + + return current_handler diff --git a/lib/python3.12/site-packages/aiohttp/client_proto.py b/lib/python3.12/site-packages/aiohttp/client_proto.py new file mode 100644 index 0000000000000000000000000000000000000000..e2fb1ce64cb6a39f5a72aecfd5840536defba519 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_proto.py @@ -0,0 +1,359 @@ +import asyncio +from contextlib import suppress +from typing import Any, Optional, Tuple, Union + +from .base_protocol import BaseProtocol +from .client_exceptions import ( + ClientConnectionError, + ClientOSError, + ClientPayloadError, + ServerDisconnectedError, + SocketTimeoutError, +) +from .helpers import ( + _EXC_SENTINEL, + EMPTY_BODY_STATUS_CODES, + BaseTimerContext, + set_exception, + set_result, +) +from .http import HttpResponseParser, RawResponseMessage +from .http_exceptions import HttpProcessingError +from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader + + +class ResponseHandler(BaseProtocol, DataQueue[Tuple[RawResponseMessage, StreamReader]]): + """Helper class to adapt between Protocol and StreamReader.""" + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + BaseProtocol.__init__(self, loop=loop) + DataQueue.__init__(self, loop) + + self._should_close = False + + self._payload: Optional[StreamReader] = None + self._skip_payload = False + self._payload_parser = None + + self._timer = None + + self._tail = b"" + self._upgraded = False + self._parser: Optional[HttpResponseParser] = None + + self._read_timeout: Optional[float] = None + self._read_timeout_handle: Optional[asyncio.TimerHandle] = None + + self._timeout_ceil_threshold: Optional[float] = 5 + + self._closed: Union[None, asyncio.Future[None]] = None + self._connection_lost_called = False + + @property + def closed(self) -> Union[None, asyncio.Future[None]]: + """Future that is set when the connection is closed. + + This property returns a Future that will be completed when the connection + is closed. The Future is created lazily on first access to avoid creating + futures that will never be awaited. + + Returns: + - A Future[None] if the connection is still open or was closed after + this property was accessed + - None if connection_lost() was already called before this property + was ever accessed (indicating no one is waiting for the closure) + """ + if self._closed is None and not self._connection_lost_called: + self._closed = self._loop.create_future() + return self._closed + + @property + def upgraded(self) -> bool: + return self._upgraded + + @property + def should_close(self) -> bool: + return bool( + self._should_close + or (self._payload is not None and not self._payload.is_eof()) + or self._upgraded + or self._exception is not None + or self._payload_parser is not None + or self._buffer + or self._tail + ) + + def force_close(self) -> None: + self._should_close = True + + def close(self) -> None: + self._exception = None # Break cyclic references + transport = self.transport + if transport is not None: + transport.close() + self.transport = None + self._payload = None + self._drop_timeout() + + def abort(self) -> None: + self._exception = None # Break cyclic references + transport = self.transport + if transport is not None: + transport.abort() + self.transport = None + self._payload = None + self._drop_timeout() + + def is_connected(self) -> bool: + return self.transport is not None and not self.transport.is_closing() + + def connection_lost(self, exc: Optional[BaseException]) -> None: + self._connection_lost_called = True + self._drop_timeout() + + original_connection_error = exc + reraised_exc = original_connection_error + + connection_closed_cleanly = original_connection_error is None + + if self._closed is not None: + # If someone is waiting for the closed future, + # we should set it to None or an exception. If + # self._closed is None, it means that + # connection_lost() was called already + # or nobody is waiting for it. + if connection_closed_cleanly: + set_result(self._closed, None) + else: + assert original_connection_error is not None + set_exception( + self._closed, + ClientConnectionError( + f"Connection lost: {original_connection_error !s}", + ), + original_connection_error, + ) + + if self._payload_parser is not None: + with suppress(Exception): # FIXME: log this somehow? + self._payload_parser.feed_eof() + + uncompleted = None + if self._parser is not None: + try: + uncompleted = self._parser.feed_eof() + except Exception as underlying_exc: + if self._payload is not None: + client_payload_exc_msg = ( + f"Response payload is not completed: {underlying_exc !r}" + ) + if not connection_closed_cleanly: + client_payload_exc_msg = ( + f"{client_payload_exc_msg !s}. " + f"{original_connection_error !r}" + ) + set_exception( + self._payload, + ClientPayloadError(client_payload_exc_msg), + underlying_exc, + ) + + if not self.is_eof(): + if isinstance(original_connection_error, OSError): + reraised_exc = ClientOSError(*original_connection_error.args) + if connection_closed_cleanly: + reraised_exc = ServerDisconnectedError(uncompleted) + # assigns self._should_close to True as side effect, + # we do it anyway below + underlying_non_eof_exc = ( + _EXC_SENTINEL + if connection_closed_cleanly + else original_connection_error + ) + assert underlying_non_eof_exc is not None + assert reraised_exc is not None + self.set_exception(reraised_exc, underlying_non_eof_exc) + + self._should_close = True + self._parser = None + self._payload = None + self._payload_parser = None + self._reading_paused = False + + super().connection_lost(reraised_exc) + + def eof_received(self) -> None: + # should call parser.feed_eof() most likely + self._drop_timeout() + + def pause_reading(self) -> None: + super().pause_reading() + self._drop_timeout() + + def resume_reading(self) -> None: + super().resume_reading() + self._reschedule_timeout() + + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + self._should_close = True + self._drop_timeout() + super().set_exception(exc, exc_cause) + + def set_parser(self, parser: Any, payload: Any) -> None: + # TODO: actual types are: + # parser: WebSocketReader + # payload: WebSocketDataQueue + # but they are not generi enough + # Need an ABC for both types + self._payload = payload + self._payload_parser = parser + + self._drop_timeout() + + if self._tail: + data, self._tail = self._tail, b"" + self.data_received(data) + + def set_response_params( + self, + *, + timer: Optional[BaseTimerContext] = None, + skip_payload: bool = False, + read_until_eof: bool = False, + auto_decompress: bool = True, + read_timeout: Optional[float] = None, + read_bufsize: int = 2**16, + timeout_ceil_threshold: float = 5, + max_line_size: int = 8190, + max_field_size: int = 8190, + ) -> None: + self._skip_payload = skip_payload + + self._read_timeout = read_timeout + + self._timeout_ceil_threshold = timeout_ceil_threshold + + self._parser = HttpResponseParser( + self, + self._loop, + read_bufsize, + timer=timer, + payload_exception=ClientPayloadError, + response_with_body=not skip_payload, + read_until_eof=read_until_eof, + auto_decompress=auto_decompress, + max_line_size=max_line_size, + max_field_size=max_field_size, + ) + + if self._tail: + data, self._tail = self._tail, b"" + self.data_received(data) + + def _drop_timeout(self) -> None: + if self._read_timeout_handle is not None: + self._read_timeout_handle.cancel() + self._read_timeout_handle = None + + def _reschedule_timeout(self) -> None: + timeout = self._read_timeout + if self._read_timeout_handle is not None: + self._read_timeout_handle.cancel() + + if timeout: + self._read_timeout_handle = self._loop.call_later( + timeout, self._on_read_timeout + ) + else: + self._read_timeout_handle = None + + def start_timeout(self) -> None: + self._reschedule_timeout() + + @property + def read_timeout(self) -> Optional[float]: + return self._read_timeout + + @read_timeout.setter + def read_timeout(self, read_timeout: Optional[float]) -> None: + self._read_timeout = read_timeout + + def _on_read_timeout(self) -> None: + exc = SocketTimeoutError("Timeout on reading data from socket") + self.set_exception(exc) + if self._payload is not None: + set_exception(self._payload, exc) + + def data_received(self, data: bytes) -> None: + self._reschedule_timeout() + + if not data: + return + + # custom payload parser - currently always WebSocketReader + if self._payload_parser is not None: + eof, tail = self._payload_parser.feed_data(data) + if eof: + self._payload = None + self._payload_parser = None + + if tail: + self.data_received(tail) + return + + if self._upgraded or self._parser is None: + # i.e. websocket connection, websocket parser is not set yet + self._tail += data + return + + # parse http messages + try: + messages, upgraded, tail = self._parser.feed_data(data) + except BaseException as underlying_exc: + if self.transport is not None: + # connection.release() could be called BEFORE + # data_received(), the transport is already + # closed in this case + self.transport.close() + # should_close is True after the call + if isinstance(underlying_exc, HttpProcessingError): + exc = HttpProcessingError( + code=underlying_exc.code, + message=underlying_exc.message, + headers=underlying_exc.headers, + ) + else: + exc = HttpProcessingError() + self.set_exception(exc, underlying_exc) + return + + self._upgraded = upgraded + + payload: Optional[StreamReader] = None + for message, payload in messages: + if message.should_close: + self._should_close = True + + self._payload = payload + + if self._skip_payload or message.code in EMPTY_BODY_STATUS_CODES: + self.feed_data((message, EMPTY_PAYLOAD), 0) + else: + self.feed_data((message, payload), 0) + + if payload is not None: + # new message(s) was processed + # register timeout handler unsubscribing + # either on end-of-stream or immediately for + # EMPTY_PAYLOAD + if payload is not EMPTY_PAYLOAD: + payload.on_eof(self._drop_timeout) + else: + self._drop_timeout() + + if upgraded and tail: + self.data_received(tail) diff --git a/lib/python3.12/site-packages/aiohttp/client_reqrep.py b/lib/python3.12/site-packages/aiohttp/client_reqrep.py new file mode 100644 index 0000000000000000000000000000000000000000..a9e0795893d6a56595f6f7b04b62cc4486d48a0f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_reqrep.py @@ -0,0 +1,1536 @@ +import asyncio +import codecs +import contextlib +import functools +import io +import re +import sys +import traceback +import warnings +from collections.abc import Mapping +from hashlib import md5, sha1, sha256 +from http.cookies import Morsel, SimpleCookie +from types import MappingProxyType, TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + List, + Literal, + NamedTuple, + Optional, + Tuple, + Type, + Union, +) + +import attr +from multidict import CIMultiDict, CIMultiDictProxy, MultiDict, MultiDictProxy +from yarl import URL + +from . import hdrs, helpers, http, multipart, payload +from ._cookie_helpers import ( + parse_cookie_header, + parse_set_cookie_headers, + preserve_morsel_with_coded_value, +) +from .abc import AbstractStreamWriter +from .client_exceptions import ( + ClientConnectionError, + ClientOSError, + ClientResponseError, + ContentTypeError, + InvalidURL, + ServerFingerprintMismatch, +) +from .compression_utils import HAS_BROTLI, HAS_ZSTD +from .formdata import FormData +from .helpers import ( + _SENTINEL, + BaseTimerContext, + BasicAuth, + HeadersMixin, + TimerNoop, + noop, + reify, + sentinel, + set_exception, + set_result, +) +from .http import ( + SERVER_SOFTWARE, + HttpVersion, + HttpVersion10, + HttpVersion11, + StreamWriter, +) +from .streams import StreamReader +from .typedefs import ( + DEFAULT_JSON_DECODER, + JSONDecoder, + LooseCookies, + LooseHeaders, + Query, + RawHeaders, +) + +if TYPE_CHECKING: + import ssl + from ssl import SSLContext +else: + try: + import ssl + from ssl import SSLContext + except ImportError: # pragma: no cover + ssl = None # type: ignore[assignment] + SSLContext = object # type: ignore[misc,assignment] + + +__all__ = ("ClientRequest", "ClientResponse", "RequestInfo", "Fingerprint") + + +if TYPE_CHECKING: + from .client import ClientSession + from .connector import Connection + from .tracing import Trace + + +_CONNECTION_CLOSED_EXCEPTION = ClientConnectionError("Connection closed") +_CONTAINS_CONTROL_CHAR_RE = re.compile(r"[^-!#$%&'*+.^_`|~0-9a-zA-Z]") +json_re = re.compile(r"^application/(?:[\w.+-]+?\+)?json") + + +def _gen_default_accept_encoding() -> str: + encodings = [ + "gzip", + "deflate", + ] + if HAS_BROTLI: + encodings.append("br") + if HAS_ZSTD: + encodings.append("zstd") + return ", ".join(encodings) + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class ContentDisposition: + type: Optional[str] + parameters: "MappingProxyType[str, str]" + filename: Optional[str] + + +class _RequestInfo(NamedTuple): + url: URL + method: str + headers: "CIMultiDictProxy[str]" + real_url: URL + + +class RequestInfo(_RequestInfo): + + def __new__( + cls, + url: URL, + method: str, + headers: "CIMultiDictProxy[str]", + real_url: Union[URL, _SENTINEL] = sentinel, + ) -> "RequestInfo": + """Create a new RequestInfo instance. + + For backwards compatibility, the real_url parameter is optional. + """ + return tuple.__new__( + cls, (url, method, headers, url if real_url is sentinel else real_url) + ) + + +class Fingerprint: + HASHFUNC_BY_DIGESTLEN = { + 16: md5, + 20: sha1, + 32: sha256, + } + + def __init__(self, fingerprint: bytes) -> None: + digestlen = len(fingerprint) + hashfunc = self.HASHFUNC_BY_DIGESTLEN.get(digestlen) + if not hashfunc: + raise ValueError("fingerprint has invalid length") + elif hashfunc is md5 or hashfunc is sha1: + raise ValueError("md5 and sha1 are insecure and not supported. Use sha256.") + self._hashfunc = hashfunc + self._fingerprint = fingerprint + + @property + def fingerprint(self) -> bytes: + return self._fingerprint + + def check(self, transport: asyncio.Transport) -> None: + if not transport.get_extra_info("sslcontext"): + return + sslobj = transport.get_extra_info("ssl_object") + cert = sslobj.getpeercert(binary_form=True) + got = self._hashfunc(cert).digest() + if got != self._fingerprint: + host, port, *_ = transport.get_extra_info("peername") + raise ServerFingerprintMismatch(self._fingerprint, got, host, port) + + +if ssl is not None: + SSL_ALLOWED_TYPES = (ssl.SSLContext, bool, Fingerprint, type(None)) +else: # pragma: no cover + SSL_ALLOWED_TYPES = (bool, type(None)) + + +def _merge_ssl_params( + ssl: Union["SSLContext", bool, Fingerprint], + verify_ssl: Optional[bool], + ssl_context: Optional["SSLContext"], + fingerprint: Optional[bytes], +) -> Union["SSLContext", bool, Fingerprint]: + if ssl is None: + ssl = True # Double check for backwards compatibility + if verify_ssl is not None and not verify_ssl: + warnings.warn( + "verify_ssl is deprecated, use ssl=False instead", + DeprecationWarning, + stacklevel=3, + ) + if ssl is not True: + raise ValueError( + "verify_ssl, ssl_context, fingerprint and ssl " + "parameters are mutually exclusive" + ) + else: + ssl = False + if ssl_context is not None: + warnings.warn( + "ssl_context is deprecated, use ssl=context instead", + DeprecationWarning, + stacklevel=3, + ) + if ssl is not True: + raise ValueError( + "verify_ssl, ssl_context, fingerprint and ssl " + "parameters are mutually exclusive" + ) + else: + ssl = ssl_context + if fingerprint is not None: + warnings.warn( + "fingerprint is deprecated, use ssl=Fingerprint(fingerprint) instead", + DeprecationWarning, + stacklevel=3, + ) + if ssl is not True: + raise ValueError( + "verify_ssl, ssl_context, fingerprint and ssl " + "parameters are mutually exclusive" + ) + else: + ssl = Fingerprint(fingerprint) + if not isinstance(ssl, SSL_ALLOWED_TYPES): + raise TypeError( + "ssl should be SSLContext, bool, Fingerprint or None, " + "got {!r} instead.".format(ssl) + ) + return ssl + + +_SSL_SCHEMES = frozenset(("https", "wss")) + + +# ConnectionKey is a NamedTuple because it is used as a key in a dict +# and a set in the connector. Since a NamedTuple is a tuple it uses +# the fast native tuple __hash__ and __eq__ implementation in CPython. +class ConnectionKey(NamedTuple): + # the key should contain an information about used proxy / TLS + # to prevent reusing wrong connections from a pool + host: str + port: Optional[int] + is_ssl: bool + ssl: Union[SSLContext, bool, Fingerprint] + proxy: Optional[URL] + proxy_auth: Optional[BasicAuth] + proxy_headers_hash: Optional[int] # hash(CIMultiDict) + + +def _is_expected_content_type( + response_content_type: str, expected_content_type: str +) -> bool: + if expected_content_type == "application/json": + return json_re.match(response_content_type) is not None + return expected_content_type in response_content_type + + +def _warn_if_unclosed_payload(payload: payload.Payload, stacklevel: int = 2) -> None: + """Warn if the payload is not closed. + + Callers must check that the body is a Payload before calling this method. + + Args: + payload: The payload to check + stacklevel: Stack level for the warning (default 2 for direct callers) + """ + if not payload.autoclose and not payload.consumed: + warnings.warn( + "The previous request body contains unclosed resources. " + "Use await request.update_body() instead of setting request.body " + "directly to properly close resources and avoid leaks.", + ResourceWarning, + stacklevel=stacklevel, + ) + + +class ClientResponse(HeadersMixin): + + # Some of these attributes are None when created, + # but will be set by the start() method. + # As the end user will likely never see the None values, we cheat the types below. + # from the Status-Line of the response + version: Optional[HttpVersion] = None # HTTP-Version + status: int = None # type: ignore[assignment] # Status-Code + reason: Optional[str] = None # Reason-Phrase + + content: StreamReader = None # type: ignore[assignment] # Payload stream + _body: Optional[bytes] = None + _headers: CIMultiDictProxy[str] = None # type: ignore[assignment] + _history: Tuple["ClientResponse", ...] = () + _raw_headers: RawHeaders = None # type: ignore[assignment] + + _connection: Optional["Connection"] = None # current connection + _cookies: Optional[SimpleCookie] = None + _raw_cookie_headers: Optional[Tuple[str, ...]] = None + _continue: Optional["asyncio.Future[bool]"] = None + _source_traceback: Optional[traceback.StackSummary] = None + _session: Optional["ClientSession"] = None + # set up by ClientRequest after ClientResponse object creation + # post-init stage allows to not change ctor signature + _closed = True # to allow __del__ for non-initialized properly response + _released = False + _in_context = False + + _resolve_charset: Callable[["ClientResponse", bytes], str] = lambda *_: "utf-8" + + __writer: Optional["asyncio.Task[None]"] = None + + def __init__( + self, + method: str, + url: URL, + *, + writer: "Optional[asyncio.Task[None]]", + continue100: Optional["asyncio.Future[bool]"], + timer: BaseTimerContext, + request_info: RequestInfo, + traces: List["Trace"], + loop: asyncio.AbstractEventLoop, + session: "ClientSession", + ) -> None: + # URL forbids subclasses, so a simple type check is enough. + assert type(url) is URL + + self.method = method + + self._real_url = url + self._url = url.with_fragment(None) if url.raw_fragment else url + if writer is not None: + self._writer = writer + if continue100 is not None: + self._continue = continue100 + self._request_info = request_info + self._timer = timer if timer is not None else TimerNoop() + self._cache: Dict[str, Any] = {} + self._traces = traces + self._loop = loop + # Save reference to _resolve_charset, so that get_encoding() will still + # work after the response has finished reading the body. + # TODO: Fix session=None in tests (see ClientRequest.__init__). + if session is not None: + # store a reference to session #1985 + self._session = session + self._resolve_charset = session._resolve_charset + if loop.get_debug(): + self._source_traceback = traceback.extract_stack(sys._getframe(1)) + + def __reset_writer(self, _: object = None) -> None: + self.__writer = None + + @property + def _writer(self) -> Optional["asyncio.Task[None]"]: + """The writer task for streaming data. + + _writer is only provided for backwards compatibility + for subclasses that may need to access it. + """ + return self.__writer + + @_writer.setter + def _writer(self, writer: Optional["asyncio.Task[None]"]) -> None: + """Set the writer task for streaming data.""" + if self.__writer is not None: + self.__writer.remove_done_callback(self.__reset_writer) + self.__writer = writer + if writer is None: + return + if writer.done(): + # The writer is already done, so we can clear it immediately. + self.__writer = None + else: + writer.add_done_callback(self.__reset_writer) + + @property + def cookies(self) -> SimpleCookie: + if self._cookies is None: + if self._raw_cookie_headers is not None: + # Parse cookies for response.cookies (SimpleCookie for backward compatibility) + cookies = SimpleCookie() + # Use parse_set_cookie_headers for more lenient parsing that handles + # malformed cookies better than SimpleCookie.load + cookies.update(parse_set_cookie_headers(self._raw_cookie_headers)) + self._cookies = cookies + else: + self._cookies = SimpleCookie() + return self._cookies + + @cookies.setter + def cookies(self, cookies: SimpleCookie) -> None: + self._cookies = cookies + # Generate raw cookie headers from the SimpleCookie + if cookies: + self._raw_cookie_headers = tuple( + morsel.OutputString() for morsel in cookies.values() + ) + else: + self._raw_cookie_headers = None + + @reify + def url(self) -> URL: + return self._url + + @reify + def url_obj(self) -> URL: + warnings.warn("Deprecated, use .url #1654", DeprecationWarning, stacklevel=2) + return self._url + + @reify + def real_url(self) -> URL: + return self._real_url + + @reify + def host(self) -> str: + assert self._url.host is not None + return self._url.host + + @reify + def headers(self) -> "CIMultiDictProxy[str]": + return self._headers + + @reify + def raw_headers(self) -> RawHeaders: + return self._raw_headers + + @reify + def request_info(self) -> RequestInfo: + return self._request_info + + @reify + def content_disposition(self) -> Optional[ContentDisposition]: + raw = self._headers.get(hdrs.CONTENT_DISPOSITION) + if raw is None: + return None + disposition_type, params_dct = multipart.parse_content_disposition(raw) + params = MappingProxyType(params_dct) + filename = multipart.content_disposition_filename(params) + return ContentDisposition(disposition_type, params, filename) + + def __del__(self, _warnings: Any = warnings) -> None: + if self._closed: + return + + if self._connection is not None: + self._connection.release() + self._cleanup_writer() + + if self._loop.get_debug(): + kwargs = {"source": self} + _warnings.warn(f"Unclosed response {self!r}", ResourceWarning, **kwargs) + context = {"client_response": self, "message": "Unclosed response"} + if self._source_traceback: + context["source_traceback"] = self._source_traceback + self._loop.call_exception_handler(context) + + def __repr__(self) -> str: + out = io.StringIO() + ascii_encodable_url = str(self.url) + if self.reason: + ascii_encodable_reason = self.reason.encode( + "ascii", "backslashreplace" + ).decode("ascii") + else: + ascii_encodable_reason = "None" + print( + "".format( + ascii_encodable_url, self.status, ascii_encodable_reason + ), + file=out, + ) + print(self.headers, file=out) + return out.getvalue() + + @property + def connection(self) -> Optional["Connection"]: + return self._connection + + @reify + def history(self) -> Tuple["ClientResponse", ...]: + """A sequence of of responses, if redirects occurred.""" + return self._history + + @reify + def links(self) -> "MultiDictProxy[MultiDictProxy[Union[str, URL]]]": + links_str = ", ".join(self.headers.getall("link", [])) + + if not links_str: + return MultiDictProxy(MultiDict()) + + links: MultiDict[MultiDictProxy[Union[str, URL]]] = MultiDict() + + for val in re.split(r",(?=\s*<)", links_str): + match = re.match(r"\s*<(.*)>(.*)", val) + if match is None: # pragma: no cover + # the check exists to suppress mypy error + continue + url, params_str = match.groups() + params = params_str.split(";")[1:] + + link: MultiDict[Union[str, URL]] = MultiDict() + + for param in params: + match = re.match(r"^\s*(\S*)\s*=\s*(['\"]?)(.*?)(\2)\s*$", param, re.M) + if match is None: # pragma: no cover + # the check exists to suppress mypy error + continue + key, _, value, _ = match.groups() + + link.add(key, value) + + key = link.get("rel", url) + + link.add("url", self.url.join(URL(url))) + + links.add(str(key), MultiDictProxy(link)) + + return MultiDictProxy(links) + + async def start(self, connection: "Connection") -> "ClientResponse": + """Start response processing.""" + self._closed = False + self._protocol = connection.protocol + self._connection = connection + + with self._timer: + while True: + # read response + try: + protocol = self._protocol + message, payload = await protocol.read() # type: ignore[union-attr] + except http.HttpProcessingError as exc: + raise ClientResponseError( + self.request_info, + self.history, + status=exc.code, + message=exc.message, + headers=exc.headers, + ) from exc + + if message.code < 100 or message.code > 199 or message.code == 101: + break + + if self._continue is not None: + set_result(self._continue, True) + self._continue = None + + # payload eof handler + payload.on_eof(self._response_eof) + + # response status + self.version = message.version + self.status = message.code + self.reason = message.reason + + # headers + self._headers = message.headers # type is CIMultiDictProxy + self._raw_headers = message.raw_headers # type is Tuple[bytes, bytes] + + # payload + self.content = payload + + # cookies + if cookie_hdrs := self.headers.getall(hdrs.SET_COOKIE, ()): + # Store raw cookie headers for CookieJar + self._raw_cookie_headers = tuple(cookie_hdrs) + return self + + def _response_eof(self) -> None: + if self._closed: + return + + # protocol could be None because connection could be detached + protocol = self._connection and self._connection.protocol + if protocol is not None and protocol.upgraded: + return + + self._closed = True + self._cleanup_writer() + self._release_connection() + + @property + def closed(self) -> bool: + return self._closed + + def close(self) -> None: + if not self._released: + self._notify_content() + + self._closed = True + if self._loop is None or self._loop.is_closed(): + return + + self._cleanup_writer() + if self._connection is not None: + self._connection.close() + self._connection = None + + def release(self) -> Any: + if not self._released: + self._notify_content() + + self._closed = True + + self._cleanup_writer() + self._release_connection() + return noop() + + @property + def ok(self) -> bool: + """Returns ``True`` if ``status`` is less than ``400``, ``False`` if not. + + This is **not** a check for ``200 OK`` but a check that the response + status is under 400. + """ + return 400 > self.status + + def raise_for_status(self) -> None: + if not self.ok: + # reason should always be not None for a started response + assert self.reason is not None + + # If we're in a context we can rely on __aexit__() to release as the + # exception propagates. + if not self._in_context: + self.release() + + raise ClientResponseError( + self.request_info, + self.history, + status=self.status, + message=self.reason, + headers=self.headers, + ) + + def _release_connection(self) -> None: + if self._connection is not None: + if self.__writer is None: + self._connection.release() + self._connection = None + else: + self.__writer.add_done_callback(lambda f: self._release_connection()) + + async def _wait_released(self) -> None: + if self.__writer is not None: + try: + await self.__writer + except asyncio.CancelledError: + if ( + sys.version_info >= (3, 11) + and (task := asyncio.current_task()) + and task.cancelling() + ): + raise + self._release_connection() + + def _cleanup_writer(self) -> None: + if self.__writer is not None: + self.__writer.cancel() + self._session = None + + def _notify_content(self) -> None: + content = self.content + if content and content.exception() is None: + set_exception(content, _CONNECTION_CLOSED_EXCEPTION) + self._released = True + + async def wait_for_close(self) -> None: + if self.__writer is not None: + try: + await self.__writer + except asyncio.CancelledError: + if ( + sys.version_info >= (3, 11) + and (task := asyncio.current_task()) + and task.cancelling() + ): + raise + self.release() + + async def read(self) -> bytes: + """Read response payload.""" + if self._body is None: + try: + self._body = await self.content.read() + for trace in self._traces: + await trace.send_response_chunk_received( + self.method, self.url, self._body + ) + except BaseException: + self.close() + raise + elif self._released: # Response explicitly released + raise ClientConnectionError("Connection closed") + + protocol = self._connection and self._connection.protocol + if protocol is None or not protocol.upgraded: + await self._wait_released() # Underlying connection released + return self._body + + def get_encoding(self) -> str: + ctype = self.headers.get(hdrs.CONTENT_TYPE, "").lower() + mimetype = helpers.parse_mimetype(ctype) + + encoding = mimetype.parameters.get("charset") + if encoding: + with contextlib.suppress(LookupError, ValueError): + return codecs.lookup(encoding).name + + if mimetype.type == "application" and ( + mimetype.subtype == "json" or mimetype.subtype == "rdap" + ): + # RFC 7159 states that the default encoding is UTF-8. + # RFC 7483 defines application/rdap+json + return "utf-8" + + if self._body is None: + raise RuntimeError( + "Cannot compute fallback encoding of a not yet read body" + ) + + return self._resolve_charset(self, self._body) + + async def text(self, encoding: Optional[str] = None, errors: str = "strict") -> str: + """Read response payload and decode.""" + if self._body is None: + await self.read() + + if encoding is None: + encoding = self.get_encoding() + + return self._body.decode(encoding, errors=errors) # type: ignore[union-attr] + + async def json( + self, + *, + encoding: Optional[str] = None, + loads: JSONDecoder = DEFAULT_JSON_DECODER, + content_type: Optional[str] = "application/json", + ) -> Any: + """Read and decodes JSON response.""" + if self._body is None: + await self.read() + + if content_type: + ctype = self.headers.get(hdrs.CONTENT_TYPE, "").lower() + if not _is_expected_content_type(ctype, content_type): + raise ContentTypeError( + self.request_info, + self.history, + status=self.status, + message=( + "Attempt to decode JSON with unexpected mimetype: %s" % ctype + ), + headers=self.headers, + ) + + stripped = self._body.strip() # type: ignore[union-attr] + if not stripped: + return None + + if encoding is None: + encoding = self.get_encoding() + + return loads(stripped.decode(encoding)) + + async def __aenter__(self) -> "ClientResponse": + self._in_context = True + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + self._in_context = False + # similar to _RequestContextManager, we do not need to check + # for exceptions, response object can close connection + # if state is broken + self.release() + await self.wait_for_close() + + +class ClientRequest: + GET_METHODS = { + hdrs.METH_GET, + hdrs.METH_HEAD, + hdrs.METH_OPTIONS, + hdrs.METH_TRACE, + } + POST_METHODS = {hdrs.METH_PATCH, hdrs.METH_POST, hdrs.METH_PUT} + ALL_METHODS = GET_METHODS.union(POST_METHODS).union({hdrs.METH_DELETE}) + + DEFAULT_HEADERS = { + hdrs.ACCEPT: "*/*", + hdrs.ACCEPT_ENCODING: _gen_default_accept_encoding(), + } + + # Type of body depends on PAYLOAD_REGISTRY, which is dynamic. + _body: Union[None, payload.Payload] = None + auth = None + response = None + + __writer: Optional["asyncio.Task[None]"] = None # async task for streaming data + + # These class defaults help create_autospec() work correctly. + # If autospec is improved in future, maybe these can be removed. + url = URL() + method = "GET" + + _continue = None # waiter future for '100 Continue' response + + _skip_auto_headers: Optional["CIMultiDict[None]"] = None + + # N.B. + # Adding __del__ method with self._writer closing doesn't make sense + # because _writer is instance method, thus it keeps a reference to self. + # Until writer has finished finalizer will not be called. + + def __init__( + self, + method: str, + url: URL, + *, + params: Query = None, + headers: Optional[LooseHeaders] = None, + skip_auto_headers: Optional[Iterable[str]] = None, + data: Any = None, + cookies: Optional[LooseCookies] = None, + auth: Optional[BasicAuth] = None, + version: http.HttpVersion = http.HttpVersion11, + compress: Union[str, bool, None] = None, + chunked: Optional[bool] = None, + expect100: bool = False, + loop: Optional[asyncio.AbstractEventLoop] = None, + response_class: Optional[Type["ClientResponse"]] = None, + proxy: Optional[URL] = None, + proxy_auth: Optional[BasicAuth] = None, + timer: Optional[BaseTimerContext] = None, + session: Optional["ClientSession"] = None, + ssl: Union[SSLContext, bool, Fingerprint] = True, + proxy_headers: Optional[LooseHeaders] = None, + traces: Optional[List["Trace"]] = None, + trust_env: bool = False, + server_hostname: Optional[str] = None, + ): + if loop is None: + loop = asyncio.get_event_loop() + if match := _CONTAINS_CONTROL_CHAR_RE.search(method): + raise ValueError( + f"Method cannot contain non-token characters {method!r} " + f"(found at least {match.group()!r})" + ) + # URL forbids subclasses, so a simple type check is enough. + assert type(url) is URL, url + if proxy is not None: + assert type(proxy) is URL, proxy + # FIXME: session is None in tests only, need to fix tests + # assert session is not None + if TYPE_CHECKING: + assert session is not None + self._session = session + if params: + url = url.extend_query(params) + self.original_url = url + self.url = url.with_fragment(None) if url.raw_fragment else url + self.method = method.upper() + self.chunked = chunked + self.compress = compress + self.loop = loop + self.length = None + if response_class is None: + real_response_class = ClientResponse + else: + real_response_class = response_class + self.response_class: Type[ClientResponse] = real_response_class + self._timer = timer if timer is not None else TimerNoop() + self._ssl = ssl if ssl is not None else True + self.server_hostname = server_hostname + + if loop.get_debug(): + self._source_traceback = traceback.extract_stack(sys._getframe(1)) + + self.update_version(version) + self.update_host(url) + self.update_headers(headers) + self.update_auto_headers(skip_auto_headers) + self.update_cookies(cookies) + self.update_content_encoding(data) + self.update_auth(auth, trust_env) + self.update_proxy(proxy, proxy_auth, proxy_headers) + + self.update_body_from_data(data) + if data is not None or self.method not in self.GET_METHODS: + self.update_transfer_encoding() + self.update_expect_continue(expect100) + self._traces = [] if traces is None else traces + + def __reset_writer(self, _: object = None) -> None: + self.__writer = None + + def _get_content_length(self) -> Optional[int]: + """Extract and validate Content-Length header value. + + Returns parsed Content-Length value or None if not set. + Raises ValueError if header exists but cannot be parsed as an integer. + """ + if hdrs.CONTENT_LENGTH not in self.headers: + return None + + content_length_hdr = self.headers[hdrs.CONTENT_LENGTH] + try: + return int(content_length_hdr) + except ValueError: + raise ValueError( + f"Invalid Content-Length header: {content_length_hdr}" + ) from None + + @property + def skip_auto_headers(self) -> CIMultiDict[None]: + return self._skip_auto_headers or CIMultiDict() + + @property + def _writer(self) -> Optional["asyncio.Task[None]"]: + return self.__writer + + @_writer.setter + def _writer(self, writer: "asyncio.Task[None]") -> None: + if self.__writer is not None: + self.__writer.remove_done_callback(self.__reset_writer) + self.__writer = writer + writer.add_done_callback(self.__reset_writer) + + def is_ssl(self) -> bool: + return self.url.scheme in _SSL_SCHEMES + + @property + def ssl(self) -> Union["SSLContext", bool, Fingerprint]: + return self._ssl + + @property + def connection_key(self) -> ConnectionKey: + if proxy_headers := self.proxy_headers: + h: Optional[int] = hash(tuple(proxy_headers.items())) + else: + h = None + url = self.url + return tuple.__new__( + ConnectionKey, + ( + url.raw_host or "", + url.port, + url.scheme in _SSL_SCHEMES, + self._ssl, + self.proxy, + self.proxy_auth, + h, + ), + ) + + @property + def host(self) -> str: + ret = self.url.raw_host + assert ret is not None + return ret + + @property + def port(self) -> Optional[int]: + return self.url.port + + @property + def body(self) -> Union[payload.Payload, Literal[b""]]: + """Request body.""" + # empty body is represented as bytes for backwards compatibility + return self._body or b"" + + @body.setter + def body(self, value: Any) -> None: + """Set request body with warning for non-autoclose payloads. + + WARNING: This setter must be called from within an event loop and is not + thread-safe. Setting body outside of an event loop may raise RuntimeError + when closing file-based payloads. + + DEPRECATED: Direct assignment to body is deprecated and will be removed + in a future version. Use await update_body() instead for proper resource + management. + """ + # Close existing payload if present + if self._body is not None: + # Warn if the payload needs manual closing + # stacklevel=3: user code -> body setter -> _warn_if_unclosed_payload + _warn_if_unclosed_payload(self._body, stacklevel=3) + # NOTE: In the future, when we remove sync close support, + # this setter will need to be removed and only the async + # update_body() method will be available. For now, we call + # _close() for backwards compatibility. + self._body._close() + self._update_body(value) + + @property + def request_info(self) -> RequestInfo: + headers: CIMultiDictProxy[str] = CIMultiDictProxy(self.headers) + # These are created on every request, so we use a NamedTuple + # for performance reasons. We don't use the RequestInfo.__new__ + # method because it has a different signature which is provided + # for backwards compatibility only. + return tuple.__new__( + RequestInfo, (self.url, self.method, headers, self.original_url) + ) + + @property + def session(self) -> "ClientSession": + """Return the ClientSession instance. + + This property provides access to the ClientSession that initiated + this request, allowing middleware to make additional requests + using the same session. + """ + return self._session + + def update_host(self, url: URL) -> None: + """Update destination host, port and connection type (ssl).""" + # get host/port + if not url.raw_host: + raise InvalidURL(url) + + # basic auth info + if url.raw_user or url.raw_password: + self.auth = helpers.BasicAuth(url.user or "", url.password or "") + + def update_version(self, version: Union[http.HttpVersion, str]) -> None: + """Convert request version to two elements tuple. + + parser HTTP version '1.1' => (1, 1) + """ + if isinstance(version, str): + v = [part.strip() for part in version.split(".", 1)] + try: + version = http.HttpVersion(int(v[0]), int(v[1])) + except ValueError: + raise ValueError( + f"Can not parse http version number: {version}" + ) from None + self.version = version + + def update_headers(self, headers: Optional[LooseHeaders]) -> None: + """Update request headers.""" + self.headers: CIMultiDict[str] = CIMultiDict() + + # Build the host header + host = self.url.host_port_subcomponent + + # host_port_subcomponent is None when the URL is a relative URL. + # but we know we do not have a relative URL here. + assert host is not None + self.headers[hdrs.HOST] = host + + if not headers: + return + + if isinstance(headers, (dict, MultiDictProxy, MultiDict)): + headers = headers.items() + + for key, value in headers: # type: ignore[misc] + # A special case for Host header + if key in hdrs.HOST_ALL: + self.headers[key] = value + else: + self.headers.add(key, value) + + def update_auto_headers(self, skip_auto_headers: Optional[Iterable[str]]) -> None: + if skip_auto_headers is not None: + self._skip_auto_headers = CIMultiDict( + (hdr, None) for hdr in sorted(skip_auto_headers) + ) + used_headers = self.headers.copy() + used_headers.extend(self._skip_auto_headers) # type: ignore[arg-type] + else: + # Fast path when there are no headers to skip + # which is the most common case. + used_headers = self.headers + + for hdr, val in self.DEFAULT_HEADERS.items(): + if hdr not in used_headers: + self.headers[hdr] = val + + if hdrs.USER_AGENT not in used_headers: + self.headers[hdrs.USER_AGENT] = SERVER_SOFTWARE + + def update_cookies(self, cookies: Optional[LooseCookies]) -> None: + """Update request cookies header.""" + if not cookies: + return + + c = SimpleCookie() + if hdrs.COOKIE in self.headers: + # parse_cookie_header for RFC 6265 compliant Cookie header parsing + c.update(parse_cookie_header(self.headers.get(hdrs.COOKIE, ""))) + del self.headers[hdrs.COOKIE] + + if isinstance(cookies, Mapping): + iter_cookies = cookies.items() + else: + iter_cookies = cookies # type: ignore[assignment] + for name, value in iter_cookies: + if isinstance(value, Morsel): + # Use helper to preserve coded_value exactly as sent by server + c[name] = preserve_morsel_with_coded_value(value) + else: + c[name] = value # type: ignore[assignment] + + self.headers[hdrs.COOKIE] = c.output(header="", sep=";").strip() + + def update_content_encoding(self, data: Any) -> None: + """Set request content encoding.""" + if not data: + # Don't compress an empty body. + self.compress = None + return + + if self.headers.get(hdrs.CONTENT_ENCODING): + if self.compress: + raise ValueError( + "compress can not be set if Content-Encoding header is set" + ) + elif self.compress: + if not isinstance(self.compress, str): + self.compress = "deflate" + self.headers[hdrs.CONTENT_ENCODING] = self.compress + self.chunked = True # enable chunked, no need to deal with length + + def update_transfer_encoding(self) -> None: + """Analyze transfer-encoding header.""" + te = self.headers.get(hdrs.TRANSFER_ENCODING, "").lower() + + if "chunked" in te: + if self.chunked: + raise ValueError( + "chunked can not be set " + 'if "Transfer-Encoding: chunked" header is set' + ) + + elif self.chunked: + if hdrs.CONTENT_LENGTH in self.headers: + raise ValueError( + "chunked can not be set if Content-Length header is set" + ) + + self.headers[hdrs.TRANSFER_ENCODING] = "chunked" + + def update_auth(self, auth: Optional[BasicAuth], trust_env: bool = False) -> None: + """Set basic auth.""" + if auth is None: + auth = self.auth + if auth is None: + return + + if not isinstance(auth, helpers.BasicAuth): + raise TypeError("BasicAuth() tuple is required instead") + + self.headers[hdrs.AUTHORIZATION] = auth.encode() + + def update_body_from_data(self, body: Any, _stacklevel: int = 3) -> None: + """Update request body from data.""" + if self._body is not None: + _warn_if_unclosed_payload(self._body, stacklevel=_stacklevel) + + if body is None: + self._body = None + # Set Content-Length to 0 when body is None for methods that expect a body + if ( + self.method not in self.GET_METHODS + and not self.chunked + and hdrs.CONTENT_LENGTH not in self.headers + ): + self.headers[hdrs.CONTENT_LENGTH] = "0" + return + + # FormData + maybe_payload = body() if isinstance(body, FormData) else body + + try: + body_payload = payload.PAYLOAD_REGISTRY.get(maybe_payload, disposition=None) + except payload.LookupError: + body_payload = FormData(maybe_payload)() # type: ignore[arg-type] + + self._body = body_payload + # enable chunked encoding if needed + if not self.chunked and hdrs.CONTENT_LENGTH not in self.headers: + if (size := body_payload.size) is not None: + self.headers[hdrs.CONTENT_LENGTH] = str(size) + else: + self.chunked = True + + # copy payload headers + assert body_payload.headers + headers = self.headers + skip_headers = self._skip_auto_headers + for key, value in body_payload.headers.items(): + if key in headers or (skip_headers is not None and key in skip_headers): + continue + headers[key] = value + + def _update_body(self, body: Any) -> None: + """Update request body after its already been set.""" + # Remove existing Content-Length header since body is changing + if hdrs.CONTENT_LENGTH in self.headers: + del self.headers[hdrs.CONTENT_LENGTH] + + # Remove existing Transfer-Encoding header to avoid conflicts + if self.chunked and hdrs.TRANSFER_ENCODING in self.headers: + del self.headers[hdrs.TRANSFER_ENCODING] + + # Now update the body using the existing method + # Called from _update_body, add 1 to stacklevel from caller + self.update_body_from_data(body, _stacklevel=4) + + # Update transfer encoding headers if needed (same logic as __init__) + if body is not None or self.method not in self.GET_METHODS: + self.update_transfer_encoding() + + async def update_body(self, body: Any) -> None: + """ + Update request body and close previous payload if needed. + + This method safely updates the request body by first closing any existing + payload to prevent resource leaks, then setting the new body. + + IMPORTANT: Always use this method instead of setting request.body directly. + Direct assignment to request.body will leak resources if the previous body + contains file handles, streams, or other resources that need cleanup. + + Args: + body: The new body content. Can be: + - bytes/bytearray: Raw binary data + - str: Text data (will be encoded using charset from Content-Type) + - FormData: Form data that will be encoded as multipart/form-data + - Payload: A pre-configured payload object + - AsyncIterable: An async iterable of bytes chunks + - File-like object: Will be read and sent as binary data + - None: Clears the body + + Usage: + # CORRECT: Use update_body + await request.update_body(b"new request data") + + # WRONG: Don't set body directly + # request.body = b"new request data" # This will leak resources! + + # Update with form data + form_data = FormData() + form_data.add_field('field', 'value') + await request.update_body(form_data) + + # Clear body + await request.update_body(None) + + Note: + This method is async because it may need to close file handles or + other resources associated with the previous payload. Always await + this method to ensure proper cleanup. + + Warning: + Setting request.body directly is highly discouraged and can lead to: + - Resource leaks (unclosed file handles, streams) + - Memory leaks (unreleased buffers) + - Unexpected behavior with streaming payloads + + It is not recommended to change the payload type in middleware. If the + body was already set (e.g., as bytes), it's best to keep the same type + rather than converting it (e.g., to str) as this may result in unexpected + behavior. + + See Also: + - update_body_from_data: Synchronous body update without cleanup + - body property: Direct body access (STRONGLY DISCOURAGED) + + """ + # Close existing payload if it exists and needs closing + if self._body is not None: + await self._body.close() + self._update_body(body) + + def update_expect_continue(self, expect: bool = False) -> None: + if expect: + self.headers[hdrs.EXPECT] = "100-continue" + elif ( + hdrs.EXPECT in self.headers + and self.headers[hdrs.EXPECT].lower() == "100-continue" + ): + expect = True + + if expect: + self._continue = self.loop.create_future() + + def update_proxy( + self, + proxy: Optional[URL], + proxy_auth: Optional[BasicAuth], + proxy_headers: Optional[LooseHeaders], + ) -> None: + self.proxy = proxy + if proxy is None: + self.proxy_auth = None + self.proxy_headers = None + return + + if proxy_auth and not isinstance(proxy_auth, helpers.BasicAuth): + raise ValueError("proxy_auth must be None or BasicAuth() tuple") + self.proxy_auth = proxy_auth + + if proxy_headers is not None and not isinstance( + proxy_headers, (MultiDict, MultiDictProxy) + ): + proxy_headers = CIMultiDict(proxy_headers) + self.proxy_headers = proxy_headers + + async def write_bytes( + self, + writer: AbstractStreamWriter, + conn: "Connection", + content_length: Optional[int] = None, + ) -> None: + """ + Write the request body to the connection stream. + + This method handles writing different types of request bodies: + 1. Payload objects (using their specialized write_with_length method) + 2. Bytes/bytearray objects + 3. Iterable body content + + Args: + writer: The stream writer to write the body to + conn: The connection being used for this request + content_length: Optional maximum number of bytes to write from the body + (None means write the entire body) + + The method properly handles: + - Waiting for 100-Continue responses if required + - Content length constraints for chunked encoding + - Error handling for network issues, cancellation, and other exceptions + - Signaling EOF and timeout management + + Raises: + ClientOSError: When there's an OS-level error writing the body + ClientConnectionError: When there's a general connection error + asyncio.CancelledError: When the operation is cancelled + + """ + # 100 response + if self._continue is not None: + # Force headers to be sent before waiting for 100-continue + writer.send_headers() + await writer.drain() + await self._continue + + protocol = conn.protocol + assert protocol is not None + try: + # This should be a rare case but the + # self._body can be set to None while + # the task is being started or we wait above + # for the 100-continue response. + # The more likely case is we have an empty + # payload, but 100-continue is still expected. + if self._body is not None: + await self._body.write_with_length(writer, content_length) + except OSError as underlying_exc: + reraised_exc = underlying_exc + + # Distinguish between timeout and other OS errors for better error reporting + exc_is_not_timeout = underlying_exc.errno is not None or not isinstance( + underlying_exc, asyncio.TimeoutError + ) + if exc_is_not_timeout: + reraised_exc = ClientOSError( + underlying_exc.errno, + f"Can not write request body for {self.url !s}", + ) + + set_exception(protocol, reraised_exc, underlying_exc) + except asyncio.CancelledError: + # Body hasn't been fully sent, so connection can't be reused + conn.close() + raise + except Exception as underlying_exc: + set_exception( + protocol, + ClientConnectionError( + "Failed to send bytes into the underlying connection " + f"{conn !s}: {underlying_exc!r}", + ), + underlying_exc, + ) + else: + # Successfully wrote the body, signal EOF and start response timeout + await writer.write_eof() + protocol.start_timeout() + + async def send(self, conn: "Connection") -> "ClientResponse": + # Specify request target: + # - CONNECT request must send authority form URI + # - not CONNECT proxy must send absolute form URI + # - most common is origin form URI + if self.method == hdrs.METH_CONNECT: + connect_host = self.url.host_subcomponent + assert connect_host is not None + path = f"{connect_host}:{self.url.port}" + elif self.proxy and not self.is_ssl(): + path = str(self.url) + else: + path = self.url.raw_path_qs + + protocol = conn.protocol + assert protocol is not None + writer = StreamWriter( + protocol, + self.loop, + on_chunk_sent=( + functools.partial(self._on_chunk_request_sent, self.method, self.url) + if self._traces + else None + ), + on_headers_sent=( + functools.partial(self._on_headers_request_sent, self.method, self.url) + if self._traces + else None + ), + ) + + if self.compress: + writer.enable_compression(self.compress) # type: ignore[arg-type] + + if self.chunked is not None: + writer.enable_chunking() + + # set default content-type + if ( + self.method in self.POST_METHODS + and ( + self._skip_auto_headers is None + or hdrs.CONTENT_TYPE not in self._skip_auto_headers + ) + and hdrs.CONTENT_TYPE not in self.headers + ): + self.headers[hdrs.CONTENT_TYPE] = "application/octet-stream" + + v = self.version + if hdrs.CONNECTION not in self.headers: + if conn._connector.force_close: + if v == HttpVersion11: + self.headers[hdrs.CONNECTION] = "close" + elif v == HttpVersion10: + self.headers[hdrs.CONNECTION] = "keep-alive" + + # status + headers + status_line = f"{self.method} {path} HTTP/{v.major}.{v.minor}" + + # Buffer headers for potential coalescing with body + await writer.write_headers(status_line, self.headers) + + task: Optional["asyncio.Task[None]"] + if self._body or self._continue is not None or protocol.writing_paused: + coro = self.write_bytes(writer, conn, self._get_content_length()) + if sys.version_info >= (3, 12): + # Optimization for Python 3.12, try to write + # bytes immediately to avoid having to schedule + # the task on the event loop. + task = asyncio.Task(coro, loop=self.loop, eager_start=True) + else: + task = self.loop.create_task(coro) + if task.done(): + task = None + else: + self._writer = task + else: + # We have nothing to write because + # - there is no body + # - the protocol does not have writing paused + # - we are not waiting for a 100-continue response + protocol.start_timeout() + writer.set_eof() + task = None + response_class = self.response_class + assert response_class is not None + self.response = response_class( + self.method, + self.original_url, + writer=task, + continue100=self._continue, + timer=self._timer, + request_info=self.request_info, + traces=self._traces, + loop=self.loop, + session=self._session, + ) + return self.response + + async def close(self) -> None: + if self.__writer is not None: + try: + await self.__writer + except asyncio.CancelledError: + if ( + sys.version_info >= (3, 11) + and (task := asyncio.current_task()) + and task.cancelling() + ): + raise + + def terminate(self) -> None: + if self.__writer is not None: + if not self.loop.is_closed(): + self.__writer.cancel() + self.__writer.remove_done_callback(self.__reset_writer) + self.__writer = None + + async def _on_chunk_request_sent(self, method: str, url: URL, chunk: bytes) -> None: + for trace in self._traces: + await trace.send_request_chunk_sent(method, url, chunk) + + async def _on_headers_request_sent( + self, method: str, url: URL, headers: "CIMultiDict[str]" + ) -> None: + for trace in self._traces: + await trace.send_request_headers(method, url, headers) diff --git a/lib/python3.12/site-packages/aiohttp/client_ws.py b/lib/python3.12/site-packages/aiohttp/client_ws.py new file mode 100644 index 0000000000000000000000000000000000000000..daa57d1930b1b0cfc03594aaae62bd93da5da165 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/client_ws.py @@ -0,0 +1,428 @@ +"""WebSocket client for asyncio.""" + +import asyncio +import sys +from types import TracebackType +from typing import Any, Optional, Type, cast + +import attr + +from ._websocket.reader import WebSocketDataQueue +from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError +from .client_reqrep import ClientResponse +from .helpers import calculate_timeout_when, set_result +from .http import ( + WS_CLOSED_MESSAGE, + WS_CLOSING_MESSAGE, + WebSocketError, + WSCloseCode, + WSMessage, + WSMsgType, +) +from .http_websocket import _INTERNAL_RECEIVE_TYPES, WebSocketWriter +from .streams import EofStream +from .typedefs import ( + DEFAULT_JSON_DECODER, + DEFAULT_JSON_ENCODER, + JSONDecoder, + JSONEncoder, +) + +if sys.version_info >= (3, 11): + import asyncio as async_timeout +else: + import async_timeout + + +@attr.s(frozen=True, slots=True) +class ClientWSTimeout: + ws_receive = attr.ib(type=Optional[float], default=None) + ws_close = attr.ib(type=Optional[float], default=None) + + +DEFAULT_WS_CLIENT_TIMEOUT = ClientWSTimeout(ws_receive=None, ws_close=10.0) + + +class ClientWebSocketResponse: + def __init__( + self, + reader: WebSocketDataQueue, + writer: WebSocketWriter, + protocol: Optional[str], + response: ClientResponse, + timeout: ClientWSTimeout, + autoclose: bool, + autoping: bool, + loop: asyncio.AbstractEventLoop, + *, + heartbeat: Optional[float] = None, + compress: int = 0, + client_notakeover: bool = False, + ) -> None: + self._response = response + self._conn = response.connection + + self._writer = writer + self._reader = reader + self._protocol = protocol + self._closed = False + self._closing = False + self._close_code: Optional[int] = None + self._timeout = timeout + self._autoclose = autoclose + self._autoping = autoping + self._heartbeat = heartbeat + self._heartbeat_cb: Optional[asyncio.TimerHandle] = None + self._heartbeat_when: float = 0.0 + if heartbeat is not None: + self._pong_heartbeat = heartbeat / 2.0 + self._pong_response_cb: Optional[asyncio.TimerHandle] = None + self._loop = loop + self._waiting: bool = False + self._close_wait: Optional[asyncio.Future[None]] = None + self._exception: Optional[BaseException] = None + self._compress = compress + self._client_notakeover = client_notakeover + self._ping_task: Optional[asyncio.Task[None]] = None + + self._reset_heartbeat() + + def _cancel_heartbeat(self) -> None: + self._cancel_pong_response_cb() + if self._heartbeat_cb is not None: + self._heartbeat_cb.cancel() + self._heartbeat_cb = None + if self._ping_task is not None: + self._ping_task.cancel() + self._ping_task = None + + def _cancel_pong_response_cb(self) -> None: + if self._pong_response_cb is not None: + self._pong_response_cb.cancel() + self._pong_response_cb = None + + def _reset_heartbeat(self) -> None: + if self._heartbeat is None: + return + self._cancel_pong_response_cb() + loop = self._loop + assert loop is not None + conn = self._conn + timeout_ceil_threshold = ( + conn._connector._timeout_ceil_threshold if conn is not None else 5 + ) + now = loop.time() + when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold) + self._heartbeat_when = when + if self._heartbeat_cb is None: + # We do not cancel the previous heartbeat_cb here because + # it generates a significant amount of TimerHandle churn + # which causes asyncio to rebuild the heap frequently. + # Instead _send_heartbeat() will reschedule the next + # heartbeat if it fires too early. + self._heartbeat_cb = loop.call_at(when, self._send_heartbeat) + + def _send_heartbeat(self) -> None: + self._heartbeat_cb = None + loop = self._loop + now = loop.time() + if now < self._heartbeat_when: + # Heartbeat fired too early, reschedule + self._heartbeat_cb = loop.call_at( + self._heartbeat_when, self._send_heartbeat + ) + return + + conn = self._conn + timeout_ceil_threshold = ( + conn._connector._timeout_ceil_threshold if conn is not None else 5 + ) + when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold) + self._cancel_pong_response_cb() + self._pong_response_cb = loop.call_at(when, self._pong_not_received) + + coro = self._writer.send_frame(b"", WSMsgType.PING) + if sys.version_info >= (3, 12): + # Optimization for Python 3.12, try to send the ping + # immediately to avoid having to schedule + # the task on the event loop. + ping_task = asyncio.Task(coro, loop=loop, eager_start=True) + else: + ping_task = loop.create_task(coro) + + if not ping_task.done(): + self._ping_task = ping_task + ping_task.add_done_callback(self._ping_task_done) + else: + self._ping_task_done(ping_task) + + def _ping_task_done(self, task: "asyncio.Task[None]") -> None: + """Callback for when the ping task completes.""" + if not task.cancelled() and (exc := task.exception()): + self._handle_ping_pong_exception(exc) + self._ping_task = None + + def _pong_not_received(self) -> None: + self._handle_ping_pong_exception( + ServerTimeoutError(f"No PONG received after {self._pong_heartbeat} seconds") + ) + + def _handle_ping_pong_exception(self, exc: BaseException) -> None: + """Handle exceptions raised during ping/pong processing.""" + if self._closed: + return + self._set_closed() + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._exception = exc + self._response.close() + if self._waiting and not self._closing: + self._reader.feed_data(WSMessage(WSMsgType.ERROR, exc, None), 0) + + def _set_closed(self) -> None: + """Set the connection to closed. + + Cancel any heartbeat timers and set the closed flag. + """ + self._closed = True + self._cancel_heartbeat() + + def _set_closing(self) -> None: + """Set the connection to closing. + + Cancel any heartbeat timers and set the closing flag. + """ + self._closing = True + self._cancel_heartbeat() + + @property + def closed(self) -> bool: + return self._closed + + @property + def close_code(self) -> Optional[int]: + return self._close_code + + @property + def protocol(self) -> Optional[str]: + return self._protocol + + @property + def compress(self) -> int: + return self._compress + + @property + def client_notakeover(self) -> bool: + return self._client_notakeover + + def get_extra_info(self, name: str, default: Any = None) -> Any: + """extra info from connection transport""" + conn = self._response.connection + if conn is None: + return default + transport = conn.transport + if transport is None: + return default + return transport.get_extra_info(name, default) + + def exception(self) -> Optional[BaseException]: + return self._exception + + async def ping(self, message: bytes = b"") -> None: + await self._writer.send_frame(message, WSMsgType.PING) + + async def pong(self, message: bytes = b"") -> None: + await self._writer.send_frame(message, WSMsgType.PONG) + + async def send_frame( + self, message: bytes, opcode: WSMsgType, compress: Optional[int] = None + ) -> None: + """Send a frame over the websocket.""" + await self._writer.send_frame(message, opcode, compress) + + async def send_str(self, data: str, compress: Optional[int] = None) -> None: + if not isinstance(data, str): + raise TypeError("data argument must be str (%r)" % type(data)) + await self._writer.send_frame( + data.encode("utf-8"), WSMsgType.TEXT, compress=compress + ) + + async def send_bytes(self, data: bytes, compress: Optional[int] = None) -> None: + if not isinstance(data, (bytes, bytearray, memoryview)): + raise TypeError("data argument must be byte-ish (%r)" % type(data)) + await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress) + + async def send_json( + self, + data: Any, + compress: Optional[int] = None, + *, + dumps: JSONEncoder = DEFAULT_JSON_ENCODER, + ) -> None: + await self.send_str(dumps(data), compress=compress) + + async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool: + # we need to break `receive()` cycle first, + # `close()` may be called from different task + if self._waiting and not self._closing: + assert self._loop is not None + self._close_wait = self._loop.create_future() + self._set_closing() + self._reader.feed_data(WS_CLOSING_MESSAGE, 0) + await self._close_wait + + if self._closed: + return False + + self._set_closed() + try: + await self._writer.close(code, message) + except asyncio.CancelledError: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._response.close() + raise + except Exception as exc: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._exception = exc + self._response.close() + return True + + if self._close_code: + self._response.close() + return True + + while True: + try: + async with async_timeout.timeout(self._timeout.ws_close): + msg = await self._reader.read() + except asyncio.CancelledError: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._response.close() + raise + except Exception as exc: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._exception = exc + self._response.close() + return True + + if msg.type is WSMsgType.CLOSE: + self._close_code = msg.data + self._response.close() + return True + + async def receive(self, timeout: Optional[float] = None) -> WSMessage: + receive_timeout = timeout or self._timeout.ws_receive + + while True: + if self._waiting: + raise RuntimeError("Concurrent call to receive() is not allowed") + + if self._closed: + return WS_CLOSED_MESSAGE + elif self._closing: + await self.close() + return WS_CLOSED_MESSAGE + + try: + self._waiting = True + try: + if receive_timeout: + # Entering the context manager and creating + # Timeout() object can take almost 50% of the + # run time in this loop so we avoid it if + # there is no read timeout. + async with async_timeout.timeout(receive_timeout): + msg = await self._reader.read() + else: + msg = await self._reader.read() + self._reset_heartbeat() + finally: + self._waiting = False + if self._close_wait: + set_result(self._close_wait, None) + except (asyncio.CancelledError, asyncio.TimeoutError): + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + raise + except EofStream: + self._close_code = WSCloseCode.OK + await self.close() + return WSMessage(WSMsgType.CLOSED, None, None) + except ClientError: + # Likely ServerDisconnectedError when connection is lost + self._set_closed() + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + return WS_CLOSED_MESSAGE + except WebSocketError as exc: + self._close_code = exc.code + await self.close(code=exc.code) + return WSMessage(WSMsgType.ERROR, exc, None) + except Exception as exc: + self._exception = exc + self._set_closing() + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + await self.close() + return WSMessage(WSMsgType.ERROR, exc, None) + + if msg.type not in _INTERNAL_RECEIVE_TYPES: + # If its not a close/closing/ping/pong message + # we can return it immediately + return msg + + if msg.type is WSMsgType.CLOSE: + self._set_closing() + self._close_code = msg.data + if not self._closed and self._autoclose: + await self.close() + elif msg.type is WSMsgType.CLOSING: + self._set_closing() + elif msg.type is WSMsgType.PING and self._autoping: + await self.pong(msg.data) + continue + elif msg.type is WSMsgType.PONG and self._autoping: + continue + + return msg + + async def receive_str(self, *, timeout: Optional[float] = None) -> str: + msg = await self.receive(timeout) + if msg.type is not WSMsgType.TEXT: + raise WSMessageTypeError( + f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT" + ) + return cast(str, msg.data) + + async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes: + msg = await self.receive(timeout) + if msg.type is not WSMsgType.BINARY: + raise WSMessageTypeError( + f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY" + ) + return cast(bytes, msg.data) + + async def receive_json( + self, + *, + loads: JSONDecoder = DEFAULT_JSON_DECODER, + timeout: Optional[float] = None, + ) -> Any: + data = await self.receive_str(timeout=timeout) + return loads(data) + + def __aiter__(self) -> "ClientWebSocketResponse": + return self + + async def __anext__(self) -> WSMessage: + msg = await self.receive() + if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): + raise StopAsyncIteration + return msg + + async def __aenter__(self) -> "ClientWebSocketResponse": + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + await self.close() diff --git a/lib/python3.12/site-packages/aiohttp/compression_utils.py b/lib/python3.12/site-packages/aiohttp/compression_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e478d24c3d70105e6796007808462c4159f8c833 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/compression_utils.py @@ -0,0 +1,348 @@ +import asyncio +import sys +import zlib +from abc import ABC, abstractmethod +from concurrent.futures import Executor +from typing import Any, Final, Optional, Protocol, TypedDict, cast + +if sys.version_info >= (3, 12): + from collections.abc import Buffer +else: + from typing import Union + + Buffer = Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"] + +try: + try: + import brotlicffi as brotli + except ImportError: + import brotli + + HAS_BROTLI = True +except ImportError: # pragma: no cover + HAS_BROTLI = False + +try: + if sys.version_info >= (3, 14): + from compression.zstd import ZstdDecompressor # noqa: I900 + else: # TODO(PY314): Remove mentions of backports.zstd across codebase + from backports.zstd import ZstdDecompressor + + HAS_ZSTD = True +except ImportError: + HAS_ZSTD = False + + +MAX_SYNC_CHUNK_SIZE = 4096 +DEFAULT_MAX_DECOMPRESS_SIZE = 2**25 # 32MiB + +# Unlimited decompression constants - different libraries use different conventions +ZLIB_MAX_LENGTH_UNLIMITED = 0 # zlib uses 0 to mean unlimited +ZSTD_MAX_LENGTH_UNLIMITED = -1 # zstd uses -1 to mean unlimited + + +class ZLibCompressObjProtocol(Protocol): + def compress(self, data: Buffer) -> bytes: ... + def flush(self, mode: int = ..., /) -> bytes: ... + + +class ZLibDecompressObjProtocol(Protocol): + def decompress(self, data: Buffer, max_length: int = ...) -> bytes: ... + def flush(self, length: int = ..., /) -> bytes: ... + + @property + def eof(self) -> bool: ... + + +class ZLibBackendProtocol(Protocol): + MAX_WBITS: int + Z_FULL_FLUSH: int + Z_SYNC_FLUSH: int + Z_BEST_SPEED: int + Z_FINISH: int + + def compressobj( + self, + level: int = ..., + method: int = ..., + wbits: int = ..., + memLevel: int = ..., + strategy: int = ..., + zdict: Optional[Buffer] = ..., + ) -> ZLibCompressObjProtocol: ... + def decompressobj( + self, wbits: int = ..., zdict: Buffer = ... + ) -> ZLibDecompressObjProtocol: ... + + def compress( + self, data: Buffer, /, level: int = ..., wbits: int = ... + ) -> bytes: ... + def decompress( + self, data: Buffer, /, wbits: int = ..., bufsize: int = ... + ) -> bytes: ... + + +class CompressObjArgs(TypedDict, total=False): + wbits: int + strategy: int + level: int + + +class ZLibBackendWrapper: + def __init__(self, _zlib_backend: ZLibBackendProtocol): + self._zlib_backend: ZLibBackendProtocol = _zlib_backend + + @property + def name(self) -> str: + return getattr(self._zlib_backend, "__name__", "undefined") + + @property + def MAX_WBITS(self) -> int: + return self._zlib_backend.MAX_WBITS + + @property + def Z_FULL_FLUSH(self) -> int: + return self._zlib_backend.Z_FULL_FLUSH + + @property + def Z_SYNC_FLUSH(self) -> int: + return self._zlib_backend.Z_SYNC_FLUSH + + @property + def Z_BEST_SPEED(self) -> int: + return self._zlib_backend.Z_BEST_SPEED + + @property + def Z_FINISH(self) -> int: + return self._zlib_backend.Z_FINISH + + def compressobj(self, *args: Any, **kwargs: Any) -> ZLibCompressObjProtocol: + return self._zlib_backend.compressobj(*args, **kwargs) + + def decompressobj(self, *args: Any, **kwargs: Any) -> ZLibDecompressObjProtocol: + return self._zlib_backend.decompressobj(*args, **kwargs) + + def compress(self, data: Buffer, *args: Any, **kwargs: Any) -> bytes: + return self._zlib_backend.compress(data, *args, **kwargs) + + def decompress(self, data: Buffer, *args: Any, **kwargs: Any) -> bytes: + return self._zlib_backend.decompress(data, *args, **kwargs) + + # Everything not explicitly listed in the Protocol we just pass through + def __getattr__(self, attrname: str) -> Any: + return getattr(self._zlib_backend, attrname) + + +ZLibBackend: ZLibBackendWrapper = ZLibBackendWrapper(zlib) + + +def set_zlib_backend(new_zlib_backend: ZLibBackendProtocol) -> None: + ZLibBackend._zlib_backend = new_zlib_backend + + +def encoding_to_mode( + encoding: Optional[str] = None, + suppress_deflate_header: bool = False, +) -> int: + if encoding == "gzip": + return 16 + ZLibBackend.MAX_WBITS + + return -ZLibBackend.MAX_WBITS if suppress_deflate_header else ZLibBackend.MAX_WBITS + + +class DecompressionBaseHandler(ABC): + def __init__( + self, + executor: Optional[Executor] = None, + max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, + ): + """Base class for decompression handlers.""" + self._executor = executor + self._max_sync_chunk_size = max_sync_chunk_size + + @abstractmethod + def decompress_sync( + self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED + ) -> bytes: + """Decompress the given data.""" + + async def decompress( + self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED + ) -> bytes: + """Decompress the given data.""" + if ( + self._max_sync_chunk_size is not None + and len(data) > self._max_sync_chunk_size + ): + return await asyncio.get_event_loop().run_in_executor( + self._executor, self.decompress_sync, data, max_length + ) + return self.decompress_sync(data, max_length) + + +class ZLibCompressor: + def __init__( + self, + encoding: Optional[str] = None, + suppress_deflate_header: bool = False, + level: Optional[int] = None, + wbits: Optional[int] = None, + strategy: Optional[int] = None, + executor: Optional[Executor] = None, + max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, + ): + self._executor = executor + self._max_sync_chunk_size = max_sync_chunk_size + self._mode = ( + encoding_to_mode(encoding, suppress_deflate_header) + if wbits is None + else wbits + ) + self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend) + + kwargs: CompressObjArgs = {} + kwargs["wbits"] = self._mode + if strategy is not None: + kwargs["strategy"] = strategy + if level is not None: + kwargs["level"] = level + self._compressor = self._zlib_backend.compressobj(**kwargs) + + def compress_sync(self, data: bytes) -> bytes: + return self._compressor.compress(data) + + async def compress(self, data: bytes) -> bytes: + """Compress the data and returned the compressed bytes. + + Note that flush() must be called after the last call to compress() + + If the data size is large than the max_sync_chunk_size, the compression + will be done in the executor. Otherwise, the compression will be done + in the event loop. + + **WARNING: This method is NOT cancellation-safe when used with flush().** + If this operation is cancelled, the compressor state may be corrupted. + The connection MUST be closed after cancellation to avoid data corruption + in subsequent compress operations. + + For cancellation-safe compression (e.g., WebSocket), the caller MUST wrap + compress() + flush() + send operations in a shield and lock to ensure atomicity. + """ + # For large payloads, offload compression to executor to avoid blocking event loop + should_use_executor = ( + self._max_sync_chunk_size is not None + and len(data) > self._max_sync_chunk_size + ) + if should_use_executor: + return await asyncio.get_running_loop().run_in_executor( + self._executor, self._compressor.compress, data + ) + return self.compress_sync(data) + + def flush(self, mode: Optional[int] = None) -> bytes: + """Flush the compressor synchronously. + + **WARNING: This method is NOT cancellation-safe when called after compress().** + The flush() operation accesses shared compressor state. If compress() was + cancelled, calling flush() may result in corrupted data. The connection MUST + be closed after compress() cancellation. + + For cancellation-safe compression (e.g., WebSocket), the caller MUST wrap + compress() + flush() + send operations in a shield and lock to ensure atomicity. + """ + return self._compressor.flush( + mode if mode is not None else self._zlib_backend.Z_FINISH + ) + + +class ZLibDecompressor(DecompressionBaseHandler): + def __init__( + self, + encoding: Optional[str] = None, + suppress_deflate_header: bool = False, + executor: Optional[Executor] = None, + max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, + ): + super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) + self._mode = encoding_to_mode(encoding, suppress_deflate_header) + self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend) + self._decompressor = self._zlib_backend.decompressobj(wbits=self._mode) + + def decompress_sync( + self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED + ) -> bytes: + return self._decompressor.decompress(data, max_length) + + def flush(self, length: int = 0) -> bytes: + return ( + self._decompressor.flush(length) + if length > 0 + else self._decompressor.flush() + ) + + @property + def eof(self) -> bool: + return self._decompressor.eof + + +class BrotliDecompressor(DecompressionBaseHandler): + # Supports both 'brotlipy' and 'Brotli' packages + # since they share an import name. The top branches + # are for 'brotlipy' and bottom branches for 'Brotli' + def __init__( + self, + executor: Optional[Executor] = None, + max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, + ) -> None: + """Decompress data using the Brotli library.""" + if not HAS_BROTLI: + raise RuntimeError( + "The brotli decompression is not available. " + "Please install `Brotli` module" + ) + self._obj = brotli.Decompressor() + super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) + + def decompress_sync( + self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED + ) -> bytes: + """Decompress the given data.""" + if hasattr(self._obj, "decompress"): + return cast(bytes, self._obj.decompress(data, max_length)) + return cast(bytes, self._obj.process(data, max_length)) + + def flush(self) -> bytes: + """Flush the decompressor.""" + if hasattr(self._obj, "flush"): + return cast(bytes, self._obj.flush()) + return b"" + + +class ZSTDDecompressor(DecompressionBaseHandler): + def __init__( + self, + executor: Optional[Executor] = None, + max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, + ) -> None: + if not HAS_ZSTD: + raise RuntimeError( + "The zstd decompression is not available. " + "Please install `backports.zstd` module" + ) + self._obj = ZstdDecompressor() + super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) + + def decompress_sync( + self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED + ) -> bytes: + # zstd uses -1 for unlimited, while zlib uses 0 for unlimited + # Convert the zlib convention (0=unlimited) to zstd convention (-1=unlimited) + zstd_max_length = ( + ZSTD_MAX_LENGTH_UNLIMITED + if max_length == ZLIB_MAX_LENGTH_UNLIMITED + else max_length + ) + return self._obj.decompress(data, zstd_max_length) + + def flush(self) -> bytes: + return b"" diff --git a/lib/python3.12/site-packages/aiohttp/connector.py b/lib/python3.12/site-packages/aiohttp/connector.py new file mode 100644 index 0000000000000000000000000000000000000000..290a42400f9af5257e1ca1511c3db7bd0570e9bd --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/connector.py @@ -0,0 +1,1842 @@ +import asyncio +import functools +import random +import socket +import sys +import traceback +import warnings +from collections import OrderedDict, defaultdict, deque +from contextlib import suppress +from http import HTTPStatus +from itertools import chain, cycle, islice +from time import monotonic +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + DefaultDict, + Deque, + Dict, + Iterator, + List, + Literal, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, + cast, +) + +import aiohappyeyeballs +from aiohappyeyeballs import AddrInfoType, SocketFactoryType + +from . import hdrs, helpers +from .abc import AbstractResolver, ResolveResult +from .client_exceptions import ( + ClientConnectionError, + ClientConnectorCertificateError, + ClientConnectorDNSError, + ClientConnectorError, + ClientConnectorSSLError, + ClientHttpProxyError, + ClientProxyConnectionError, + ServerFingerprintMismatch, + UnixClientConnectorError, + cert_errors, + ssl_errors, +) +from .client_proto import ResponseHandler +from .client_reqrep import ClientRequest, Fingerprint, _merge_ssl_params +from .helpers import ( + _SENTINEL, + ceil_timeout, + is_ip_address, + noop, + sentinel, + set_exception, + set_result, +) +from .log import client_logger +from .resolver import DefaultResolver + +if sys.version_info >= (3, 12): + from collections.abc import Buffer +else: + Buffer = Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"] + +if TYPE_CHECKING: + import ssl + + SSLContext = ssl.SSLContext +else: + try: + import ssl + + SSLContext = ssl.SSLContext + except ImportError: # pragma: no cover + ssl = None # type: ignore[assignment] + SSLContext = object # type: ignore[misc,assignment] + +EMPTY_SCHEMA_SET = frozenset({""}) +HTTP_SCHEMA_SET = frozenset({"http", "https"}) +WS_SCHEMA_SET = frozenset({"ws", "wss"}) + +HTTP_AND_EMPTY_SCHEMA_SET = HTTP_SCHEMA_SET | EMPTY_SCHEMA_SET +HIGH_LEVEL_SCHEMA_SET = HTTP_AND_EMPTY_SCHEMA_SET | WS_SCHEMA_SET + +NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < ( + 3, + 13, + 1, +) or sys.version_info < (3, 12, 7) +# Cleanup closed is no longer needed after https://github.com/python/cpython/pull/118960 +# which first appeared in Python 3.12.7 and 3.13.1 + + +__all__ = ( + "BaseConnector", + "TCPConnector", + "UnixConnector", + "NamedPipeConnector", + "AddrInfoType", + "SocketFactoryType", +) + + +if TYPE_CHECKING: + from .client import ClientTimeout + from .client_reqrep import ConnectionKey + from .tracing import Trace + + +class _DeprecationWaiter: + __slots__ = ("_awaitable", "_awaited") + + def __init__(self, awaitable: Awaitable[Any]) -> None: + self._awaitable = awaitable + self._awaited = False + + def __await__(self) -> Any: + self._awaited = True + return self._awaitable.__await__() + + def __del__(self) -> None: + if not self._awaited: + warnings.warn( + "Connector.close() is a coroutine, " + "please use await connector.close()", + DeprecationWarning, + ) + + +async def _wait_for_close(waiters: List[Awaitable[object]]) -> None: + """Wait for all waiters to finish closing.""" + results = await asyncio.gather(*waiters, return_exceptions=True) + for res in results: + if isinstance(res, Exception): + client_logger.debug("Error while closing connector: %r", res) + + +class Connection: + + _source_traceback = None + + def __init__( + self, + connector: "BaseConnector", + key: "ConnectionKey", + protocol: ResponseHandler, + loop: asyncio.AbstractEventLoop, + ) -> None: + self._key = key + self._connector = connector + self._loop = loop + self._protocol: Optional[ResponseHandler] = protocol + self._callbacks: List[Callable[[], None]] = [] + + if loop.get_debug(): + self._source_traceback = traceback.extract_stack(sys._getframe(1)) + + def __repr__(self) -> str: + return f"Connection<{self._key}>" + + def __del__(self, _warnings: Any = warnings) -> None: + if self._protocol is not None: + kwargs = {"source": self} + _warnings.warn(f"Unclosed connection {self!r}", ResourceWarning, **kwargs) + if self._loop.is_closed(): + return + + self._connector._release(self._key, self._protocol, should_close=True) + + context = {"client_connection": self, "message": "Unclosed connection"} + if self._source_traceback is not None: + context["source_traceback"] = self._source_traceback + self._loop.call_exception_handler(context) + + def __bool__(self) -> Literal[True]: + """Force subclasses to not be falsy, to make checks simpler.""" + return True + + @property + def loop(self) -> asyncio.AbstractEventLoop: + warnings.warn( + "connector.loop property is deprecated", DeprecationWarning, stacklevel=2 + ) + return self._loop + + @property + def transport(self) -> Optional[asyncio.Transport]: + if self._protocol is None: + return None + return self._protocol.transport + + @property + def protocol(self) -> Optional[ResponseHandler]: + return self._protocol + + def add_callback(self, callback: Callable[[], None]) -> None: + if callback is not None: + self._callbacks.append(callback) + + def _notify_release(self) -> None: + callbacks, self._callbacks = self._callbacks[:], [] + + for cb in callbacks: + with suppress(Exception): + cb() + + def close(self) -> None: + self._notify_release() + + if self._protocol is not None: + self._connector._release(self._key, self._protocol, should_close=True) + self._protocol = None + + def release(self) -> None: + self._notify_release() + + if self._protocol is not None: + self._connector._release(self._key, self._protocol) + self._protocol = None + + @property + def closed(self) -> bool: + return self._protocol is None or not self._protocol.is_connected() + + +class _ConnectTunnelConnection(Connection): + """Special connection wrapper for CONNECT tunnels that must never be pooled. + + This connection wraps the proxy connection that will be upgraded with TLS. + It must never be released to the pool because: + 1. Its 'closed' future will never complete, causing session.close() to hang + 2. It represents an intermediate state, not a reusable connection + 3. The real connection (with TLS) will be created separately + """ + + def release(self) -> None: + """Do nothing - don't pool or close the connection. + + These connections are an intermediate state during the CONNECT tunnel + setup and will be cleaned up naturally after the TLS upgrade. If they + were to be pooled, they would never be properly closed, causing + session.close() to wait forever for their 'closed' future. + """ + + +class _TransportPlaceholder: + """placeholder for BaseConnector.connect function""" + + __slots__ = ("closed", "transport") + + def __init__(self, closed_future: asyncio.Future[Optional[Exception]]) -> None: + """Initialize a placeholder for a transport.""" + self.closed = closed_future + self.transport = None + + def close(self) -> None: + """Close the placeholder.""" + + def abort(self) -> None: + """Abort the placeholder (does nothing).""" + + +class BaseConnector: + """Base connector class. + + keepalive_timeout - (optional) Keep-alive timeout. + force_close - Set to True to force close and do reconnect + after each request (and between redirects). + limit - The total number of simultaneous connections. + limit_per_host - Number of simultaneous connections to one host. + enable_cleanup_closed - Enables clean-up closed ssl transports. + Disabled by default. + timeout_ceil_threshold - Trigger ceiling of timeout values when + it's above timeout_ceil_threshold. + loop - Optional event loop. + """ + + _closed = True # prevent AttributeError in __del__ if ctor was failed + _source_traceback = None + + # abort transport after 2 seconds (cleanup broken connections) + _cleanup_closed_period = 2.0 + + allowed_protocol_schema_set = HIGH_LEVEL_SCHEMA_SET + + def __init__( + self, + *, + keepalive_timeout: Union[object, None, float] = sentinel, + force_close: bool = False, + limit: int = 100, + limit_per_host: int = 0, + enable_cleanup_closed: bool = False, + loop: Optional[asyncio.AbstractEventLoop] = None, + timeout_ceil_threshold: float = 5, + ) -> None: + + if force_close: + if keepalive_timeout is not None and keepalive_timeout is not sentinel: + raise ValueError( + "keepalive_timeout cannot be set if force_close is True" + ) + else: + if keepalive_timeout is sentinel: + keepalive_timeout = 15.0 + + loop = loop or asyncio.get_running_loop() + self._timeout_ceil_threshold = timeout_ceil_threshold + + self._closed = False + if loop.get_debug(): + self._source_traceback = traceback.extract_stack(sys._getframe(1)) + + # Connection pool of reusable connections. + # We use a deque to store connections because it has O(1) popleft() + # and O(1) append() operations to implement a FIFO queue. + self._conns: DefaultDict[ + ConnectionKey, Deque[Tuple[ResponseHandler, float]] + ] = defaultdict(deque) + self._limit = limit + self._limit_per_host = limit_per_host + self._acquired: Set[ResponseHandler] = set() + self._acquired_per_host: DefaultDict[ConnectionKey, Set[ResponseHandler]] = ( + defaultdict(set) + ) + self._keepalive_timeout = cast(float, keepalive_timeout) + self._force_close = force_close + + # {host_key: FIFO list of waiters} + # The FIFO is implemented with an OrderedDict with None keys because + # python does not have an ordered set. + self._waiters: DefaultDict[ + ConnectionKey, OrderedDict[asyncio.Future[None], None] + ] = defaultdict(OrderedDict) + + self._loop = loop + self._factory = functools.partial(ResponseHandler, loop=loop) + + # start keep-alive connection cleanup task + self._cleanup_handle: Optional[asyncio.TimerHandle] = None + + # start cleanup closed transports task + self._cleanup_closed_handle: Optional[asyncio.TimerHandle] = None + + if enable_cleanup_closed and not NEEDS_CLEANUP_CLOSED: + warnings.warn( + "enable_cleanup_closed ignored because " + "https://github.com/python/cpython/pull/118960 is fixed " + f"in Python version {sys.version_info}", + DeprecationWarning, + stacklevel=2, + ) + enable_cleanup_closed = False + + self._cleanup_closed_disabled = not enable_cleanup_closed + self._cleanup_closed_transports: List[Optional[asyncio.Transport]] = [] + self._placeholder_future: asyncio.Future[Optional[Exception]] = ( + loop.create_future() + ) + self._placeholder_future.set_result(None) + self._cleanup_closed() + + def __del__(self, _warnings: Any = warnings) -> None: + if self._closed: + return + if not self._conns: + return + + conns = [repr(c) for c in self._conns.values()] + + self._close() + + kwargs = {"source": self} + _warnings.warn(f"Unclosed connector {self!r}", ResourceWarning, **kwargs) + context = { + "connector": self, + "connections": conns, + "message": "Unclosed connector", + } + if self._source_traceback is not None: + context["source_traceback"] = self._source_traceback + self._loop.call_exception_handler(context) + + def __enter__(self) -> "BaseConnector": + warnings.warn( + '"with Connector():" is deprecated, ' + 'use "async with Connector():" instead', + DeprecationWarning, + ) + return self + + def __exit__(self, *exc: Any) -> None: + self._close() + + async def __aenter__(self) -> "BaseConnector": + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]] = None, + exc_value: Optional[BaseException] = None, + exc_traceback: Optional[TracebackType] = None, + ) -> None: + await self.close() + + @property + def force_close(self) -> bool: + """Ultimately close connection on releasing if True.""" + return self._force_close + + @property + def limit(self) -> int: + """The total number for simultaneous connections. + + If limit is 0 the connector has no limit. + The default limit size is 100. + """ + return self._limit + + @property + def limit_per_host(self) -> int: + """The limit for simultaneous connections to the same endpoint. + + Endpoints are the same if they are have equal + (host, port, is_ssl) triple. + """ + return self._limit_per_host + + def _cleanup(self) -> None: + """Cleanup unused transports.""" + if self._cleanup_handle: + self._cleanup_handle.cancel() + # _cleanup_handle should be unset, otherwise _release() will not + # recreate it ever! + self._cleanup_handle = None + + now = monotonic() + timeout = self._keepalive_timeout + + if self._conns: + connections = defaultdict(deque) + deadline = now - timeout + for key, conns in self._conns.items(): + alive: Deque[Tuple[ResponseHandler, float]] = deque() + for proto, use_time in conns: + if proto.is_connected() and use_time - deadline >= 0: + alive.append((proto, use_time)) + continue + transport = proto.transport + proto.close() + if not self._cleanup_closed_disabled and key.is_ssl: + self._cleanup_closed_transports.append(transport) + + if alive: + connections[key] = alive + + self._conns = connections + + if self._conns: + self._cleanup_handle = helpers.weakref_handle( + self, + "_cleanup", + timeout, + self._loop, + timeout_ceil_threshold=self._timeout_ceil_threshold, + ) + + def _cleanup_closed(self) -> None: + """Double confirmation for transport close. + + Some broken ssl servers may leave socket open without proper close. + """ + if self._cleanup_closed_handle: + self._cleanup_closed_handle.cancel() + + for transport in self._cleanup_closed_transports: + if transport is not None: + transport.abort() + + self._cleanup_closed_transports = [] + + if not self._cleanup_closed_disabled: + self._cleanup_closed_handle = helpers.weakref_handle( + self, + "_cleanup_closed", + self._cleanup_closed_period, + self._loop, + timeout_ceil_threshold=self._timeout_ceil_threshold, + ) + + def close(self, *, abort_ssl: bool = False) -> Awaitable[None]: + """Close all opened transports. + + :param abort_ssl: If True, SSL connections will be aborted immediately + without performing the shutdown handshake. This provides + faster cleanup at the cost of less graceful disconnection. + """ + if not (waiters := self._close(abort_ssl=abort_ssl)): + # If there are no connections to close, we can return a noop + # awaitable to avoid scheduling a task on the event loop. + return _DeprecationWaiter(noop()) + coro = _wait_for_close(waiters) + if sys.version_info >= (3, 12): + # Optimization for Python 3.12, try to close connections + # immediately to avoid having to schedule the task on the event loop. + task = asyncio.Task(coro, loop=self._loop, eager_start=True) + else: + task = self._loop.create_task(coro) + return _DeprecationWaiter(task) + + def _close(self, *, abort_ssl: bool = False) -> List[Awaitable[object]]: + waiters: List[Awaitable[object]] = [] + + if self._closed: + return waiters + + self._closed = True + + try: + if self._loop.is_closed(): + return waiters + + # cancel cleanup task + if self._cleanup_handle: + self._cleanup_handle.cancel() + + # cancel cleanup close task + if self._cleanup_closed_handle: + self._cleanup_closed_handle.cancel() + + for data in self._conns.values(): + for proto, _ in data: + if ( + abort_ssl + and proto.transport + and proto.transport.get_extra_info("sslcontext") is not None + ): + proto.abort() + else: + proto.close() + if closed := proto.closed: + waiters.append(closed) + + for proto in self._acquired: + if ( + abort_ssl + and proto.transport + and proto.transport.get_extra_info("sslcontext") is not None + ): + proto.abort() + else: + proto.close() + if closed := proto.closed: + waiters.append(closed) + + for transport in self._cleanup_closed_transports: + if transport is not None: + transport.abort() + + return waiters + + finally: + self._conns.clear() + self._acquired.clear() + for keyed_waiters in self._waiters.values(): + for keyed_waiter in keyed_waiters: + keyed_waiter.cancel() + self._waiters.clear() + self._cleanup_handle = None + self._cleanup_closed_transports.clear() + self._cleanup_closed_handle = None + + @property + def closed(self) -> bool: + """Is connector closed. + + A readonly property. + """ + return self._closed + + def _available_connections(self, key: "ConnectionKey") -> int: + """ + Return number of available connections. + + The limit, limit_per_host and the connection key are taken into account. + + If it returns less than 1 means that there are no connections + available. + """ + # check total available connections + # If there are no limits, this will always return 1 + total_remain = 1 + + if self._limit and (total_remain := self._limit - len(self._acquired)) <= 0: + return total_remain + + # check limit per host + if host_remain := self._limit_per_host: + if acquired := self._acquired_per_host.get(key): + host_remain -= len(acquired) + if total_remain > host_remain: + return host_remain + + return total_remain + + def _update_proxy_auth_header_and_build_proxy_req( + self, req: ClientRequest + ) -> ClientRequest: + """Set Proxy-Authorization header for non-SSL proxy requests and builds the proxy request for SSL proxy requests.""" + url = req.proxy + assert url is not None + headers: Dict[str, str] = {} + if req.proxy_headers is not None: + headers = req.proxy_headers # type: ignore[assignment] + headers[hdrs.HOST] = req.headers[hdrs.HOST] + proxy_req = ClientRequest( + hdrs.METH_GET, + url, + headers=headers, + auth=req.proxy_auth, + loop=self._loop, + ssl=req.ssl, + ) + auth = proxy_req.headers.pop(hdrs.AUTHORIZATION, None) + if auth is not None: + if not req.is_ssl(): + req.headers[hdrs.PROXY_AUTHORIZATION] = auth + else: + proxy_req.headers[hdrs.PROXY_AUTHORIZATION] = auth + return proxy_req + + async def connect( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> Connection: + """Get from pool or create new connection.""" + key = req.connection_key + if (conn := await self._get(key, traces)) is not None: + # If we do not have to wait and we can get a connection from the pool + # we can avoid the timeout ceil logic and directly return the connection + if req.proxy: + self._update_proxy_auth_header_and_build_proxy_req(req) + return conn + + async with ceil_timeout(timeout.connect, timeout.ceil_threshold): + if self._available_connections(key) <= 0: + await self._wait_for_available_connection(key, traces) + if (conn := await self._get(key, traces)) is not None: + if req.proxy: + self._update_proxy_auth_header_and_build_proxy_req(req) + return conn + + placeholder = cast( + ResponseHandler, _TransportPlaceholder(self._placeholder_future) + ) + self._acquired.add(placeholder) + if self._limit_per_host: + self._acquired_per_host[key].add(placeholder) + + try: + # Traces are done inside the try block to ensure that the + # that the placeholder is still cleaned up if an exception + # is raised. + if traces: + for trace in traces: + await trace.send_connection_create_start() + proto = await self._create_connection(req, traces, timeout) + if traces: + for trace in traces: + await trace.send_connection_create_end() + except BaseException: + self._release_acquired(key, placeholder) + raise + else: + if self._closed: + proto.close() + raise ClientConnectionError("Connector is closed.") + + # The connection was successfully created, drop the placeholder + # and add the real connection to the acquired set. There should + # be no awaits after the proto is added to the acquired set + # to ensure that the connection is not left in the acquired set + # on cancellation. + self._acquired.remove(placeholder) + self._acquired.add(proto) + if self._limit_per_host: + acquired_per_host = self._acquired_per_host[key] + acquired_per_host.remove(placeholder) + acquired_per_host.add(proto) + return Connection(self, key, proto, self._loop) + + async def _wait_for_available_connection( + self, key: "ConnectionKey", traces: List["Trace"] + ) -> None: + """Wait for an available connection slot.""" + # We loop here because there is a race between + # the connection limit check and the connection + # being acquired. If the connection is acquired + # between the check and the await statement, we + # need to loop again to check if the connection + # slot is still available. + attempts = 0 + while True: + fut: asyncio.Future[None] = self._loop.create_future() + keyed_waiters = self._waiters[key] + keyed_waiters[fut] = None + if attempts: + # If we have waited before, we need to move the waiter + # to the front of the queue as otherwise we might get + # starved and hit the timeout. + keyed_waiters.move_to_end(fut, last=False) + + try: + # Traces happen in the try block to ensure that the + # the waiter is still cleaned up if an exception is raised. + if traces: + for trace in traces: + await trace.send_connection_queued_start() + await fut + if traces: + for trace in traces: + await trace.send_connection_queued_end() + finally: + # pop the waiter from the queue if its still + # there and not already removed by _release_waiter + keyed_waiters.pop(fut, None) + if not self._waiters.get(key, True): + del self._waiters[key] + + if self._available_connections(key) > 0: + break + attempts += 1 + + async def _get( + self, key: "ConnectionKey", traces: List["Trace"] + ) -> Optional[Connection]: + """Get next reusable connection for the key or None. + + The connection will be marked as acquired. + """ + if (conns := self._conns.get(key)) is None: + return None + + t1 = monotonic() + while conns: + proto, t0 = conns.popleft() + # We will we reuse the connection if its connected and + # the keepalive timeout has not been exceeded + if proto.is_connected() and t1 - t0 <= self._keepalive_timeout: + if not conns: + # The very last connection was reclaimed: drop the key + del self._conns[key] + self._acquired.add(proto) + if self._limit_per_host: + self._acquired_per_host[key].add(proto) + if traces: + for trace in traces: + try: + await trace.send_connection_reuseconn() + except BaseException: + self._release_acquired(key, proto) + raise + return Connection(self, key, proto, self._loop) + + # Connection cannot be reused, close it + transport = proto.transport + proto.close() + # only for SSL transports + if not self._cleanup_closed_disabled and key.is_ssl: + self._cleanup_closed_transports.append(transport) + + # No more connections: drop the key + del self._conns[key] + return None + + def _release_waiter(self) -> None: + """ + Iterates over all waiters until one to be released is found. + + The one to be released is not finished and + belongs to a host that has available connections. + """ + if not self._waiters: + return + + # Having the dict keys ordered this avoids to iterate + # at the same order at each call. + queues = list(self._waiters) + random.shuffle(queues) + + for key in queues: + if self._available_connections(key) < 1: + continue + + waiters = self._waiters[key] + while waiters: + waiter, _ = waiters.popitem(last=False) + if not waiter.done(): + waiter.set_result(None) + return + + def _release_acquired(self, key: "ConnectionKey", proto: ResponseHandler) -> None: + """Release acquired connection.""" + if self._closed: + # acquired connection is already released on connector closing + return + + self._acquired.discard(proto) + if self._limit_per_host and (conns := self._acquired_per_host.get(key)): + conns.discard(proto) + if not conns: + del self._acquired_per_host[key] + self._release_waiter() + + def _release( + self, + key: "ConnectionKey", + protocol: ResponseHandler, + *, + should_close: bool = False, + ) -> None: + if self._closed: + # acquired connection is already released on connector closing + return + + self._release_acquired(key, protocol) + + if self._force_close or should_close or protocol.should_close: + transport = protocol.transport + protocol.close() + + if key.is_ssl and not self._cleanup_closed_disabled: + self._cleanup_closed_transports.append(transport) + return + + self._conns[key].append((protocol, monotonic())) + + if self._cleanup_handle is None: + self._cleanup_handle = helpers.weakref_handle( + self, + "_cleanup", + self._keepalive_timeout, + self._loop, + timeout_ceil_threshold=self._timeout_ceil_threshold, + ) + + async def _create_connection( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> ResponseHandler: + raise NotImplementedError() + + +class _DNSCacheTable: + def __init__(self, ttl: Optional[float] = None) -> None: + self._addrs_rr: Dict[Tuple[str, int], Tuple[Iterator[ResolveResult], int]] = {} + self._timestamps: Dict[Tuple[str, int], float] = {} + self._ttl = ttl + + def __contains__(self, host: object) -> bool: + return host in self._addrs_rr + + def add(self, key: Tuple[str, int], addrs: List[ResolveResult]) -> None: + self._addrs_rr[key] = (cycle(addrs), len(addrs)) + + if self._ttl is not None: + self._timestamps[key] = monotonic() + + def remove(self, key: Tuple[str, int]) -> None: + self._addrs_rr.pop(key, None) + + if self._ttl is not None: + self._timestamps.pop(key, None) + + def clear(self) -> None: + self._addrs_rr.clear() + self._timestamps.clear() + + def next_addrs(self, key: Tuple[str, int]) -> List[ResolveResult]: + loop, length = self._addrs_rr[key] + addrs = list(islice(loop, length)) + # Consume one more element to shift internal state of `cycle` + next(loop) + return addrs + + def expired(self, key: Tuple[str, int]) -> bool: + if self._ttl is None: + return False + + return self._timestamps[key] + self._ttl < monotonic() + + +def _make_ssl_context(verified: bool) -> SSLContext: + """Create SSL context. + + This method is not async-friendly and should be called from a thread + because it will load certificates from disk and do other blocking I/O. + """ + if ssl is None: + # No ssl support + return None + if verified: + sslcontext = ssl.create_default_context() + else: + sslcontext = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + sslcontext.options |= ssl.OP_NO_SSLv2 + sslcontext.options |= ssl.OP_NO_SSLv3 + sslcontext.check_hostname = False + sslcontext.verify_mode = ssl.CERT_NONE + sslcontext.options |= ssl.OP_NO_COMPRESSION + sslcontext.set_default_verify_paths() + sslcontext.set_alpn_protocols(("http/1.1",)) + return sslcontext + + +# The default SSLContext objects are created at import time +# since they do blocking I/O to load certificates from disk, +# and imports should always be done before the event loop starts +# or in a thread. +_SSL_CONTEXT_VERIFIED = _make_ssl_context(True) +_SSL_CONTEXT_UNVERIFIED = _make_ssl_context(False) + + +class TCPConnector(BaseConnector): + """TCP connector. + + verify_ssl - Set to True to check ssl certifications. + fingerprint - Pass the binary sha256 + digest of the expected certificate in DER format to verify + that the certificate the server presents matches. See also + https://en.wikipedia.org/wiki/HTTP_Public_Key_Pinning + resolver - Enable DNS lookups and use this + resolver + use_dns_cache - Use memory cache for DNS lookups. + ttl_dns_cache - Max seconds having cached a DNS entry, None forever. + family - socket address family + local_addr - local tuple of (host, port) to bind socket to + + keepalive_timeout - (optional) Keep-alive timeout. + force_close - Set to True to force close and do reconnect + after each request (and between redirects). + limit - The total number of simultaneous connections. + limit_per_host - Number of simultaneous connections to one host. + enable_cleanup_closed - Enables clean-up closed ssl transports. + Disabled by default. + happy_eyeballs_delay - This is the “Connection Attempt Delay” + as defined in RFC 8305. To disable + the happy eyeballs algorithm, set to None. + interleave - “First Address Family Count” as defined in RFC 8305 + loop - Optional event loop. + socket_factory - A SocketFactoryType function that, if supplied, + will be used to create sockets given an + AddrInfoType. + ssl_shutdown_timeout - DEPRECATED. Will be removed in aiohttp 4.0. + Grace period for SSL shutdown handshake on TLS + connections. Default is 0 seconds (immediate abort). + This parameter allowed for a clean SSL shutdown by + notifying the remote peer of connection closure, + while avoiding excessive delays during connector cleanup. + Note: Only takes effect on Python 3.11+. + """ + + allowed_protocol_schema_set = HIGH_LEVEL_SCHEMA_SET | frozenset({"tcp"}) + + def __init__( + self, + *, + verify_ssl: bool = True, + fingerprint: Optional[bytes] = None, + use_dns_cache: bool = True, + ttl_dns_cache: Optional[int] = 10, + family: socket.AddressFamily = socket.AddressFamily.AF_UNSPEC, + ssl_context: Optional[SSLContext] = None, + ssl: Union[bool, Fingerprint, SSLContext] = True, + local_addr: Optional[Tuple[str, int]] = None, + resolver: Optional[AbstractResolver] = None, + keepalive_timeout: Union[None, float, object] = sentinel, + force_close: bool = False, + limit: int = 100, + limit_per_host: int = 0, + enable_cleanup_closed: bool = False, + loop: Optional[asyncio.AbstractEventLoop] = None, + timeout_ceil_threshold: float = 5, + happy_eyeballs_delay: Optional[float] = 0.25, + interleave: Optional[int] = None, + socket_factory: Optional[SocketFactoryType] = None, + ssl_shutdown_timeout: Union[_SENTINEL, None, float] = sentinel, + ): + super().__init__( + keepalive_timeout=keepalive_timeout, + force_close=force_close, + limit=limit, + limit_per_host=limit_per_host, + enable_cleanup_closed=enable_cleanup_closed, + loop=loop, + timeout_ceil_threshold=timeout_ceil_threshold, + ) + + self._ssl = _merge_ssl_params(ssl, verify_ssl, ssl_context, fingerprint) + + self._resolver: AbstractResolver + if resolver is None: + self._resolver = DefaultResolver(loop=self._loop) + self._resolver_owner = True + else: + self._resolver = resolver + self._resolver_owner = False + + self._use_dns_cache = use_dns_cache + self._cached_hosts = _DNSCacheTable(ttl=ttl_dns_cache) + self._throttle_dns_futures: Dict[ + Tuple[str, int], Set["asyncio.Future[None]"] + ] = {} + self._family = family + self._local_addr_infos = aiohappyeyeballs.addr_to_addr_infos(local_addr) + self._happy_eyeballs_delay = happy_eyeballs_delay + self._interleave = interleave + self._resolve_host_tasks: Set["asyncio.Task[List[ResolveResult]]"] = set() + self._socket_factory = socket_factory + self._ssl_shutdown_timeout: Optional[float] + # Handle ssl_shutdown_timeout with warning for Python < 3.11 + if ssl_shutdown_timeout is sentinel: + self._ssl_shutdown_timeout = 0 + else: + # Deprecation warning for ssl_shutdown_timeout parameter + warnings.warn( + "The ssl_shutdown_timeout parameter is deprecated and will be removed in aiohttp 4.0", + DeprecationWarning, + stacklevel=2, + ) + if ( + sys.version_info < (3, 11) + and ssl_shutdown_timeout is not None + and ssl_shutdown_timeout != 0 + ): + warnings.warn( + f"ssl_shutdown_timeout={ssl_shutdown_timeout} is ignored on Python < 3.11; " + "only ssl_shutdown_timeout=0 is supported. The timeout will be ignored.", + RuntimeWarning, + stacklevel=2, + ) + self._ssl_shutdown_timeout = ssl_shutdown_timeout + + def _close(self, *, abort_ssl: bool = False) -> List[Awaitable[object]]: + """Close all ongoing DNS calls.""" + for fut in chain.from_iterable(self._throttle_dns_futures.values()): + fut.cancel() + + waiters = super()._close(abort_ssl=abort_ssl) + + for t in self._resolve_host_tasks: + t.cancel() + waiters.append(t) + + return waiters + + async def close(self, *, abort_ssl: bool = False) -> None: + """ + Close all opened transports. + + :param abort_ssl: If True, SSL connections will be aborted immediately + without performing the shutdown handshake. If False (default), + the behavior is determined by ssl_shutdown_timeout: + - If ssl_shutdown_timeout=0: connections are aborted + - If ssl_shutdown_timeout>0: graceful shutdown is performed + """ + if self._resolver_owner: + await self._resolver.close() + # Use abort_ssl param if explicitly set, otherwise use ssl_shutdown_timeout default + await super().close(abort_ssl=abort_ssl or self._ssl_shutdown_timeout == 0) + + @property + def family(self) -> int: + """Socket family like AF_INET.""" + return self._family + + @property + def use_dns_cache(self) -> bool: + """True if local DNS caching is enabled.""" + return self._use_dns_cache + + def clear_dns_cache( + self, host: Optional[str] = None, port: Optional[int] = None + ) -> None: + """Remove specified host/port or clear all dns local cache.""" + if host is not None and port is not None: + self._cached_hosts.remove((host, port)) + elif host is not None or port is not None: + raise ValueError("either both host and port or none of them are allowed") + else: + self._cached_hosts.clear() + + async def _resolve_host( + self, host: str, port: int, traces: Optional[Sequence["Trace"]] = None + ) -> List[ResolveResult]: + """Resolve host and return list of addresses.""" + if is_ip_address(host): + return [ + { + "hostname": host, + "host": host, + "port": port, + "family": self._family, + "proto": 0, + "flags": 0, + } + ] + + if not self._use_dns_cache: + + if traces: + for trace in traces: + await trace.send_dns_resolvehost_start(host) + + res = await self._resolver.resolve(host, port, family=self._family) + + if traces: + for trace in traces: + await trace.send_dns_resolvehost_end(host) + + return res + + key = (host, port) + if key in self._cached_hosts and not self._cached_hosts.expired(key): + # get result early, before any await (#4014) + result = self._cached_hosts.next_addrs(key) + + if traces: + for trace in traces: + await trace.send_dns_cache_hit(host) + return result + + futures: Set["asyncio.Future[None]"] + # + # If multiple connectors are resolving the same host, we wait + # for the first one to resolve and then use the result for all of them. + # We use a throttle to ensure that we only resolve the host once + # and then use the result for all the waiters. + # + if key in self._throttle_dns_futures: + # get futures early, before any await (#4014) + futures = self._throttle_dns_futures[key] + future: asyncio.Future[None] = self._loop.create_future() + futures.add(future) + if traces: + for trace in traces: + await trace.send_dns_cache_hit(host) + try: + await future + finally: + futures.discard(future) + return self._cached_hosts.next_addrs(key) + + # update dict early, before any await (#4014) + self._throttle_dns_futures[key] = futures = set() + # In this case we need to create a task to ensure that we can shield + # the task from cancellation as cancelling this lookup should not cancel + # the underlying lookup or else the cancel event will get broadcast to + # all the waiters across all connections. + # + coro = self._resolve_host_with_throttle(key, host, port, futures, traces) + loop = asyncio.get_running_loop() + if sys.version_info >= (3, 12): + # Optimization for Python 3.12, try to send immediately + resolved_host_task = asyncio.Task(coro, loop=loop, eager_start=True) + else: + resolved_host_task = loop.create_task(coro) + + if not resolved_host_task.done(): + self._resolve_host_tasks.add(resolved_host_task) + resolved_host_task.add_done_callback(self._resolve_host_tasks.discard) + + try: + return await asyncio.shield(resolved_host_task) + except asyncio.CancelledError: + + def drop_exception(fut: "asyncio.Future[List[ResolveResult]]") -> None: + with suppress(Exception, asyncio.CancelledError): + fut.result() + + resolved_host_task.add_done_callback(drop_exception) + raise + + async def _resolve_host_with_throttle( + self, + key: Tuple[str, int], + host: str, + port: int, + futures: Set["asyncio.Future[None]"], + traces: Optional[Sequence["Trace"]], + ) -> List[ResolveResult]: + """Resolve host and set result for all waiters. + + This method must be run in a task and shielded from cancellation + to avoid cancelling the underlying lookup. + """ + try: + if traces: + for trace in traces: + await trace.send_dns_cache_miss(host) + + for trace in traces: + await trace.send_dns_resolvehost_start(host) + + addrs = await self._resolver.resolve(host, port, family=self._family) + if traces: + for trace in traces: + await trace.send_dns_resolvehost_end(host) + + self._cached_hosts.add(key, addrs) + for fut in futures: + set_result(fut, None) + except BaseException as e: + # any DNS exception is set for the waiters to raise the same exception. + # This coro is always run in task that is shielded from cancellation so + # we should never be propagating cancellation here. + for fut in futures: + set_exception(fut, e) + raise + finally: + self._throttle_dns_futures.pop(key) + + return self._cached_hosts.next_addrs(key) + + async def _create_connection( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> ResponseHandler: + """Create connection. + + Has same keyword arguments as BaseEventLoop.create_connection. + """ + if req.proxy: + _, proto = await self._create_proxy_connection(req, traces, timeout) + else: + _, proto = await self._create_direct_connection(req, traces, timeout) + + return proto + + def _get_ssl_context(self, req: ClientRequest) -> Optional[SSLContext]: + """Logic to get the correct SSL context + + 0. if req.ssl is false, return None + + 1. if ssl_context is specified in req, use it + 2. if _ssl_context is specified in self, use it + 3. otherwise: + 1. if verify_ssl is not specified in req, use self.ssl_context + (will generate a default context according to self.verify_ssl) + 2. if verify_ssl is True in req, generate a default SSL context + 3. if verify_ssl is False in req, generate a SSL context that + won't verify + """ + if not req.is_ssl(): + return None + + if ssl is None: # pragma: no cover + raise RuntimeError("SSL is not supported.") + sslcontext = req.ssl + if isinstance(sslcontext, ssl.SSLContext): + return sslcontext + if sslcontext is not True: + # not verified or fingerprinted + return _SSL_CONTEXT_UNVERIFIED + sslcontext = self._ssl + if isinstance(sslcontext, ssl.SSLContext): + return sslcontext + if sslcontext is not True: + # not verified or fingerprinted + return _SSL_CONTEXT_UNVERIFIED + return _SSL_CONTEXT_VERIFIED + + def _get_fingerprint(self, req: ClientRequest) -> Optional["Fingerprint"]: + ret = req.ssl + if isinstance(ret, Fingerprint): + return ret + ret = self._ssl + if isinstance(ret, Fingerprint): + return ret + return None + + async def _wrap_create_connection( + self, + *args: Any, + addr_infos: List[AddrInfoType], + req: ClientRequest, + timeout: "ClientTimeout", + client_error: Type[Exception] = ClientConnectorError, + **kwargs: Any, + ) -> Tuple[asyncio.Transport, ResponseHandler]: + try: + async with ceil_timeout( + timeout.sock_connect, ceil_threshold=timeout.ceil_threshold + ): + sock = await aiohappyeyeballs.start_connection( + addr_infos=addr_infos, + local_addr_infos=self._local_addr_infos, + happy_eyeballs_delay=self._happy_eyeballs_delay, + interleave=self._interleave, + loop=self._loop, + socket_factory=self._socket_factory, + ) + # Add ssl_shutdown_timeout for Python 3.11+ when SSL is used + if ( + kwargs.get("ssl") + and self._ssl_shutdown_timeout + and sys.version_info >= (3, 11) + ): + kwargs["ssl_shutdown_timeout"] = self._ssl_shutdown_timeout + return await self._loop.create_connection(*args, **kwargs, sock=sock) + except cert_errors as exc: + raise ClientConnectorCertificateError(req.connection_key, exc) from exc + except ssl_errors as exc: + raise ClientConnectorSSLError(req.connection_key, exc) from exc + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise client_error(req.connection_key, exc) from exc + + async def _wrap_existing_connection( + self, + *args: Any, + req: ClientRequest, + timeout: "ClientTimeout", + client_error: Type[Exception] = ClientConnectorError, + **kwargs: Any, + ) -> Tuple[asyncio.Transport, ResponseHandler]: + try: + async with ceil_timeout( + timeout.sock_connect, ceil_threshold=timeout.ceil_threshold + ): + return await self._loop.create_connection(*args, **kwargs) + except cert_errors as exc: + raise ClientConnectorCertificateError(req.connection_key, exc) from exc + except ssl_errors as exc: + raise ClientConnectorSSLError(req.connection_key, exc) from exc + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise client_error(req.connection_key, exc) from exc + + def _fail_on_no_start_tls(self, req: "ClientRequest") -> None: + """Raise a :py:exc:`RuntimeError` on missing ``start_tls()``. + + It is necessary for TLS-in-TLS so that it is possible to + send HTTPS queries through HTTPS proxies. + + This doesn't affect regular HTTP requests, though. + """ + if not req.is_ssl(): + return + + proxy_url = req.proxy + assert proxy_url is not None + if proxy_url.scheme != "https": + return + + self._check_loop_for_start_tls() + + def _check_loop_for_start_tls(self) -> None: + try: + self._loop.start_tls + except AttributeError as attr_exc: + raise RuntimeError( + "An HTTPS request is being sent through an HTTPS proxy. " + "This needs support for TLS in TLS but it is not implemented " + "in your runtime for the stdlib asyncio.\n\n" + "Please upgrade to Python 3.11 or higher. For more details, " + "please see:\n" + "* https://bugs.python.org/issue37179\n" + "* https://github.com/python/cpython/pull/28073\n" + "* https://docs.aiohttp.org/en/stable/" + "client_advanced.html#proxy-support\n" + "* https://github.com/aio-libs/aiohttp/discussions/6044\n", + ) from attr_exc + + def _loop_supports_start_tls(self) -> bool: + try: + self._check_loop_for_start_tls() + except RuntimeError: + return False + else: + return True + + def _warn_about_tls_in_tls( + self, + underlying_transport: asyncio.Transport, + req: ClientRequest, + ) -> None: + """Issue a warning if the requested URL has HTTPS scheme.""" + if req.request_info.url.scheme != "https": + return + + # Check if uvloop is being used, which supports TLS in TLS, + # otherwise assume that asyncio's native transport is being used. + if type(underlying_transport).__module__.startswith("uvloop"): + return + + # Support in asyncio was added in Python 3.11 (bpo-44011) + asyncio_supports_tls_in_tls = sys.version_info >= (3, 11) or getattr( + underlying_transport, + "_start_tls_compatible", + False, + ) + + if asyncio_supports_tls_in_tls: + return + + warnings.warn( + "An HTTPS request is being sent through an HTTPS proxy. " + "This support for TLS in TLS is known to be disabled " + "in the stdlib asyncio (Python <3.11). This is why you'll probably see " + "an error in the log below.\n\n" + "It is possible to enable it via monkeypatching. " + "For more details, see:\n" + "* https://bugs.python.org/issue37179\n" + "* https://github.com/python/cpython/pull/28073\n\n" + "You can temporarily patch this as follows:\n" + "* https://docs.aiohttp.org/en/stable/client_advanced.html#proxy-support\n" + "* https://github.com/aio-libs/aiohttp/discussions/6044\n", + RuntimeWarning, + source=self, + # Why `4`? At least 3 of the calls in the stack originate + # from the methods in this class. + stacklevel=3, + ) + + async def _start_tls_connection( + self, + underlying_transport: asyncio.Transport, + req: ClientRequest, + timeout: "ClientTimeout", + client_error: Type[Exception] = ClientConnectorError, + ) -> Tuple[asyncio.BaseTransport, ResponseHandler]: + """Wrap the raw TCP transport with TLS.""" + tls_proto = self._factory() # Create a brand new proto for TLS + sslcontext = self._get_ssl_context(req) + if TYPE_CHECKING: + # _start_tls_connection is unreachable in the current code path + # if sslcontext is None. + assert sslcontext is not None + + try: + async with ceil_timeout( + timeout.sock_connect, ceil_threshold=timeout.ceil_threshold + ): + try: + # ssl_shutdown_timeout is only available in Python 3.11+ + if sys.version_info >= (3, 11) and self._ssl_shutdown_timeout: + tls_transport = await self._loop.start_tls( + underlying_transport, + tls_proto, + sslcontext, + server_hostname=req.server_hostname or req.host, + ssl_handshake_timeout=timeout.total, + ssl_shutdown_timeout=self._ssl_shutdown_timeout, + ) + else: + tls_transport = await self._loop.start_tls( + underlying_transport, + tls_proto, + sslcontext, + server_hostname=req.server_hostname or req.host, + ssl_handshake_timeout=timeout.total, + ) + except BaseException: + # We need to close the underlying transport since + # `start_tls()` probably failed before it had a + # chance to do this: + if self._ssl_shutdown_timeout == 0: + underlying_transport.abort() + else: + underlying_transport.close() + raise + if isinstance(tls_transport, asyncio.Transport): + fingerprint = self._get_fingerprint(req) + if fingerprint: + try: + fingerprint.check(tls_transport) + except ServerFingerprintMismatch: + tls_transport.close() + if not self._cleanup_closed_disabled: + self._cleanup_closed_transports.append(tls_transport) + raise + except cert_errors as exc: + raise ClientConnectorCertificateError(req.connection_key, exc) from exc + except ssl_errors as exc: + raise ClientConnectorSSLError(req.connection_key, exc) from exc + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise client_error(req.connection_key, exc) from exc + except TypeError as type_err: + # Example cause looks like this: + # TypeError: transport is not supported by start_tls() + + raise ClientConnectionError( + "Cannot initialize a TLS-in-TLS connection to host " + f"{req.host!s}:{req.port:d} through an underlying connection " + f"to an HTTPS proxy {req.proxy!s} ssl:{req.ssl or 'default'} " + f"[{type_err!s}]" + ) from type_err + else: + if tls_transport is None: + msg = "Failed to start TLS (possibly caused by closing transport)" + raise client_error(req.connection_key, OSError(msg)) + tls_proto.connection_made( + tls_transport + ) # Kick the state machine of the new TLS protocol + + return tls_transport, tls_proto + + def _convert_hosts_to_addr_infos( + self, hosts: List[ResolveResult] + ) -> List[AddrInfoType]: + """Converts the list of hosts to a list of addr_infos. + + The list of hosts is the result of a DNS lookup. The list of + addr_infos is the result of a call to `socket.getaddrinfo()`. + """ + addr_infos: List[AddrInfoType] = [] + for hinfo in hosts: + host = hinfo["host"] + is_ipv6 = ":" in host + family = socket.AF_INET6 if is_ipv6 else socket.AF_INET + if self._family and self._family != family: + continue + addr = (host, hinfo["port"], 0, 0) if is_ipv6 else (host, hinfo["port"]) + addr_infos.append( + (family, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", addr) + ) + return addr_infos + + async def _create_direct_connection( + self, + req: ClientRequest, + traces: List["Trace"], + timeout: "ClientTimeout", + *, + client_error: Type[Exception] = ClientConnectorError, + ) -> Tuple[asyncio.Transport, ResponseHandler]: + sslcontext = self._get_ssl_context(req) + fingerprint = self._get_fingerprint(req) + + host = req.url.raw_host + assert host is not None + # Replace multiple trailing dots with a single one. + # A trailing dot is only present for fully-qualified domain names. + # See https://github.com/aio-libs/aiohttp/pull/7364. + if host.endswith(".."): + host = host.rstrip(".") + "." + port = req.port + assert port is not None + try: + # Cancelling this lookup should not cancel the underlying lookup + # or else the cancel event will get broadcast to all the waiters + # across all connections. + hosts = await self._resolve_host(host, port, traces=traces) + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + # in case of proxy it is not ClientProxyConnectionError + # it is problem of resolving proxy ip itself + raise ClientConnectorDNSError(req.connection_key, exc) from exc + + last_exc: Optional[Exception] = None + addr_infos = self._convert_hosts_to_addr_infos(hosts) + while addr_infos: + # Strip trailing dots, certificates contain FQDN without dots. + # See https://github.com/aio-libs/aiohttp/issues/3636 + server_hostname = ( + (req.server_hostname or host).rstrip(".") if sslcontext else None + ) + + try: + transp, proto = await self._wrap_create_connection( + self._factory, + timeout=timeout, + ssl=sslcontext, + addr_infos=addr_infos, + server_hostname=server_hostname, + req=req, + client_error=client_error, + ) + except (ClientConnectorError, asyncio.TimeoutError) as exc: + last_exc = exc + aiohappyeyeballs.pop_addr_infos_interleave(addr_infos, self._interleave) + continue + + if req.is_ssl() and fingerprint: + try: + fingerprint.check(transp) + except ServerFingerprintMismatch as exc: + transp.close() + if not self._cleanup_closed_disabled: + self._cleanup_closed_transports.append(transp) + last_exc = exc + # Remove the bad peer from the list of addr_infos + sock: socket.socket = transp.get_extra_info("socket") + bad_peer = sock.getpeername() + aiohappyeyeballs.remove_addr_infos(addr_infos, bad_peer) + continue + + return transp, proto + else: + assert last_exc is not None + raise last_exc + + async def _create_proxy_connection( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> Tuple[asyncio.BaseTransport, ResponseHandler]: + self._fail_on_no_start_tls(req) + runtime_has_start_tls = self._loop_supports_start_tls() + proxy_req = self._update_proxy_auth_header_and_build_proxy_req(req) + + # create connection to proxy server + transport, proto = await self._create_direct_connection( + proxy_req, [], timeout, client_error=ClientProxyConnectionError + ) + + if req.is_ssl(): + if runtime_has_start_tls: + self._warn_about_tls_in_tls(transport, req) + + # For HTTPS requests over HTTP proxy + # we must notify proxy to tunnel connection + # so we send CONNECT command: + # CONNECT www.python.org:443 HTTP/1.1 + # Host: www.python.org + # + # next we must do TLS handshake and so on + # to do this we must wrap raw socket into secure one + # asyncio handles this perfectly + proxy_req.method = hdrs.METH_CONNECT + proxy_req.url = req.url + key = req.connection_key._replace( + proxy=None, proxy_auth=None, proxy_headers_hash=None + ) + conn = _ConnectTunnelConnection(self, key, proto, self._loop) + proxy_resp = await proxy_req.send(conn) + try: + protocol = conn._protocol + assert protocol is not None + + # read_until_eof=True will ensure the connection isn't closed + # once the response is received and processed allowing + # START_TLS to work on the connection below. + protocol.set_response_params( + read_until_eof=runtime_has_start_tls, + timeout_ceil_threshold=self._timeout_ceil_threshold, + ) + resp = await proxy_resp.start(conn) + except BaseException: + proxy_resp.close() + conn.close() + raise + else: + conn._protocol = None + try: + if resp.status != 200: + message = resp.reason + if message is None: + message = HTTPStatus(resp.status).phrase + raise ClientHttpProxyError( + proxy_resp.request_info, + resp.history, + status=resp.status, + message=message, + headers=resp.headers, + ) + if not runtime_has_start_tls: + rawsock = transport.get_extra_info("socket", default=None) + if rawsock is None: + raise RuntimeError( + "Transport does not expose socket instance" + ) + # Duplicate the socket, so now we can close proxy transport + rawsock = rawsock.dup() + except BaseException: + # It shouldn't be closed in `finally` because it's fed to + # `loop.start_tls()` and the docs say not to touch it after + # passing there. + transport.close() + raise + finally: + if not runtime_has_start_tls: + transport.close() + + if not runtime_has_start_tls: + # HTTP proxy with support for upgrade to HTTPS + sslcontext = self._get_ssl_context(req) + return await self._wrap_existing_connection( + self._factory, + timeout=timeout, + ssl=sslcontext, + sock=rawsock, + server_hostname=req.host, + req=req, + ) + + return await self._start_tls_connection( + # Access the old transport for the last time before it's + # closed and forgotten forever: + transport, + req=req, + timeout=timeout, + ) + finally: + proxy_resp.close() + + return transport, proto + + +class UnixConnector(BaseConnector): + """Unix socket connector. + + path - Unix socket path. + keepalive_timeout - (optional) Keep-alive timeout. + force_close - Set to True to force close and do reconnect + after each request (and between redirects). + limit - The total number of simultaneous connections. + limit_per_host - Number of simultaneous connections to one host. + loop - Optional event loop. + """ + + allowed_protocol_schema_set = HIGH_LEVEL_SCHEMA_SET | frozenset({"unix"}) + + def __init__( + self, + path: str, + force_close: bool = False, + keepalive_timeout: Union[object, float, None] = sentinel, + limit: int = 100, + limit_per_host: int = 0, + loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> None: + super().__init__( + force_close=force_close, + keepalive_timeout=keepalive_timeout, + limit=limit, + limit_per_host=limit_per_host, + loop=loop, + ) + self._path = path + + @property + def path(self) -> str: + """Path to unix socket.""" + return self._path + + async def _create_connection( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> ResponseHandler: + try: + async with ceil_timeout( + timeout.sock_connect, ceil_threshold=timeout.ceil_threshold + ): + _, proto = await self._loop.create_unix_connection( + self._factory, self._path + ) + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise UnixClientConnectorError(self.path, req.connection_key, exc) from exc + + return proto + + +class NamedPipeConnector(BaseConnector): + """Named pipe connector. + + Only supported by the proactor event loop. + See also: https://docs.python.org/3/library/asyncio-eventloop.html + + path - Windows named pipe path. + keepalive_timeout - (optional) Keep-alive timeout. + force_close - Set to True to force close and do reconnect + after each request (and between redirects). + limit - The total number of simultaneous connections. + limit_per_host - Number of simultaneous connections to one host. + loop - Optional event loop. + """ + + allowed_protocol_schema_set = HIGH_LEVEL_SCHEMA_SET | frozenset({"npipe"}) + + def __init__( + self, + path: str, + force_close: bool = False, + keepalive_timeout: Union[object, float, None] = sentinel, + limit: int = 100, + limit_per_host: int = 0, + loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> None: + super().__init__( + force_close=force_close, + keepalive_timeout=keepalive_timeout, + limit=limit, + limit_per_host=limit_per_host, + loop=loop, + ) + if not isinstance( + self._loop, + asyncio.ProactorEventLoop, # type: ignore[attr-defined] + ): + raise RuntimeError( + "Named Pipes only available in proactor loop under windows" + ) + self._path = path + + @property + def path(self) -> str: + """Path to the named pipe.""" + return self._path + + async def _create_connection( + self, req: ClientRequest, traces: List["Trace"], timeout: "ClientTimeout" + ) -> ResponseHandler: + try: + async with ceil_timeout( + timeout.sock_connect, ceil_threshold=timeout.ceil_threshold + ): + _, proto = await self._loop.create_pipe_connection( # type: ignore[attr-defined] + self._factory, self._path + ) + # the drain is required so that the connection_made is called + # and transport is set otherwise it is not set before the + # `assert conn.transport is not None` + # in client.py's _request method + await asyncio.sleep(0) + # other option is to manually set transport like + # `proto.transport = trans` + except OSError as exc: + if exc.errno is None and isinstance(exc, asyncio.TimeoutError): + raise + raise ClientConnectorError(req.connection_key, exc) from exc + + return cast(ResponseHandler, proto) diff --git a/lib/python3.12/site-packages/aiohttp/cookiejar.py b/lib/python3.12/site-packages/aiohttp/cookiejar.py new file mode 100644 index 0000000000000000000000000000000000000000..193648d4309fc9252332c0c1241ab6aafbd2fe05 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/cookiejar.py @@ -0,0 +1,522 @@ +import asyncio +import calendar +import contextlib +import datetime +import heapq +import itertools +import os # noqa +import pathlib +import pickle +import re +import time +import warnings +from collections import defaultdict +from collections.abc import Mapping +from http.cookies import BaseCookie, Morsel, SimpleCookie +from typing import ( + DefaultDict, + Dict, + Iterable, + Iterator, + List, + Optional, + Set, + Tuple, + Union, +) + +from yarl import URL + +from ._cookie_helpers import preserve_morsel_with_coded_value +from .abc import AbstractCookieJar, ClearCookiePredicate +from .helpers import is_ip_address +from .typedefs import LooseCookies, PathLike, StrOrURL + +__all__ = ("CookieJar", "DummyCookieJar") + + +CookieItem = Union[str, "Morsel[str]"] + +# We cache these string methods here as their use is in performance critical code. +_FORMAT_PATH = "{}/{}".format +_FORMAT_DOMAIN_REVERSED = "{1}.{0}".format + +# The minimum number of scheduled cookie expirations before we start cleaning up +# the expiration heap. This is a performance optimization to avoid cleaning up the +# heap too often when there are only a few scheduled expirations. +_MIN_SCHEDULED_COOKIE_EXPIRATION = 100 +_SIMPLE_COOKIE = SimpleCookie() + + +class CookieJar(AbstractCookieJar): + """Implements cookie storage adhering to RFC 6265.""" + + DATE_TOKENS_RE = re.compile( + r"[\x09\x20-\x2F\x3B-\x40\x5B-\x60\x7B-\x7E]*" + r"(?P[\x00-\x08\x0A-\x1F\d:a-zA-Z\x7F-\xFF]+)" + ) + + DATE_HMS_TIME_RE = re.compile(r"(\d{1,2}):(\d{1,2}):(\d{1,2})") + + DATE_DAY_OF_MONTH_RE = re.compile(r"(\d{1,2})") + + DATE_MONTH_RE = re.compile( + "(jan)|(feb)|(mar)|(apr)|(may)|(jun)|(jul)|(aug)|(sep)|(oct)|(nov)|(dec)", + re.I, + ) + + DATE_YEAR_RE = re.compile(r"(\d{2,4})") + + # calendar.timegm() fails for timestamps after datetime.datetime.max + # Minus one as a loss of precision occurs when timestamp() is called. + MAX_TIME = ( + int(datetime.datetime.max.replace(tzinfo=datetime.timezone.utc).timestamp()) - 1 + ) + try: + calendar.timegm(time.gmtime(MAX_TIME)) + except (OSError, ValueError): + # Hit the maximum representable time on Windows + # https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/localtime-localtime32-localtime64 + # Throws ValueError on PyPy 3.9, OSError elsewhere + MAX_TIME = calendar.timegm((3000, 12, 31, 23, 59, 59, -1, -1, -1)) + except OverflowError: + # #4515: datetime.max may not be representable on 32-bit platforms + MAX_TIME = 2**31 - 1 + # Avoid minuses in the future, 3x faster + SUB_MAX_TIME = MAX_TIME - 1 + + def __init__( + self, + *, + unsafe: bool = False, + quote_cookie: bool = True, + treat_as_secure_origin: Union[StrOrURL, List[StrOrURL], None] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> None: + super().__init__(loop=loop) + self._cookies: DefaultDict[Tuple[str, str], SimpleCookie] = defaultdict( + SimpleCookie + ) + self._morsel_cache: DefaultDict[Tuple[str, str], Dict[str, Morsel[str]]] = ( + defaultdict(dict) + ) + self._host_only_cookies: Set[Tuple[str, str]] = set() + self._unsafe = unsafe + self._quote_cookie = quote_cookie + if treat_as_secure_origin is None: + treat_as_secure_origin = [] + elif isinstance(treat_as_secure_origin, URL): + treat_as_secure_origin = [treat_as_secure_origin.origin()] + elif isinstance(treat_as_secure_origin, str): + treat_as_secure_origin = [URL(treat_as_secure_origin).origin()] + else: + treat_as_secure_origin = [ + URL(url).origin() if isinstance(url, str) else url.origin() + for url in treat_as_secure_origin + ] + self._treat_as_secure_origin = treat_as_secure_origin + self._expire_heap: List[Tuple[float, Tuple[str, str, str]]] = [] + self._expirations: Dict[Tuple[str, str, str], float] = {} + + @property + def quote_cookie(self) -> bool: + return self._quote_cookie + + def save(self, file_path: PathLike) -> None: + file_path = pathlib.Path(file_path) + with file_path.open(mode="wb") as f: + pickle.dump(self._cookies, f, pickle.HIGHEST_PROTOCOL) + + def load(self, file_path: PathLike) -> None: + file_path = pathlib.Path(file_path) + with file_path.open(mode="rb") as f: + self._cookies = pickle.load(f) + + def clear(self, predicate: Optional[ClearCookiePredicate] = None) -> None: + if predicate is None: + self._expire_heap.clear() + self._cookies.clear() + self._morsel_cache.clear() + self._host_only_cookies.clear() + self._expirations.clear() + return + + now = time.time() + to_del = [ + key + for (domain, path), cookie in self._cookies.items() + for name, morsel in cookie.items() + if ( + (key := (domain, path, name)) in self._expirations + and self._expirations[key] <= now + ) + or predicate(morsel) + ] + if to_del: + self._delete_cookies(to_del) + + def clear_domain(self, domain: str) -> None: + self.clear(lambda x: self._is_domain_match(domain, x["domain"])) + + def __iter__(self) -> "Iterator[Morsel[str]]": + self._do_expiration() + for val in self._cookies.values(): + yield from val.values() + + def __len__(self) -> int: + """Return number of cookies. + + This function does not iterate self to avoid unnecessary expiration + checks. + """ + return sum(len(cookie.values()) for cookie in self._cookies.values()) + + def _do_expiration(self) -> None: + """Remove expired cookies.""" + if not (expire_heap_len := len(self._expire_heap)): + return + + # If the expiration heap grows larger than the number expirations + # times two, we clean it up to avoid keeping expired entries in + # the heap and consuming memory. We guard this with a minimum + # threshold to avoid cleaning up the heap too often when there are + # only a few scheduled expirations. + if ( + expire_heap_len > _MIN_SCHEDULED_COOKIE_EXPIRATION + and expire_heap_len > len(self._expirations) * 2 + ): + # Remove any expired entries from the expiration heap + # that do not match the expiration time in the expirations + # as it means the cookie has been re-added to the heap + # with a different expiration time. + self._expire_heap = [ + entry + for entry in self._expire_heap + if self._expirations.get(entry[1]) == entry[0] + ] + heapq.heapify(self._expire_heap) + + now = time.time() + to_del: List[Tuple[str, str, str]] = [] + # Find any expired cookies and add them to the to-delete list + while self._expire_heap: + when, cookie_key = self._expire_heap[0] + if when > now: + break + heapq.heappop(self._expire_heap) + # Check if the cookie hasn't been re-added to the heap + # with a different expiration time as it will be removed + # later when it reaches the top of the heap and its + # expiration time is met. + if self._expirations.get(cookie_key) == when: + to_del.append(cookie_key) + + if to_del: + self._delete_cookies(to_del) + + def _delete_cookies(self, to_del: List[Tuple[str, str, str]]) -> None: + for domain, path, name in to_del: + self._host_only_cookies.discard((domain, name)) + self._cookies[(domain, path)].pop(name, None) + self._morsel_cache[(domain, path)].pop(name, None) + self._expirations.pop((domain, path, name), None) + + def _expire_cookie(self, when: float, domain: str, path: str, name: str) -> None: + cookie_key = (domain, path, name) + if self._expirations.get(cookie_key) == when: + # Avoid adding duplicates to the heap + return + heapq.heappush(self._expire_heap, (when, cookie_key)) + self._expirations[cookie_key] = when + + def update_cookies(self, cookies: LooseCookies, response_url: URL = URL()) -> None: + """Update cookies.""" + hostname = response_url.raw_host + + if not self._unsafe and is_ip_address(hostname): + # Don't accept cookies from IPs + return + + if isinstance(cookies, Mapping): + cookies = cookies.items() + + for name, cookie in cookies: + if not isinstance(cookie, Morsel): + tmp = SimpleCookie() + tmp[name] = cookie # type: ignore[assignment] + cookie = tmp[name] + + domain = cookie["domain"] + + # ignore domains with trailing dots + if domain and domain[-1] == ".": + domain = "" + del cookie["domain"] + + if not domain and hostname is not None: + # Set the cookie's domain to the response hostname + # and set its host-only-flag + self._host_only_cookies.add((hostname, name)) + domain = cookie["domain"] = hostname + + if domain and domain[0] == ".": + # Remove leading dot + domain = domain[1:] + cookie["domain"] = domain + + if hostname and not self._is_domain_match(domain, hostname): + # Setting cookies for different domains is not allowed + continue + + path = cookie["path"] + if not path or path[0] != "/": + # Set the cookie's path to the response path + path = response_url.path + if not path.startswith("/"): + path = "/" + else: + # Cut everything from the last slash to the end + path = "/" + path[1 : path.rfind("/")] + cookie["path"] = path + path = path.rstrip("/") + + if max_age := cookie["max-age"]: + try: + delta_seconds = int(max_age) + max_age_expiration = min(time.time() + delta_seconds, self.MAX_TIME) + self._expire_cookie(max_age_expiration, domain, path, name) + except ValueError: + cookie["max-age"] = "" + + elif expires := cookie["expires"]: + if expire_time := self._parse_date(expires): + self._expire_cookie(expire_time, domain, path, name) + else: + cookie["expires"] = "" + + key = (domain, path) + if self._cookies[key].get(name) != cookie: + # Don't blow away the cache if the same + # cookie gets set again + self._cookies[key][name] = cookie + self._morsel_cache[key].pop(name, None) + + self._do_expiration() + + def filter_cookies(self, request_url: URL = URL()) -> "BaseCookie[str]": + """Returns this jar's cookies filtered by their attributes.""" + # We always use BaseCookie now since all + # cookies set on on filtered are fully constructed + # Morsels, not just names and values. + filtered: BaseCookie[str] = BaseCookie() + if not self._cookies: + # Skip do_expiration() if there are no cookies. + return filtered + self._do_expiration() + if not self._cookies: + # Skip rest of function if no non-expired cookies. + return filtered + if type(request_url) is not URL: + warnings.warn( + "filter_cookies expects yarl.URL instances only," + f"and will stop working in 4.x, got {type(request_url)}", + DeprecationWarning, + stacklevel=2, + ) + request_url = URL(request_url) + hostname = request_url.raw_host or "" + + is_not_secure = request_url.scheme not in ("https", "wss") + if is_not_secure and self._treat_as_secure_origin: + request_origin = URL() + with contextlib.suppress(ValueError): + request_origin = request_url.origin() + is_not_secure = request_origin not in self._treat_as_secure_origin + + # Send shared cookie + key = ("", "") + for c in self._cookies[key].values(): + # Check cache first + if c.key in self._morsel_cache[key]: + filtered[c.key] = self._morsel_cache[key][c.key] + continue + + # Build and cache the morsel + mrsl_val = self._build_morsel(c) + self._morsel_cache[key][c.key] = mrsl_val + filtered[c.key] = mrsl_val + + if is_ip_address(hostname): + if not self._unsafe: + return filtered + domains: Iterable[str] = (hostname,) + else: + # Get all the subdomains that might match a cookie (e.g. "foo.bar.com", "bar.com", "com") + domains = itertools.accumulate( + reversed(hostname.split(".")), _FORMAT_DOMAIN_REVERSED + ) + + # Get all the path prefixes that might match a cookie (e.g. "", "/foo", "/foo/bar") + paths = itertools.accumulate(request_url.path.split("/"), _FORMAT_PATH) + # Create every combination of (domain, path) pairs. + pairs = itertools.product(domains, paths) + + path_len = len(request_url.path) + # Point 2: https://www.rfc-editor.org/rfc/rfc6265.html#section-5.4 + for p in pairs: + if p not in self._cookies: + continue + for name, cookie in self._cookies[p].items(): + domain = cookie["domain"] + + if (domain, name) in self._host_only_cookies and domain != hostname: + continue + + # Skip edge case when the cookie has a trailing slash but request doesn't. + if len(cookie["path"]) > path_len: + continue + + if is_not_secure and cookie["secure"]: + continue + + # We already built the Morsel so reuse it here + if name in self._morsel_cache[p]: + filtered[name] = self._morsel_cache[p][name] + continue + + # Build and cache the morsel + mrsl_val = self._build_morsel(cookie) + self._morsel_cache[p][name] = mrsl_val + filtered[name] = mrsl_val + + return filtered + + def _build_morsel(self, cookie: Morsel[str]) -> Morsel[str]: + """Build a morsel for sending, respecting quote_cookie setting.""" + if self._quote_cookie and cookie.coded_value and cookie.coded_value[0] == '"': + return preserve_morsel_with_coded_value(cookie) + morsel: Morsel[str] = Morsel() + if self._quote_cookie: + value, coded_value = _SIMPLE_COOKIE.value_encode(cookie.value) + else: + coded_value = value = cookie.value + # We use __setstate__ instead of the public set() API because it allows us to + # bypass validation and set already validated state. This is more stable than + # setting protected attributes directly and unlikely to change since it would + # break pickling. + morsel.__setstate__({"key": cookie.key, "value": value, "coded_value": coded_value}) # type: ignore[attr-defined] + return morsel + + @staticmethod + def _is_domain_match(domain: str, hostname: str) -> bool: + """Implements domain matching adhering to RFC 6265.""" + if hostname == domain: + return True + + if not hostname.endswith(domain): + return False + + non_matching = hostname[: -len(domain)] + + if not non_matching.endswith("."): + return False + + return not is_ip_address(hostname) + + @classmethod + def _parse_date(cls, date_str: str) -> Optional[int]: + """Implements date string parsing adhering to RFC 6265.""" + if not date_str: + return None + + found_time = False + found_day = False + found_month = False + found_year = False + + hour = minute = second = 0 + day = 0 + month = 0 + year = 0 + + for token_match in cls.DATE_TOKENS_RE.finditer(date_str): + + token = token_match.group("token") + + if not found_time: + time_match = cls.DATE_HMS_TIME_RE.match(token) + if time_match: + found_time = True + hour, minute, second = (int(s) for s in time_match.groups()) + continue + + if not found_day: + day_match = cls.DATE_DAY_OF_MONTH_RE.match(token) + if day_match: + found_day = True + day = int(day_match.group()) + continue + + if not found_month: + month_match = cls.DATE_MONTH_RE.match(token) + if month_match: + found_month = True + assert month_match.lastindex is not None + month = month_match.lastindex + continue + + if not found_year: + year_match = cls.DATE_YEAR_RE.match(token) + if year_match: + found_year = True + year = int(year_match.group()) + + if 70 <= year <= 99: + year += 1900 + elif 0 <= year <= 69: + year += 2000 + + if False in (found_day, found_month, found_year, found_time): + return None + + if not 1 <= day <= 31: + return None + + if year < 1601 or hour > 23 or minute > 59 or second > 59: + return None + + return calendar.timegm((year, month, day, hour, minute, second, -1, -1, -1)) + + +class DummyCookieJar(AbstractCookieJar): + """Implements a dummy cookie storage. + + It can be used with the ClientSession when no cookie processing is needed. + + """ + + def __init__(self, *, loop: Optional[asyncio.AbstractEventLoop] = None) -> None: + super().__init__(loop=loop) + + def __iter__(self) -> "Iterator[Morsel[str]]": + while False: + yield None + + def __len__(self) -> int: + return 0 + + @property + def quote_cookie(self) -> bool: + return True + + def clear(self, predicate: Optional[ClearCookiePredicate] = None) -> None: + pass + + def clear_domain(self, domain: str) -> None: + pass + + def update_cookies(self, cookies: LooseCookies, response_url: URL = URL()) -> None: + pass + + def filter_cookies(self, request_url: URL) -> "BaseCookie[str]": + return SimpleCookie() diff --git a/lib/python3.12/site-packages/aiohttp/formdata.py b/lib/python3.12/site-packages/aiohttp/formdata.py new file mode 100644 index 0000000000000000000000000000000000000000..a5a4f603e190cdc66fe34e95a47983377f94666f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/formdata.py @@ -0,0 +1,179 @@ +import io +import warnings +from typing import Any, Iterable, List, Optional +from urllib.parse import urlencode + +from multidict import MultiDict, MultiDictProxy + +from . import hdrs, multipart, payload +from .helpers import guess_filename +from .payload import Payload + +__all__ = ("FormData",) + + +class FormData: + """Helper class for form body generation. + + Supports multipart/form-data and application/x-www-form-urlencoded. + """ + + def __init__( + self, + fields: Iterable[Any] = (), + quote_fields: bool = True, + charset: Optional[str] = None, + *, + default_to_multipart: bool = False, + ) -> None: + self._writer = multipart.MultipartWriter("form-data") + self._fields: List[Any] = [] + self._is_multipart = default_to_multipart + self._quote_fields = quote_fields + self._charset = charset + + if isinstance(fields, dict): + fields = list(fields.items()) + elif not isinstance(fields, (list, tuple)): + fields = (fields,) + self.add_fields(*fields) + + @property + def is_multipart(self) -> bool: + return self._is_multipart + + def add_field( + self, + name: str, + value: Any, + *, + content_type: Optional[str] = None, + filename: Optional[str] = None, + content_transfer_encoding: Optional[str] = None, + ) -> None: + + if isinstance(value, io.IOBase): + self._is_multipart = True + elif isinstance(value, (bytes, bytearray, memoryview)): + msg = ( + "In v4, passing bytes will no longer create a file field. " + "Please explicitly use the filename parameter or pass a BytesIO object." + ) + if filename is None and content_transfer_encoding is None: + warnings.warn(msg, DeprecationWarning) + filename = name + + type_options: MultiDict[str] = MultiDict({"name": name}) + if filename is not None and not isinstance(filename, str): + raise TypeError("filename must be an instance of str. Got: %s" % filename) + if filename is None and isinstance(value, io.IOBase): + filename = guess_filename(value, name) + if filename is not None: + type_options["filename"] = filename + self._is_multipart = True + + headers = {} + if content_type is not None: + if not isinstance(content_type, str): + raise TypeError( + "content_type must be an instance of str. Got: %s" % content_type + ) + headers[hdrs.CONTENT_TYPE] = content_type + self._is_multipart = True + if content_transfer_encoding is not None: + if not isinstance(content_transfer_encoding, str): + raise TypeError( + "content_transfer_encoding must be an instance" + " of str. Got: %s" % content_transfer_encoding + ) + msg = ( + "content_transfer_encoding is deprecated. " + "To maintain compatibility with v4 please pass a BytesPayload." + ) + warnings.warn(msg, DeprecationWarning) + self._is_multipart = True + + self._fields.append((type_options, headers, value)) + + def add_fields(self, *fields: Any) -> None: + to_add = list(fields) + + while to_add: + rec = to_add.pop(0) + + if isinstance(rec, io.IOBase): + k = guess_filename(rec, "unknown") + self.add_field(k, rec) # type: ignore[arg-type] + + elif isinstance(rec, (MultiDictProxy, MultiDict)): + to_add.extend(rec.items()) + + elif isinstance(rec, (list, tuple)) and len(rec) == 2: + k, fp = rec + self.add_field(k, fp) + + else: + raise TypeError( + "Only io.IOBase, multidict and (name, file) " + "pairs allowed, use .add_field() for passing " + "more complex parameters, got {!r}".format(rec) + ) + + def _gen_form_urlencoded(self) -> payload.BytesPayload: + # form data (x-www-form-urlencoded) + data = [] + for type_options, _, value in self._fields: + data.append((type_options["name"], value)) + + charset = self._charset if self._charset is not None else "utf-8" + + if charset == "utf-8": + content_type = "application/x-www-form-urlencoded" + else: + content_type = "application/x-www-form-urlencoded; charset=%s" % charset + + return payload.BytesPayload( + urlencode(data, doseq=True, encoding=charset).encode(), + content_type=content_type, + ) + + def _gen_form_data(self) -> multipart.MultipartWriter: + """Encode a list of fields using the multipart/form-data MIME format""" + for dispparams, headers, value in self._fields: + try: + if hdrs.CONTENT_TYPE in headers: + part = payload.get_payload( + value, + content_type=headers[hdrs.CONTENT_TYPE], + headers=headers, + encoding=self._charset, + ) + else: + part = payload.get_payload( + value, headers=headers, encoding=self._charset + ) + except Exception as exc: + raise TypeError( + "Can not serialize value type: %r\n " + "headers: %r\n value: %r" % (type(value), headers, value) + ) from exc + + if dispparams: + part.set_content_disposition( + "form-data", quote_fields=self._quote_fields, **dispparams + ) + # FIXME cgi.FieldStorage doesn't likes body parts with + # Content-Length which were sent via chunked transfer encoding + assert part.headers is not None + part.headers.popall(hdrs.CONTENT_LENGTH, None) + + self._writer.append_payload(part) + + self._fields.clear() + return self._writer + + def __call__(self) -> Payload: + if self._is_multipart: + return self._gen_form_data() + else: + return self._gen_form_urlencoded() diff --git a/lib/python3.12/site-packages/aiohttp/hdrs.py b/lib/python3.12/site-packages/aiohttp/hdrs.py new file mode 100644 index 0000000000000000000000000000000000000000..c8d6b35f33ae4be537d7776ce4085982618d1305 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/hdrs.py @@ -0,0 +1,121 @@ +"""HTTP Headers constants.""" + +# After changing the file content call ./tools/gen.py +# to regenerate the headers parser +import itertools +from typing import Final, Set + +from multidict import istr + +METH_ANY: Final[str] = "*" +METH_CONNECT: Final[str] = "CONNECT" +METH_HEAD: Final[str] = "HEAD" +METH_GET: Final[str] = "GET" +METH_DELETE: Final[str] = "DELETE" +METH_OPTIONS: Final[str] = "OPTIONS" +METH_PATCH: Final[str] = "PATCH" +METH_POST: Final[str] = "POST" +METH_PUT: Final[str] = "PUT" +METH_TRACE: Final[str] = "TRACE" + +METH_ALL: Final[Set[str]] = { + METH_CONNECT, + METH_HEAD, + METH_GET, + METH_DELETE, + METH_OPTIONS, + METH_PATCH, + METH_POST, + METH_PUT, + METH_TRACE, +} + +ACCEPT: Final[istr] = istr("Accept") +ACCEPT_CHARSET: Final[istr] = istr("Accept-Charset") +ACCEPT_ENCODING: Final[istr] = istr("Accept-Encoding") +ACCEPT_LANGUAGE: Final[istr] = istr("Accept-Language") +ACCEPT_RANGES: Final[istr] = istr("Accept-Ranges") +ACCESS_CONTROL_MAX_AGE: Final[istr] = istr("Access-Control-Max-Age") +ACCESS_CONTROL_ALLOW_CREDENTIALS: Final[istr] = istr("Access-Control-Allow-Credentials") +ACCESS_CONTROL_ALLOW_HEADERS: Final[istr] = istr("Access-Control-Allow-Headers") +ACCESS_CONTROL_ALLOW_METHODS: Final[istr] = istr("Access-Control-Allow-Methods") +ACCESS_CONTROL_ALLOW_ORIGIN: Final[istr] = istr("Access-Control-Allow-Origin") +ACCESS_CONTROL_EXPOSE_HEADERS: Final[istr] = istr("Access-Control-Expose-Headers") +ACCESS_CONTROL_REQUEST_HEADERS: Final[istr] = istr("Access-Control-Request-Headers") +ACCESS_CONTROL_REQUEST_METHOD: Final[istr] = istr("Access-Control-Request-Method") +AGE: Final[istr] = istr("Age") +ALLOW: Final[istr] = istr("Allow") +AUTHORIZATION: Final[istr] = istr("Authorization") +CACHE_CONTROL: Final[istr] = istr("Cache-Control") +CONNECTION: Final[istr] = istr("Connection") +CONTENT_DISPOSITION: Final[istr] = istr("Content-Disposition") +CONTENT_ENCODING: Final[istr] = istr("Content-Encoding") +CONTENT_LANGUAGE: Final[istr] = istr("Content-Language") +CONTENT_LENGTH: Final[istr] = istr("Content-Length") +CONTENT_LOCATION: Final[istr] = istr("Content-Location") +CONTENT_MD5: Final[istr] = istr("Content-MD5") +CONTENT_RANGE: Final[istr] = istr("Content-Range") +CONTENT_TRANSFER_ENCODING: Final[istr] = istr("Content-Transfer-Encoding") +CONTENT_TYPE: Final[istr] = istr("Content-Type") +COOKIE: Final[istr] = istr("Cookie") +DATE: Final[istr] = istr("Date") +DESTINATION: Final[istr] = istr("Destination") +DIGEST: Final[istr] = istr("Digest") +ETAG: Final[istr] = istr("Etag") +EXPECT: Final[istr] = istr("Expect") +EXPIRES: Final[istr] = istr("Expires") +FORWARDED: Final[istr] = istr("Forwarded") +FROM: Final[istr] = istr("From") +HOST: Final[istr] = istr("Host") +IF_MATCH: Final[istr] = istr("If-Match") +IF_MODIFIED_SINCE: Final[istr] = istr("If-Modified-Since") +IF_NONE_MATCH: Final[istr] = istr("If-None-Match") +IF_RANGE: Final[istr] = istr("If-Range") +IF_UNMODIFIED_SINCE: Final[istr] = istr("If-Unmodified-Since") +KEEP_ALIVE: Final[istr] = istr("Keep-Alive") +LAST_EVENT_ID: Final[istr] = istr("Last-Event-ID") +LAST_MODIFIED: Final[istr] = istr("Last-Modified") +LINK: Final[istr] = istr("Link") +LOCATION: Final[istr] = istr("Location") +MAX_FORWARDS: Final[istr] = istr("Max-Forwards") +ORIGIN: Final[istr] = istr("Origin") +PRAGMA: Final[istr] = istr("Pragma") +PROXY_AUTHENTICATE: Final[istr] = istr("Proxy-Authenticate") +PROXY_AUTHORIZATION: Final[istr] = istr("Proxy-Authorization") +RANGE: Final[istr] = istr("Range") +REFERER: Final[istr] = istr("Referer") +RETRY_AFTER: Final[istr] = istr("Retry-After") +SEC_WEBSOCKET_ACCEPT: Final[istr] = istr("Sec-WebSocket-Accept") +SEC_WEBSOCKET_VERSION: Final[istr] = istr("Sec-WebSocket-Version") +SEC_WEBSOCKET_PROTOCOL: Final[istr] = istr("Sec-WebSocket-Protocol") +SEC_WEBSOCKET_EXTENSIONS: Final[istr] = istr("Sec-WebSocket-Extensions") +SEC_WEBSOCKET_KEY: Final[istr] = istr("Sec-WebSocket-Key") +SEC_WEBSOCKET_KEY1: Final[istr] = istr("Sec-WebSocket-Key1") +SERVER: Final[istr] = istr("Server") +SET_COOKIE: Final[istr] = istr("Set-Cookie") +TE: Final[istr] = istr("TE") +TRAILER: Final[istr] = istr("Trailer") +TRANSFER_ENCODING: Final[istr] = istr("Transfer-Encoding") +UPGRADE: Final[istr] = istr("Upgrade") +URI: Final[istr] = istr("URI") +USER_AGENT: Final[istr] = istr("User-Agent") +VARY: Final[istr] = istr("Vary") +VIA: Final[istr] = istr("Via") +WANT_DIGEST: Final[istr] = istr("Want-Digest") +WARNING: Final[istr] = istr("Warning") +WWW_AUTHENTICATE: Final[istr] = istr("WWW-Authenticate") +X_FORWARDED_FOR: Final[istr] = istr("X-Forwarded-For") +X_FORWARDED_HOST: Final[istr] = istr("X-Forwarded-Host") +X_FORWARDED_PROTO: Final[istr] = istr("X-Forwarded-Proto") + +# These are the upper/lower case variants of the headers/methods +# Example: {'hOst', 'host', 'HoST', 'HOSt', 'hOsT', 'HosT', 'hoSt', ...} +METH_HEAD_ALL: Final = frozenset( + map("".join, itertools.product(*zip(METH_HEAD.upper(), METH_HEAD.lower()))) +) +METH_CONNECT_ALL: Final = frozenset( + map("".join, itertools.product(*zip(METH_CONNECT.upper(), METH_CONNECT.lower()))) +) +HOST_ALL: Final = frozenset( + map("".join, itertools.product(*zip(HOST.upper(), HOST.lower()))) +) diff --git a/lib/python3.12/site-packages/aiohttp/helpers.py b/lib/python3.12/site-packages/aiohttp/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..dfab9877d398a9c8807ad2614397b3e54316ed1f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/helpers.py @@ -0,0 +1,986 @@ +"""Various helper functions""" + +import asyncio +import base64 +import binascii +import contextlib +import datetime +import enum +import functools +import inspect +import netrc +import os +import platform +import re +import sys +import time +import weakref +from collections import namedtuple +from contextlib import suppress +from email.message import EmailMessage +from email.parser import HeaderParser +from email.policy import HTTP +from email.utils import parsedate +from math import ceil +from pathlib import Path +from types import MappingProxyType, TracebackType +from typing import ( + Any, + Callable, + ContextManager, + Dict, + Generator, + Generic, + Iterable, + Iterator, + List, + Mapping, + Optional, + Protocol, + Tuple, + Type, + TypeVar, + Union, + get_args, + overload, +) +from urllib.parse import quote +from urllib.request import getproxies, proxy_bypass + +import attr +from multidict import MultiDict, MultiDictProxy, MultiMapping +from propcache.api import under_cached_property as reify +from yarl import URL + +from . import hdrs +from .log import client_logger + +if sys.version_info >= (3, 11): + import asyncio as async_timeout +else: + import async_timeout + +__all__ = ("BasicAuth", "ChainMapProxy", "ETag", "reify") + +IS_MACOS = platform.system() == "Darwin" +IS_WINDOWS = platform.system() == "Windows" + +PY_310 = sys.version_info >= (3, 10) +PY_311 = sys.version_info >= (3, 11) + + +_T = TypeVar("_T") +_S = TypeVar("_S") + +_SENTINEL = enum.Enum("_SENTINEL", "sentinel") +sentinel = _SENTINEL.sentinel + +NO_EXTENSIONS = bool(os.environ.get("AIOHTTP_NO_EXTENSIONS")) + +# https://datatracker.ietf.org/doc/html/rfc9112#section-6.3-2.1 +EMPTY_BODY_STATUS_CODES = frozenset((204, 304, *range(100, 200))) +# https://datatracker.ietf.org/doc/html/rfc9112#section-6.3-2.1 +# https://datatracker.ietf.org/doc/html/rfc9112#section-6.3-2.2 +EMPTY_BODY_METHODS = hdrs.METH_HEAD_ALL + +DEBUG = sys.flags.dev_mode or ( + not sys.flags.ignore_environment and bool(os.environ.get("PYTHONASYNCIODEBUG")) +) + + +CHAR = {chr(i) for i in range(0, 128)} +CTL = {chr(i) for i in range(0, 32)} | { + chr(127), +} +SEPARATORS = { + "(", + ")", + "<", + ">", + "@", + ",", + ";", + ":", + "\\", + '"', + "/", + "[", + "]", + "?", + "=", + "{", + "}", + " ", + chr(9), +} +TOKEN = CHAR ^ CTL ^ SEPARATORS + + +class noop: + def __await__(self) -> Generator[None, None, None]: + yield + + +class BasicAuth(namedtuple("BasicAuth", ["login", "password", "encoding"])): + """Http basic authentication helper.""" + + def __new__( + cls, login: str, password: str = "", encoding: str = "latin1" + ) -> "BasicAuth": + if login is None: + raise ValueError("None is not allowed as login value") + + if password is None: + raise ValueError("None is not allowed as password value") + + if ":" in login: + raise ValueError('A ":" is not allowed in login (RFC 1945#section-11.1)') + + return super().__new__(cls, login, password, encoding) + + @classmethod + def decode(cls, auth_header: str, encoding: str = "latin1") -> "BasicAuth": + """Create a BasicAuth object from an Authorization HTTP header.""" + try: + auth_type, encoded_credentials = auth_header.split(" ", 1) + except ValueError: + raise ValueError("Could not parse authorization header.") + + if auth_type.lower() != "basic": + raise ValueError("Unknown authorization method %s" % auth_type) + + try: + decoded = base64.b64decode( + encoded_credentials.encode("ascii"), validate=True + ).decode(encoding) + except binascii.Error: + raise ValueError("Invalid base64 encoding.") + + try: + # RFC 2617 HTTP Authentication + # https://www.ietf.org/rfc/rfc2617.txt + # the colon must be present, but the username and password may be + # otherwise blank. + username, password = decoded.split(":", 1) + except ValueError: + raise ValueError("Invalid credentials.") + + return cls(username, password, encoding=encoding) + + @classmethod + def from_url(cls, url: URL, *, encoding: str = "latin1") -> Optional["BasicAuth"]: + """Create BasicAuth from url.""" + if not isinstance(url, URL): + raise TypeError("url should be yarl.URL instance") + # Check raw_user and raw_password first as yarl is likely + # to already have these values parsed from the netloc in the cache. + if url.raw_user is None and url.raw_password is None: + return None + return cls(url.user or "", url.password or "", encoding=encoding) + + def encode(self) -> str: + """Encode credentials.""" + creds = (f"{self.login}:{self.password}").encode(self.encoding) + return "Basic %s" % base64.b64encode(creds).decode(self.encoding) + + +def strip_auth_from_url(url: URL) -> Tuple[URL, Optional[BasicAuth]]: + """Remove user and password from URL if present and return BasicAuth object.""" + # Check raw_user and raw_password first as yarl is likely + # to already have these values parsed from the netloc in the cache. + if url.raw_user is None and url.raw_password is None: + return url, None + return url.with_user(None), BasicAuth(url.user or "", url.password or "") + + +def netrc_from_env() -> Optional[netrc.netrc]: + """Load netrc from file. + + Attempt to load it from the path specified by the env-var + NETRC or in the default location in the user's home directory. + + Returns None if it couldn't be found or fails to parse. + """ + netrc_env = os.environ.get("NETRC") + + if netrc_env is not None: + netrc_path = Path(netrc_env) + else: + try: + home_dir = Path.home() + except RuntimeError as e: # pragma: no cover + # if pathlib can't resolve home, it may raise a RuntimeError + client_logger.debug( + "Could not resolve home directory when " + "trying to look for .netrc file: %s", + e, + ) + return None + + netrc_path = home_dir / ("_netrc" if IS_WINDOWS else ".netrc") + + try: + return netrc.netrc(str(netrc_path)) + except netrc.NetrcParseError as e: + client_logger.warning("Could not parse .netrc file: %s", e) + except OSError as e: + netrc_exists = False + with contextlib.suppress(OSError): + netrc_exists = netrc_path.is_file() + # we couldn't read the file (doesn't exist, permissions, etc.) + if netrc_env or netrc_exists: + # only warn if the environment wanted us to load it, + # or it appears like the default file does actually exist + client_logger.warning("Could not read .netrc file: %s", e) + + return None + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class ProxyInfo: + proxy: URL + proxy_auth: Optional[BasicAuth] + + +def basicauth_from_netrc(netrc_obj: Optional[netrc.netrc], host: str) -> BasicAuth: + """ + Return :py:class:`~aiohttp.BasicAuth` credentials for ``host`` from ``netrc_obj``. + + :raises LookupError: if ``netrc_obj`` is :py:data:`None` or if no + entry is found for the ``host``. + """ + if netrc_obj is None: + raise LookupError("No .netrc file found") + auth_from_netrc = netrc_obj.authenticators(host) + + if auth_from_netrc is None: + raise LookupError(f"No entry for {host!s} found in the `.netrc` file.") + login, account, password = auth_from_netrc + + # TODO(PY311): username = login or account + # Up to python 3.10, account could be None if not specified, + # and login will be empty string if not specified. From 3.11, + # login and account will be empty string if not specified. + username = login if (login or account is None) else account + + # TODO(PY311): Remove this, as password will be empty string + # if not specified + if password is None: + password = "" + + return BasicAuth(username, password) + + +def proxies_from_env() -> Dict[str, ProxyInfo]: + proxy_urls = { + k: URL(v) + for k, v in getproxies().items() + if k in ("http", "https", "ws", "wss") + } + netrc_obj = netrc_from_env() + stripped = {k: strip_auth_from_url(v) for k, v in proxy_urls.items()} + ret = {} + for proto, val in stripped.items(): + proxy, auth = val + if proxy.scheme in ("https", "wss"): + client_logger.warning( + "%s proxies %s are not supported, ignoring", proxy.scheme.upper(), proxy + ) + continue + if netrc_obj and auth is None: + if proxy.host is not None: + try: + auth = basicauth_from_netrc(netrc_obj, proxy.host) + except LookupError: + auth = None + ret[proto] = ProxyInfo(proxy, auth) + return ret + + +def get_env_proxy_for_url(url: URL) -> Tuple[URL, Optional[BasicAuth]]: + """Get a permitted proxy for the given URL from the env.""" + if url.host is not None and proxy_bypass(url.host): + raise LookupError(f"Proxying is disallowed for `{url.host!r}`") + + proxies_in_env = proxies_from_env() + try: + proxy_info = proxies_in_env[url.scheme] + except KeyError: + raise LookupError(f"No proxies found for `{url!s}` in the env") + else: + return proxy_info.proxy, proxy_info.proxy_auth + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class MimeType: + type: str + subtype: str + suffix: str + parameters: "MultiDictProxy[str]" + + +@functools.lru_cache(maxsize=56) +def parse_mimetype(mimetype: str) -> MimeType: + """Parses a MIME type into its components. + + mimetype is a MIME type string. + + Returns a MimeType object. + + Example: + + >>> parse_mimetype('text/html; charset=utf-8') + MimeType(type='text', subtype='html', suffix='', + parameters={'charset': 'utf-8'}) + + """ + if not mimetype: + return MimeType( + type="", subtype="", suffix="", parameters=MultiDictProxy(MultiDict()) + ) + + parts = mimetype.split(";") + params: MultiDict[str] = MultiDict() + for item in parts[1:]: + if not item: + continue + key, _, value = item.partition("=") + params.add(key.lower().strip(), value.strip(' "')) + + fulltype = parts[0].strip().lower() + if fulltype == "*": + fulltype = "*/*" + + mtype, _, stype = fulltype.partition("/") + stype, _, suffix = stype.partition("+") + + return MimeType( + type=mtype, subtype=stype, suffix=suffix, parameters=MultiDictProxy(params) + ) + + +class EnsureOctetStream(EmailMessage): + def __init__(self) -> None: + super().__init__() + # https://www.rfc-editor.org/rfc/rfc9110#section-8.3-5 + self.set_default_type("application/octet-stream") + + def get_content_type(self) -> str: + """Re-implementation from Message + + Returns application/octet-stream in place of plain/text when + value is wrong. + + The way this class is used guarantees that content-type will + be present so simplify the checks wrt to the base implementation. + """ + value = self.get("content-type", "").lower() + + # Based on the implementation of _splitparam in the standard library + ctype, _, _ = value.partition(";") + ctype = ctype.strip() + if ctype.count("/") != 1: + return self.get_default_type() + return ctype + + +@functools.lru_cache(maxsize=56) +def parse_content_type(raw: str) -> Tuple[str, MappingProxyType[str, str]]: + """Parse Content-Type header. + + Returns a tuple of the parsed content type and a + MappingProxyType of parameters. The default returned value + is `application/octet-stream` + """ + msg = HeaderParser(EnsureOctetStream, policy=HTTP).parsestr(f"Content-Type: {raw}") + content_type = msg.get_content_type() + params = msg.get_params(()) + content_dict = dict(params[1:]) # First element is content type again + return content_type, MappingProxyType(content_dict) + + +def guess_filename(obj: Any, default: Optional[str] = None) -> Optional[str]: + name = getattr(obj, "name", None) + if name and isinstance(name, str) and name[0] != "<" and name[-1] != ">": + return Path(name).name + return default + + +not_qtext_re = re.compile(r"[^\041\043-\133\135-\176]") +QCONTENT = {chr(i) for i in range(0x20, 0x7F)} | {"\t"} + + +def quoted_string(content: str) -> str: + """Return 7-bit content as quoted-string. + + Format content into a quoted-string as defined in RFC5322 for + Internet Message Format. Notice that this is not the 8-bit HTTP + format, but the 7-bit email format. Content must be in usascii or + a ValueError is raised. + """ + if not (QCONTENT > set(content)): + raise ValueError(f"bad content for quoted-string {content!r}") + return not_qtext_re.sub(lambda x: "\\" + x.group(0), content) + + +def content_disposition_header( + disptype: str, quote_fields: bool = True, _charset: str = "utf-8", **params: str +) -> str: + """Sets ``Content-Disposition`` header for MIME. + + This is the MIME payload Content-Disposition header from RFC 2183 + and RFC 7579 section 4.2, not the HTTP Content-Disposition from + RFC 6266. + + disptype is a disposition type: inline, attachment, form-data. + Should be valid extension token (see RFC 2183) + + quote_fields performs value quoting to 7-bit MIME headers + according to RFC 7578. Set to quote_fields to False if recipient + can take 8-bit file names and field values. + + _charset specifies the charset to use when quote_fields is True. + + params is a dict with disposition params. + """ + if not disptype or not (TOKEN > set(disptype)): + raise ValueError(f"bad content disposition type {disptype!r}") + + value = disptype + if params: + lparams = [] + for key, val in params.items(): + if not key or not (TOKEN > set(key)): + raise ValueError(f"bad content disposition parameter {key!r}={val!r}") + if quote_fields: + if key.lower() == "filename": + qval = quote(val, "", encoding=_charset) + lparams.append((key, '"%s"' % qval)) + else: + try: + qval = quoted_string(val) + except ValueError: + qval = "".join( + (_charset, "''", quote(val, "", encoding=_charset)) + ) + lparams.append((key + "*", qval)) + else: + lparams.append((key, '"%s"' % qval)) + else: + qval = val.replace("\\", "\\\\").replace('"', '\\"') + lparams.append((key, '"%s"' % qval)) + sparams = "; ".join("=".join(pair) for pair in lparams) + value = "; ".join((value, sparams)) + return value + + +def is_ip_address(host: Optional[str]) -> bool: + """Check if host looks like an IP Address. + + This check is only meant as a heuristic to ensure that + a host is not a domain name. + """ + if not host: + return False + # For a host to be an ipv4 address, it must be all numeric. + # The host must contain a colon to be an IPv6 address. + return ":" in host or host.replace(".", "").isdigit() + + +_cached_current_datetime: Optional[int] = None +_cached_formatted_datetime = "" + + +def rfc822_formatted_time() -> str: + global _cached_current_datetime + global _cached_formatted_datetime + + now = int(time.time()) + if now != _cached_current_datetime: + # Weekday and month names for HTTP date/time formatting; + # always English! + # Tuples are constants stored in codeobject! + _weekdayname = ("Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun") + _monthname = ( + "", # Dummy so we can use 1-based month numbers + "Jan", + "Feb", + "Mar", + "Apr", + "May", + "Jun", + "Jul", + "Aug", + "Sep", + "Oct", + "Nov", + "Dec", + ) + + year, month, day, hh, mm, ss, wd, *tail = time.gmtime(now) + _cached_formatted_datetime = "%s, %02d %3s %4d %02d:%02d:%02d GMT" % ( + _weekdayname[wd], + day, + _monthname[month], + year, + hh, + mm, + ss, + ) + _cached_current_datetime = now + return _cached_formatted_datetime + + +def _weakref_handle(info: "Tuple[weakref.ref[object], str]") -> None: + ref, name = info + ob = ref() + if ob is not None: + with suppress(Exception): + getattr(ob, name)() + + +def weakref_handle( + ob: object, + name: str, + timeout: float, + loop: asyncio.AbstractEventLoop, + timeout_ceil_threshold: float = 5, +) -> Optional[asyncio.TimerHandle]: + if timeout is not None and timeout > 0: + when = loop.time() + timeout + if timeout >= timeout_ceil_threshold: + when = ceil(when) + + return loop.call_at(when, _weakref_handle, (weakref.ref(ob), name)) + return None + + +def call_later( + cb: Callable[[], Any], + timeout: float, + loop: asyncio.AbstractEventLoop, + timeout_ceil_threshold: float = 5, +) -> Optional[asyncio.TimerHandle]: + if timeout is None or timeout <= 0: + return None + now = loop.time() + when = calculate_timeout_when(now, timeout, timeout_ceil_threshold) + return loop.call_at(when, cb) + + +def calculate_timeout_when( + loop_time: float, + timeout: float, + timeout_ceiling_threshold: float, +) -> float: + """Calculate when to execute a timeout.""" + when = loop_time + timeout + if timeout > timeout_ceiling_threshold: + return ceil(when) + return when + + +class TimeoutHandle: + """Timeout handle""" + + __slots__ = ("_timeout", "_loop", "_ceil_threshold", "_callbacks") + + def __init__( + self, + loop: asyncio.AbstractEventLoop, + timeout: Optional[float], + ceil_threshold: float = 5, + ) -> None: + self._timeout = timeout + self._loop = loop + self._ceil_threshold = ceil_threshold + self._callbacks: List[ + Tuple[Callable[..., None], Tuple[Any, ...], Dict[str, Any]] + ] = [] + + def register( + self, callback: Callable[..., None], *args: Any, **kwargs: Any + ) -> None: + self._callbacks.append((callback, args, kwargs)) + + def close(self) -> None: + self._callbacks.clear() + + def start(self) -> Optional[asyncio.TimerHandle]: + timeout = self._timeout + if timeout is not None and timeout > 0: + when = self._loop.time() + timeout + if timeout >= self._ceil_threshold: + when = ceil(when) + return self._loop.call_at(when, self.__call__) + else: + return None + + def timer(self) -> "BaseTimerContext": + if self._timeout is not None and self._timeout > 0: + timer = TimerContext(self._loop) + self.register(timer.timeout) + return timer + else: + return TimerNoop() + + def __call__(self) -> None: + for cb, args, kwargs in self._callbacks: + with suppress(Exception): + cb(*args, **kwargs) + + self._callbacks.clear() + + +class BaseTimerContext(ContextManager["BaseTimerContext"]): + + __slots__ = () + + def assert_timeout(self) -> None: + """Raise TimeoutError if timeout has been exceeded.""" + + +class TimerNoop(BaseTimerContext): + + __slots__ = () + + def __enter__(self) -> BaseTimerContext: + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + return + + +class TimerContext(BaseTimerContext): + """Low resolution timeout context manager""" + + __slots__ = ("_loop", "_tasks", "_cancelled", "_cancelling") + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop = loop + self._tasks: List[asyncio.Task[Any]] = [] + self._cancelled = False + self._cancelling = 0 + + def assert_timeout(self) -> None: + """Raise TimeoutError if timer has already been cancelled.""" + if self._cancelled: + raise asyncio.TimeoutError from None + + def __enter__(self) -> BaseTimerContext: + task = asyncio.current_task(loop=self._loop) + if task is None: + raise RuntimeError("Timeout context manager should be used inside a task") + + if sys.version_info >= (3, 11): + # Remember if the task was already cancelling + # so when we __exit__ we can decide if we should + # raise asyncio.TimeoutError or let the cancellation propagate + self._cancelling = task.cancelling() + + if self._cancelled: + raise asyncio.TimeoutError from None + + self._tasks.append(task) + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> Optional[bool]: + enter_task: Optional[asyncio.Task[Any]] = None + if self._tasks: + enter_task = self._tasks.pop() + + if exc_type is asyncio.CancelledError and self._cancelled: + assert enter_task is not None + # The timeout was hit, and the task was cancelled + # so we need to uncancel the last task that entered the context manager + # since the cancellation should not leak out of the context manager + if sys.version_info >= (3, 11): + # If the task was already cancelling don't raise + # asyncio.TimeoutError and instead return None + # to allow the cancellation to propagate + if enter_task.uncancel() > self._cancelling: + return None + raise asyncio.TimeoutError from exc_val + return None + + def timeout(self) -> None: + if not self._cancelled: + for task in set(self._tasks): + task.cancel() + + self._cancelled = True + + +def ceil_timeout( + delay: Optional[float], ceil_threshold: float = 5 +) -> async_timeout.Timeout: + if delay is None or delay <= 0: + return async_timeout.timeout(None) + + loop = asyncio.get_running_loop() + now = loop.time() + when = now + delay + if delay > ceil_threshold: + when = ceil(when) + return async_timeout.timeout_at(when) + + +class HeadersMixin: + """Mixin for handling headers.""" + + ATTRS = frozenset(["_content_type", "_content_dict", "_stored_content_type"]) + + _headers: MultiMapping[str] + _content_type: Optional[str] = None + _content_dict: Optional[Dict[str, str]] = None + _stored_content_type: Union[str, None, _SENTINEL] = sentinel + + def _parse_content_type(self, raw: Optional[str]) -> None: + self._stored_content_type = raw + if raw is None: + # default value according to RFC 2616 + self._content_type = "application/octet-stream" + self._content_dict = {} + else: + content_type, content_mapping_proxy = parse_content_type(raw) + self._content_type = content_type + # _content_dict needs to be mutable so we can update it + self._content_dict = content_mapping_proxy.copy() + + @property + def content_type(self) -> str: + """The value of content part for Content-Type HTTP header.""" + raw = self._headers.get(hdrs.CONTENT_TYPE) + if self._stored_content_type != raw: + self._parse_content_type(raw) + assert self._content_type is not None + return self._content_type + + @property + def charset(self) -> Optional[str]: + """The value of charset part for Content-Type HTTP header.""" + raw = self._headers.get(hdrs.CONTENT_TYPE) + if self._stored_content_type != raw: + self._parse_content_type(raw) + assert self._content_dict is not None + return self._content_dict.get("charset") + + @property + def content_length(self) -> Optional[int]: + """The value of Content-Length HTTP header.""" + content_length = self._headers.get(hdrs.CONTENT_LENGTH) + return None if content_length is None else int(content_length) + + +def set_result(fut: "asyncio.Future[_T]", result: _T) -> None: + if not fut.done(): + fut.set_result(result) + + +_EXC_SENTINEL = BaseException() + + +class ErrorableProtocol(Protocol): + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = ..., + ) -> None: ... # pragma: no cover + + +def set_exception( + fut: "asyncio.Future[_T] | ErrorableProtocol", + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, +) -> None: + """Set future exception. + + If the future is marked as complete, this function is a no-op. + + :param exc_cause: An exception that is a direct cause of ``exc``. + Only set if provided. + """ + if asyncio.isfuture(fut) and fut.done(): + return + + exc_is_sentinel = exc_cause is _EXC_SENTINEL + exc_causes_itself = exc is exc_cause + if not exc_is_sentinel and not exc_causes_itself: + exc.__cause__ = exc_cause + + fut.set_exception(exc) + + +@functools.total_ordering +class AppKey(Generic[_T]): + """Keys for static typing support in Application.""" + + __slots__ = ("_name", "_t", "__orig_class__") + + # This may be set by Python when instantiating with a generic type. We need to + # support this, in order to support types that are not concrete classes, + # like Iterable, which can't be passed as the second parameter to __init__. + __orig_class__: Type[object] + + def __init__(self, name: str, t: Optional[Type[_T]] = None): + # Prefix with module name to help deduplicate key names. + frame = inspect.currentframe() + while frame: + if frame.f_code.co_name == "": + module: str = frame.f_globals["__name__"] + break + frame = frame.f_back + + self._name = module + "." + name + self._t = t + + def __lt__(self, other: object) -> bool: + if isinstance(other, AppKey): + return self._name < other._name + return True # Order AppKey above other types. + + def __repr__(self) -> str: + t = self._t + if t is None: + with suppress(AttributeError): + # Set to type arg. + t = get_args(self.__orig_class__)[0] + + if t is None: + t_repr = "<>" + elif isinstance(t, type): + if t.__module__ == "builtins": + t_repr = t.__qualname__ + else: + t_repr = f"{t.__module__}.{t.__qualname__}" + else: + t_repr = repr(t) + return f"" + + +class ChainMapProxy(Mapping[Union[str, AppKey[Any]], Any]): + __slots__ = ("_maps",) + + def __init__(self, maps: Iterable[Mapping[Union[str, AppKey[Any]], Any]]) -> None: + self._maps = tuple(maps) + + def __init_subclass__(cls) -> None: + raise TypeError( + "Inheritance class {} from ChainMapProxy " + "is forbidden".format(cls.__name__) + ) + + @overload # type: ignore[override] + def __getitem__(self, key: AppKey[_T]) -> _T: ... + + @overload + def __getitem__(self, key: str) -> Any: ... + + def __getitem__(self, key: Union[str, AppKey[_T]]) -> Any: + for mapping in self._maps: + try: + return mapping[key] + except KeyError: + pass + raise KeyError(key) + + @overload # type: ignore[override] + def get(self, key: AppKey[_T], default: _S) -> Union[_T, _S]: ... + + @overload + def get(self, key: AppKey[_T], default: None = ...) -> Optional[_T]: ... + + @overload + def get(self, key: str, default: Any = ...) -> Any: ... + + def get(self, key: Union[str, AppKey[_T]], default: Any = None) -> Any: + try: + return self[key] + except KeyError: + return default + + def __len__(self) -> int: + # reuses stored hash values if possible + return len(set().union(*self._maps)) + + def __iter__(self) -> Iterator[Union[str, AppKey[Any]]]: + d: Dict[Union[str, AppKey[Any]], Any] = {} + for mapping in reversed(self._maps): + # reuses stored hash values if possible + d.update(mapping) + return iter(d) + + def __contains__(self, key: object) -> bool: + return any(key in m for m in self._maps) + + def __bool__(self) -> bool: + return any(self._maps) + + def __repr__(self) -> str: + content = ", ".join(map(repr, self._maps)) + return f"ChainMapProxy({content})" + + +# https://tools.ietf.org/html/rfc7232#section-2.3 +_ETAGC = r"[!\x23-\x7E\x80-\xff]+" +_ETAGC_RE = re.compile(_ETAGC) +_QUOTED_ETAG = rf'(W/)?"({_ETAGC})"' +QUOTED_ETAG_RE = re.compile(_QUOTED_ETAG) +LIST_QUOTED_ETAG_RE = re.compile(rf"({_QUOTED_ETAG})(?:\s*,\s*|$)|(.)") + +ETAG_ANY = "*" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class ETag: + value: str + is_weak: bool = False + + +def validate_etag_value(value: str) -> None: + if value != ETAG_ANY and not _ETAGC_RE.fullmatch(value): + raise ValueError( + f"Value {value!r} is not a valid etag. Maybe it contains '\"'?" + ) + + +def parse_http_date(date_str: Optional[str]) -> Optional[datetime.datetime]: + """Process a date string, return a datetime object""" + if date_str is not None: + timetuple = parsedate(date_str) + if timetuple is not None: + with suppress(ValueError): + return datetime.datetime(*timetuple[:6], tzinfo=datetime.timezone.utc) + return None + + +@functools.lru_cache +def must_be_empty_body(method: str, code: int) -> bool: + """Check if a request must return an empty body.""" + return ( + code in EMPTY_BODY_STATUS_CODES + or method in EMPTY_BODY_METHODS + or (200 <= code < 300 and method in hdrs.METH_CONNECT_ALL) + ) + + +def should_remove_content_length(method: str, code: int) -> bool: + """Check if a Content-Length header should be removed. + + This should always be a subset of must_be_empty_body + """ + # https://www.rfc-editor.org/rfc/rfc9110.html#section-8.6-8 + # https://www.rfc-editor.org/rfc/rfc9110.html#section-15.4.5-4 + return code in EMPTY_BODY_STATUS_CODES or ( + 200 <= code < 300 and method in hdrs.METH_CONNECT_ALL + ) diff --git a/lib/python3.12/site-packages/aiohttp/http.py b/lib/python3.12/site-packages/aiohttp/http.py new file mode 100644 index 0000000000000000000000000000000000000000..a1feae2d9b8fe631d539a15dbf8e5ea2914d70d5 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/http.py @@ -0,0 +1,72 @@ +import sys +from http import HTTPStatus +from typing import Mapping, Tuple + +from . import __version__ +from .http_exceptions import HttpProcessingError as HttpProcessingError +from .http_parser import ( + HeadersParser as HeadersParser, + HttpParser as HttpParser, + HttpRequestParser as HttpRequestParser, + HttpResponseParser as HttpResponseParser, + RawRequestMessage as RawRequestMessage, + RawResponseMessage as RawResponseMessage, +) +from .http_websocket import ( + WS_CLOSED_MESSAGE as WS_CLOSED_MESSAGE, + WS_CLOSING_MESSAGE as WS_CLOSING_MESSAGE, + WS_KEY as WS_KEY, + WebSocketError as WebSocketError, + WebSocketReader as WebSocketReader, + WebSocketWriter as WebSocketWriter, + WSCloseCode as WSCloseCode, + WSMessage as WSMessage, + WSMsgType as WSMsgType, + ws_ext_gen as ws_ext_gen, + ws_ext_parse as ws_ext_parse, +) +from .http_writer import ( + HttpVersion as HttpVersion, + HttpVersion10 as HttpVersion10, + HttpVersion11 as HttpVersion11, + StreamWriter as StreamWriter, +) + +__all__ = ( + "HttpProcessingError", + "RESPONSES", + "SERVER_SOFTWARE", + # .http_writer + "StreamWriter", + "HttpVersion", + "HttpVersion10", + "HttpVersion11", + # .http_parser + "HeadersParser", + "HttpParser", + "HttpRequestParser", + "HttpResponseParser", + "RawRequestMessage", + "RawResponseMessage", + # .http_websocket + "WS_CLOSED_MESSAGE", + "WS_CLOSING_MESSAGE", + "WS_KEY", + "WebSocketReader", + "WebSocketWriter", + "ws_ext_gen", + "ws_ext_parse", + "WSMessage", + "WebSocketError", + "WSMsgType", + "WSCloseCode", +) + + +SERVER_SOFTWARE: str = "Python/{0[0]}.{0[1]} aiohttp/{1}".format( + sys.version_info, __version__ +) + +RESPONSES: Mapping[int, Tuple[str, str]] = { + v: (v.phrase, v.description) for v in HTTPStatus.__members__.values() +} diff --git a/lib/python3.12/site-packages/aiohttp/http_exceptions.py b/lib/python3.12/site-packages/aiohttp/http_exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..0b5867c7861ec0ba2558784be068476a129954cd --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/http_exceptions.py @@ -0,0 +1,116 @@ +"""Low-level http related exceptions.""" + +from textwrap import indent +from typing import Optional, Union + +from .typedefs import _CIMultiDict + +__all__ = ("HttpProcessingError",) + + +class HttpProcessingError(Exception): + """HTTP error. + + Shortcut for raising HTTP errors with custom code, message and headers. + + code: HTTP Error code. + message: (optional) Error message. + headers: (optional) Headers to be sent in response, a list of pairs + """ + + code = 0 + message = "" + headers = None + + def __init__( + self, + *, + code: Optional[int] = None, + message: str = "", + headers: Optional[_CIMultiDict] = None, + ) -> None: + if code is not None: + self.code = code + self.headers = headers + self.message = message + + def __str__(self) -> str: + msg = indent(self.message, " ") + return f"{self.code}, message:\n{msg}" + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}: {self.code}, message={self.message!r}>" + + +class BadHttpMessage(HttpProcessingError): + + code = 400 + message = "Bad Request" + + def __init__(self, message: str, *, headers: Optional[_CIMultiDict] = None) -> None: + super().__init__(message=message, headers=headers) + self.args = (message,) + + +class HttpBadRequest(BadHttpMessage): + + code = 400 + message = "Bad Request" + + +class PayloadEncodingError(BadHttpMessage): + """Base class for payload errors""" + + +class ContentEncodingError(PayloadEncodingError): + """Content encoding error.""" + + +class TransferEncodingError(PayloadEncodingError): + """transfer encoding error.""" + + +class ContentLengthError(PayloadEncodingError): + """Not enough data to satisfy content length header.""" + + +class DecompressSizeError(PayloadEncodingError): + """Decompressed size exceeds the configured limit.""" + + +class LineTooLong(BadHttpMessage): + def __init__( + self, line: str, limit: str = "Unknown", actual_size: str = "Unknown" + ) -> None: + super().__init__( + f"Got more than {limit} bytes ({actual_size}) when reading {line}." + ) + self.args = (line, limit, actual_size) + + +class InvalidHeader(BadHttpMessage): + def __init__(self, hdr: Union[bytes, str]) -> None: + hdr_s = hdr.decode(errors="backslashreplace") if isinstance(hdr, bytes) else hdr + super().__init__(f"Invalid HTTP header: {hdr!r}") + self.hdr = hdr_s + self.args = (hdr,) + + +class BadStatusLine(BadHttpMessage): + def __init__(self, line: str = "", error: Optional[str] = None) -> None: + if not isinstance(line, str): + line = repr(line) + super().__init__(error or f"Bad status line {line!r}") + self.args = (line,) + self.line = line + + +class BadHttpMethod(BadStatusLine): + """Invalid HTTP method in status line.""" + + def __init__(self, line: str = "", error: Optional[str] = None) -> None: + super().__init__(line, error or f"Bad HTTP method in status line {line!r}") + + +class InvalidURLError(BadHttpMessage): + pass diff --git a/lib/python3.12/site-packages/aiohttp/http_parser.py b/lib/python3.12/site-packages/aiohttp/http_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..393e76a15865f3f536a4483ab7a3b1cfcd0702e5 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/http_parser.py @@ -0,0 +1,1086 @@ +import abc +import asyncio +import re +import string +from contextlib import suppress +from enum import IntEnum +from typing import ( + Any, + ClassVar, + Final, + Generic, + List, + Literal, + NamedTuple, + Optional, + Pattern, + Set, + Tuple, + Type, + TypeVar, + Union, +) + +from multidict import CIMultiDict, CIMultiDictProxy, istr +from yarl import URL + +from . import hdrs +from .base_protocol import BaseProtocol +from .compression_utils import ( + DEFAULT_MAX_DECOMPRESS_SIZE, + HAS_BROTLI, + HAS_ZSTD, + BrotliDecompressor, + ZLibDecompressor, + ZSTDDecompressor, +) +from .helpers import ( + _EXC_SENTINEL, + DEBUG, + EMPTY_BODY_METHODS, + EMPTY_BODY_STATUS_CODES, + NO_EXTENSIONS, + BaseTimerContext, + set_exception, +) +from .http_exceptions import ( + BadHttpMessage, + BadHttpMethod, + BadStatusLine, + ContentEncodingError, + ContentLengthError, + DecompressSizeError, + InvalidHeader, + InvalidURLError, + LineTooLong, + TransferEncodingError, +) +from .http_writer import HttpVersion, HttpVersion10 +from .streams import EMPTY_PAYLOAD, StreamReader +from .typedefs import RawHeaders + +__all__ = ( + "HeadersParser", + "HttpParser", + "HttpRequestParser", + "HttpResponseParser", + "RawRequestMessage", + "RawResponseMessage", +) + +_SEP = Literal[b"\r\n", b"\n"] + +ASCIISET: Final[Set[str]] = set(string.printable) + +# See https://www.rfc-editor.org/rfc/rfc9110.html#name-overview +# and https://www.rfc-editor.org/rfc/rfc9110.html#name-tokens +# +# method = token +# tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." / +# "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA +# token = 1*tchar +_TCHAR_SPECIALS: Final[str] = re.escape("!#$%&'*+-.^_`|~") +TOKENRE: Final[Pattern[str]] = re.compile(f"[0-9A-Za-z{_TCHAR_SPECIALS}]+") +VERSRE: Final[Pattern[str]] = re.compile(r"HTTP/(\d)\.(\d)", re.ASCII) +DIGITS: Final[Pattern[str]] = re.compile(r"\d+", re.ASCII) +HEXDIGITS: Final[Pattern[bytes]] = re.compile(rb"[0-9a-fA-F]+") + + +class RawRequestMessage(NamedTuple): + method: str + path: str + version: HttpVersion + headers: "CIMultiDictProxy[str]" + raw_headers: RawHeaders + should_close: bool + compression: Optional[str] + upgrade: bool + chunked: bool + url: URL + + +class RawResponseMessage(NamedTuple): + version: HttpVersion + code: int + reason: str + headers: CIMultiDictProxy[str] + raw_headers: RawHeaders + should_close: bool + compression: Optional[str] + upgrade: bool + chunked: bool + + +_MsgT = TypeVar("_MsgT", RawRequestMessage, RawResponseMessage) + + +class ParseState(IntEnum): + + PARSE_NONE = 0 + PARSE_LENGTH = 1 + PARSE_CHUNKED = 2 + PARSE_UNTIL_EOF = 3 + + +class ChunkState(IntEnum): + PARSE_CHUNKED_SIZE = 0 + PARSE_CHUNKED_CHUNK = 1 + PARSE_CHUNKED_CHUNK_EOF = 2 + PARSE_MAYBE_TRAILERS = 3 + PARSE_TRAILERS = 4 + + +class HeadersParser: + def __init__( + self, + max_line_size: int = 8190, + max_headers: int = 32768, + max_field_size: int = 8190, + lax: bool = False, + ) -> None: + self.max_line_size = max_line_size + self.max_headers = max_headers + self.max_field_size = max_field_size + self._lax = lax + + def parse_headers( + self, lines: List[bytes] + ) -> Tuple["CIMultiDictProxy[str]", RawHeaders]: + headers: CIMultiDict[str] = CIMultiDict() + # note: "raw" does not mean inclusion of OWS before/after the field value + raw_headers = [] + + lines_idx = 0 + line = lines[lines_idx] + line_count = len(lines) + + while line: + # Parse initial header name : value pair. + try: + bname, bvalue = line.split(b":", 1) + except ValueError: + raise InvalidHeader(line) from None + + if len(bname) == 0: + raise InvalidHeader(bname) + + # https://www.rfc-editor.org/rfc/rfc9112.html#section-5.1-2 + if {bname[0], bname[-1]} & {32, 9}: # {" ", "\t"} + raise InvalidHeader(line) + + bvalue = bvalue.lstrip(b" \t") + if len(bname) > self.max_field_size: + raise LineTooLong( + "request header name {}".format( + bname.decode("utf8", "backslashreplace") + ), + str(self.max_field_size), + str(len(bname)), + ) + name = bname.decode("utf-8", "surrogateescape") + if not TOKENRE.fullmatch(name): + raise InvalidHeader(bname) + + header_length = len(bvalue) + + # next line + lines_idx += 1 + line = lines[lines_idx] + + # consume continuation lines + continuation = self._lax and line and line[0] in (32, 9) # (' ', '\t') + + # Deprecated: https://www.rfc-editor.org/rfc/rfc9112.html#name-obsolete-line-folding + if continuation: + bvalue_lst = [bvalue] + while continuation: + header_length += len(line) + if header_length > self.max_field_size: + raise LineTooLong( + "request header field {}".format( + bname.decode("utf8", "backslashreplace") + ), + str(self.max_field_size), + str(header_length), + ) + bvalue_lst.append(line) + + # next line + lines_idx += 1 + if lines_idx < line_count: + line = lines[lines_idx] + if line: + continuation = line[0] in (32, 9) # (' ', '\t') + else: + line = b"" + break + bvalue = b"".join(bvalue_lst) + else: + if header_length > self.max_field_size: + raise LineTooLong( + "request header field {}".format( + bname.decode("utf8", "backslashreplace") + ), + str(self.max_field_size), + str(header_length), + ) + + bvalue = bvalue.strip(b" \t") + value = bvalue.decode("utf-8", "surrogateescape") + + # https://www.rfc-editor.org/rfc/rfc9110.html#section-5.5-5 + if "\n" in value or "\r" in value or "\x00" in value: + raise InvalidHeader(bvalue) + + headers.add(name, value) + raw_headers.append((bname, bvalue)) + + return (CIMultiDictProxy(headers), tuple(raw_headers)) + + +def _is_supported_upgrade(headers: CIMultiDictProxy[str]) -> bool: + """Check if the upgrade header is supported.""" + u = headers.get(hdrs.UPGRADE, "") + # .lower() can transform non-ascii characters. + return u.isascii() and u.lower() in {"tcp", "websocket"} + + +class HttpParser(abc.ABC, Generic[_MsgT]): + lax: ClassVar[bool] = False + + def __init__( + self, + protocol: Optional[BaseProtocol] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + limit: int = 2**16, + max_line_size: int = 8190, + max_headers: int = 32768, + max_field_size: int = 8190, + timer: Optional[BaseTimerContext] = None, + code: Optional[int] = None, + method: Optional[str] = None, + payload_exception: Optional[Type[BaseException]] = None, + response_with_body: bool = True, + read_until_eof: bool = False, + auto_decompress: bool = True, + ) -> None: + self.protocol = protocol + self.loop = loop + self.max_line_size = max_line_size + self.max_headers = max_headers + self.max_field_size = max_field_size + self.timer = timer + self.code = code + self.method = method + self.payload_exception = payload_exception + self.response_with_body = response_with_body + self.read_until_eof = read_until_eof + + self._lines: List[bytes] = [] + self._tail = b"" + self._upgraded = False + self._payload = None + self._payload_parser: Optional[HttpPayloadParser] = None + self._auto_decompress = auto_decompress + self._limit = limit + self._headers_parser = HeadersParser( + max_line_size, max_headers, max_field_size, self.lax + ) + + @abc.abstractmethod + def parse_message(self, lines: List[bytes]) -> _MsgT: ... + + @abc.abstractmethod + def _is_chunked_te(self, te: str) -> bool: ... + + def feed_eof(self) -> Optional[_MsgT]: + if self._payload_parser is not None: + self._payload_parser.feed_eof() + self._payload_parser = None + else: + # try to extract partial message + if self._tail: + self._lines.append(self._tail) + + if self._lines: + if self._lines[-1] != "\r\n": + self._lines.append(b"") + with suppress(Exception): + return self.parse_message(self._lines) + return None + + def feed_data( + self, + data: bytes, + SEP: _SEP = b"\r\n", + EMPTY: bytes = b"", + CONTENT_LENGTH: istr = hdrs.CONTENT_LENGTH, + METH_CONNECT: str = hdrs.METH_CONNECT, + SEC_WEBSOCKET_KEY1: istr = hdrs.SEC_WEBSOCKET_KEY1, + ) -> Tuple[List[Tuple[_MsgT, StreamReader]], bool, bytes]: + + messages = [] + + if self._tail: + data, self._tail = self._tail + data, b"" + + data_len = len(data) + start_pos = 0 + loop = self.loop + + should_close = False + while start_pos < data_len: + + # read HTTP message (request/response line + headers), \r\n\r\n + # and split by lines + if self._payload_parser is None and not self._upgraded: + pos = data.find(SEP, start_pos) + # consume \r\n + if pos == start_pos and not self._lines: + start_pos = pos + len(SEP) + continue + + if pos >= start_pos: + if should_close: + raise BadHttpMessage("Data after `Connection: close`") + + # line found + line = data[start_pos:pos] + if SEP == b"\n": # For lax response parsing + line = line.rstrip(b"\r") + self._lines.append(line) + start_pos = pos + len(SEP) + + # \r\n\r\n found + if self._lines[-1] == EMPTY: + try: + msg: _MsgT = self.parse_message(self._lines) + finally: + self._lines.clear() + + def get_content_length() -> Optional[int]: + # payload length + length_hdr = msg.headers.get(CONTENT_LENGTH) + if length_hdr is None: + return None + + # Shouldn't allow +/- or other number formats. + # https://www.rfc-editor.org/rfc/rfc9110#section-8.6-2 + # msg.headers is already stripped of leading/trailing wsp + if not DIGITS.fullmatch(length_hdr): + raise InvalidHeader(CONTENT_LENGTH) + + return int(length_hdr) + + length = get_content_length() + # do not support old websocket spec + if SEC_WEBSOCKET_KEY1 in msg.headers: + raise InvalidHeader(SEC_WEBSOCKET_KEY1) + + self._upgraded = msg.upgrade and _is_supported_upgrade( + msg.headers + ) + + method = getattr(msg, "method", self.method) + # code is only present on responses + code = getattr(msg, "code", 0) + + assert self.protocol is not None + # calculate payload + empty_body = code in EMPTY_BODY_STATUS_CODES or bool( + method and method in EMPTY_BODY_METHODS + ) + if not empty_body and ( + ((length is not None and length > 0) or msg.chunked) + and not self._upgraded + ): + payload = StreamReader( + self.protocol, + timer=self.timer, + loop=loop, + limit=self._limit, + ) + payload_parser = HttpPayloadParser( + payload, + length=length, + chunked=msg.chunked, + method=method, + compression=msg.compression, + code=self.code, + response_with_body=self.response_with_body, + auto_decompress=self._auto_decompress, + lax=self.lax, + headers_parser=self._headers_parser, + ) + if not payload_parser.done: + self._payload_parser = payload_parser + elif method == METH_CONNECT: + assert isinstance(msg, RawRequestMessage) + payload = StreamReader( + self.protocol, + timer=self.timer, + loop=loop, + limit=self._limit, + ) + self._upgraded = True + self._payload_parser = HttpPayloadParser( + payload, + method=msg.method, + compression=msg.compression, + auto_decompress=self._auto_decompress, + lax=self.lax, + headers_parser=self._headers_parser, + ) + elif not empty_body and length is None and self.read_until_eof: + payload = StreamReader( + self.protocol, + timer=self.timer, + loop=loop, + limit=self._limit, + ) + payload_parser = HttpPayloadParser( + payload, + length=length, + chunked=msg.chunked, + method=method, + compression=msg.compression, + code=self.code, + response_with_body=self.response_with_body, + auto_decompress=self._auto_decompress, + lax=self.lax, + headers_parser=self._headers_parser, + ) + if not payload_parser.done: + self._payload_parser = payload_parser + else: + payload = EMPTY_PAYLOAD + + messages.append((msg, payload)) + should_close = msg.should_close + else: + self._tail = data[start_pos:] + data = EMPTY + break + + # no parser, just store + elif self._payload_parser is None and self._upgraded: + assert not self._lines + break + + # feed payload + elif data and start_pos < data_len: + assert not self._lines + assert self._payload_parser is not None + try: + eof, data = self._payload_parser.feed_data(data[start_pos:], SEP) + except BaseException as underlying_exc: + reraised_exc = underlying_exc + if self.payload_exception is not None: + reraised_exc = self.payload_exception(str(underlying_exc)) + + set_exception( + self._payload_parser.payload, + reraised_exc, + underlying_exc, + ) + + eof = True + data = b"" + if isinstance( + underlying_exc, (InvalidHeader, TransferEncodingError) + ): + raise + + if eof: + start_pos = 0 + data_len = len(data) + self._payload_parser = None + continue + else: + break + + if data and start_pos < data_len: + data = data[start_pos:] + else: + data = EMPTY + + return messages, self._upgraded, data + + def parse_headers( + self, lines: List[bytes] + ) -> Tuple[ + "CIMultiDictProxy[str]", RawHeaders, Optional[bool], Optional[str], bool, bool + ]: + """Parses RFC 5322 headers from a stream. + + Line continuations are supported. Returns list of header name + and value pairs. Header name is in upper case. + """ + headers, raw_headers = self._headers_parser.parse_headers(lines) + close_conn = None + encoding = None + upgrade = False + chunked = False + + # https://www.rfc-editor.org/rfc/rfc9110.html#section-5.5-6 + # https://www.rfc-editor.org/rfc/rfc9110.html#name-collected-abnf + singletons = ( + hdrs.CONTENT_LENGTH, + hdrs.CONTENT_LOCATION, + hdrs.CONTENT_RANGE, + hdrs.CONTENT_TYPE, + hdrs.ETAG, + hdrs.HOST, + hdrs.MAX_FORWARDS, + hdrs.SERVER, + hdrs.TRANSFER_ENCODING, + hdrs.USER_AGENT, + ) + bad_hdr = next((h for h in singletons if len(headers.getall(h, ())) > 1), None) + if bad_hdr is not None: + raise BadHttpMessage(f"Duplicate '{bad_hdr}' header found.") + + # keep-alive + conn = headers.get(hdrs.CONNECTION) + if conn: + v = conn.lower() + if v == "close": + close_conn = True + elif v == "keep-alive": + close_conn = False + # https://www.rfc-editor.org/rfc/rfc9110.html#name-101-switching-protocols + elif v == "upgrade" and headers.get(hdrs.UPGRADE): + upgrade = True + + # encoding + enc = headers.get(hdrs.CONTENT_ENCODING, "") + if enc.isascii() and enc.lower() in {"gzip", "deflate", "br", "zstd"}: + encoding = enc + + # chunking + te = headers.get(hdrs.TRANSFER_ENCODING) + if te is not None: + if self._is_chunked_te(te): + chunked = True + + if hdrs.CONTENT_LENGTH in headers: + raise BadHttpMessage( + "Transfer-Encoding can't be present with Content-Length", + ) + + return (headers, raw_headers, close_conn, encoding, upgrade, chunked) + + def set_upgraded(self, val: bool) -> None: + """Set connection upgraded (to websocket) mode. + + :param bool val: new state. + """ + self._upgraded = val + + +class HttpRequestParser(HttpParser[RawRequestMessage]): + """Read request status line. + + Exception .http_exceptions.BadStatusLine + could be raised in case of any errors in status line. + Returns RawRequestMessage. + """ + + def parse_message(self, lines: List[bytes]) -> RawRequestMessage: + # request line + line = lines[0].decode("utf-8", "surrogateescape") + try: + method, path, version = line.split(" ", maxsplit=2) + except ValueError: + raise BadHttpMethod(line) from None + + if len(path) > self.max_line_size: + raise LineTooLong( + "Status line is too long", str(self.max_line_size), str(len(path)) + ) + + # method + if not TOKENRE.fullmatch(method): + raise BadHttpMethod(method) + + # version + match = VERSRE.fullmatch(version) + if match is None: + raise BadStatusLine(line) + version_o = HttpVersion(int(match.group(1)), int(match.group(2))) + + if method == "CONNECT": + # authority-form, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.3 + url = URL.build(authority=path, encoded=True) + elif path.startswith("/"): + # origin-form, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.1 + path_part, _hash_separator, url_fragment = path.partition("#") + path_part, _question_mark_separator, qs_part = path_part.partition("?") + + # NOTE: `yarl.URL.build()` is used to mimic what the Cython-based + # NOTE: parser does, otherwise it results into the same + # NOTE: HTTP Request-Line input producing different + # NOTE: `yarl.URL()` objects + url = URL.build( + path=path_part, + query_string=qs_part, + fragment=url_fragment, + encoded=True, + ) + elif path == "*" and method == "OPTIONS": + # asterisk-form, + url = URL(path, encoded=True) + else: + # absolute-form for proxy maybe, + # https://datatracker.ietf.org/doc/html/rfc7230#section-5.3.2 + url = URL(path, encoded=True) + if url.scheme == "": + # not absolute-form + raise InvalidURLError( + path.encode(errors="surrogateescape").decode("latin1") + ) + + # read headers + ( + headers, + raw_headers, + close, + compression, + upgrade, + chunked, + ) = self.parse_headers(lines[1:]) + + if close is None: # then the headers weren't set in the request + if version_o <= HttpVersion10: # HTTP 1.0 must asks to not close + close = True + else: # HTTP 1.1 must ask to close. + close = False + + return RawRequestMessage( + method, + path, + version_o, + headers, + raw_headers, + close, + compression, + upgrade, + chunked, + url, + ) + + def _is_chunked_te(self, te: str) -> bool: + te = te.rsplit(",", maxsplit=1)[-1].strip(" \t") + # .lower() transforms some non-ascii chars, so must check first. + if te.isascii() and te.lower() == "chunked": + return True + # https://www.rfc-editor.org/rfc/rfc9112#section-6.3-2.4.3 + raise BadHttpMessage("Request has invalid `Transfer-Encoding`") + + +class HttpResponseParser(HttpParser[RawResponseMessage]): + """Read response status line and headers. + + BadStatusLine could be raised in case of any errors in status line. + Returns RawResponseMessage. + """ + + # Lax mode should only be enabled on response parser. + lax = not DEBUG + + def feed_data( + self, + data: bytes, + SEP: Optional[_SEP] = None, + *args: Any, + **kwargs: Any, + ) -> Tuple[List[Tuple[RawResponseMessage, StreamReader]], bool, bytes]: + if SEP is None: + SEP = b"\r\n" if DEBUG else b"\n" + return super().feed_data(data, SEP, *args, **kwargs) + + def parse_message(self, lines: List[bytes]) -> RawResponseMessage: + line = lines[0].decode("utf-8", "surrogateescape") + try: + version, status = line.split(maxsplit=1) + except ValueError: + raise BadStatusLine(line) from None + + try: + status, reason = status.split(maxsplit=1) + except ValueError: + status = status.strip() + reason = "" + + if len(reason) > self.max_line_size: + raise LineTooLong( + "Status line is too long", str(self.max_line_size), str(len(reason)) + ) + + # version + match = VERSRE.fullmatch(version) + if match is None: + raise BadStatusLine(line) + version_o = HttpVersion(int(match.group(1)), int(match.group(2))) + + # The status code is a three-digit ASCII number, no padding + if len(status) != 3 or not DIGITS.fullmatch(status): + raise BadStatusLine(line) + status_i = int(status) + + # read headers + ( + headers, + raw_headers, + close, + compression, + upgrade, + chunked, + ) = self.parse_headers(lines[1:]) + + if close is None: + if version_o <= HttpVersion10: + close = True + # https://www.rfc-editor.org/rfc/rfc9112.html#name-message-body-length + elif 100 <= status_i < 200 or status_i in {204, 304}: + close = False + elif hdrs.CONTENT_LENGTH in headers or hdrs.TRANSFER_ENCODING in headers: + close = False + else: + # https://www.rfc-editor.org/rfc/rfc9112.html#section-6.3-2.8 + close = True + + return RawResponseMessage( + version_o, + status_i, + reason.strip(), + headers, + raw_headers, + close, + compression, + upgrade, + chunked, + ) + + def _is_chunked_te(self, te: str) -> bool: + # https://www.rfc-editor.org/rfc/rfc9112#section-6.3-2.4.2 + return te.rsplit(",", maxsplit=1)[-1].strip(" \t").lower() == "chunked" + + +class HttpPayloadParser: + def __init__( + self, + payload: StreamReader, + length: Optional[int] = None, + chunked: bool = False, + compression: Optional[str] = None, + code: Optional[int] = None, + method: Optional[str] = None, + response_with_body: bool = True, + auto_decompress: bool = True, + lax: bool = False, + *, + headers_parser: HeadersParser, + ) -> None: + self._length = 0 + self._type = ParseState.PARSE_UNTIL_EOF + self._chunk = ChunkState.PARSE_CHUNKED_SIZE + self._chunk_size = 0 + self._chunk_tail = b"" + self._auto_decompress = auto_decompress + self._lax = lax + self._headers_parser = headers_parser + self._trailer_lines: list[bytes] = [] + self.done = False + + # payload decompression wrapper + if response_with_body and compression and self._auto_decompress: + real_payload: Union[StreamReader, DeflateBuffer] = DeflateBuffer( + payload, compression + ) + else: + real_payload = payload + + # payload parser + if not response_with_body: + # don't parse payload if it's not expected to be received + self._type = ParseState.PARSE_NONE + real_payload.feed_eof() + self.done = True + elif chunked: + self._type = ParseState.PARSE_CHUNKED + elif length is not None: + self._type = ParseState.PARSE_LENGTH + self._length = length + if self._length == 0: + real_payload.feed_eof() + self.done = True + + self.payload = real_payload + + def feed_eof(self) -> None: + if self._type == ParseState.PARSE_UNTIL_EOF: + self.payload.feed_eof() + elif self._type == ParseState.PARSE_LENGTH: + raise ContentLengthError( + "Not enough data to satisfy content length header." + ) + elif self._type == ParseState.PARSE_CHUNKED: + raise TransferEncodingError( + "Not enough data to satisfy transfer length header." + ) + + def feed_data( + self, chunk: bytes, SEP: _SEP = b"\r\n", CHUNK_EXT: bytes = b";" + ) -> Tuple[bool, bytes]: + # Read specified amount of bytes + if self._type == ParseState.PARSE_LENGTH: + required = self._length + chunk_len = len(chunk) + + if required >= chunk_len: + self._length = required - chunk_len + self.payload.feed_data(chunk, chunk_len) + if self._length == 0: + self.payload.feed_eof() + return True, b"" + else: + self._length = 0 + self.payload.feed_data(chunk[:required], required) + self.payload.feed_eof() + return True, chunk[required:] + + # Chunked transfer encoding parser + elif self._type == ParseState.PARSE_CHUNKED: + if self._chunk_tail: + chunk = self._chunk_tail + chunk + self._chunk_tail = b"" + + while chunk: + + # read next chunk size + if self._chunk == ChunkState.PARSE_CHUNKED_SIZE: + pos = chunk.find(SEP) + if pos >= 0: + i = chunk.find(CHUNK_EXT, 0, pos) + if i >= 0: + size_b = chunk[:i] # strip chunk-extensions + # Verify no LF in the chunk-extension + if b"\n" in (ext := chunk[i:pos]): + exc = TransferEncodingError( + f"Unexpected LF in chunk-extension: {ext!r}" + ) + set_exception(self.payload, exc) + raise exc + else: + size_b = chunk[:pos] + + if self._lax: # Allow whitespace in lax mode. + size_b = size_b.strip() + + if not re.fullmatch(HEXDIGITS, size_b): + exc = TransferEncodingError( + chunk[:pos].decode("ascii", "surrogateescape") + ) + set_exception(self.payload, exc) + raise exc + size = int(bytes(size_b), 16) + + chunk = chunk[pos + len(SEP) :] + if size == 0: # eof marker + self._chunk = ChunkState.PARSE_TRAILERS + if self._lax and chunk.startswith(b"\r"): + chunk = chunk[1:] + else: + self._chunk = ChunkState.PARSE_CHUNKED_CHUNK + self._chunk_size = size + self.payload.begin_http_chunk_receiving() + else: + self._chunk_tail = chunk + return False, b"" + + # read chunk and feed buffer + if self._chunk == ChunkState.PARSE_CHUNKED_CHUNK: + required = self._chunk_size + chunk_len = len(chunk) + + if required > chunk_len: + self._chunk_size = required - chunk_len + self.payload.feed_data(chunk, chunk_len) + return False, b"" + else: + self._chunk_size = 0 + self.payload.feed_data(chunk[:required], required) + chunk = chunk[required:] + self._chunk = ChunkState.PARSE_CHUNKED_CHUNK_EOF + self.payload.end_http_chunk_receiving() + + # toss the CRLF at the end of the chunk + if self._chunk == ChunkState.PARSE_CHUNKED_CHUNK_EOF: + if self._lax and chunk.startswith(b"\r"): + chunk = chunk[1:] + if chunk[: len(SEP)] == SEP: + chunk = chunk[len(SEP) :] + self._chunk = ChunkState.PARSE_CHUNKED_SIZE + else: + self._chunk_tail = chunk + return False, b"" + + if self._chunk == ChunkState.PARSE_TRAILERS: + pos = chunk.find(SEP) + if pos < 0: # No line found + self._chunk_tail = chunk + return False, b"" + + line = chunk[:pos] + chunk = chunk[pos + len(SEP) :] + if SEP == b"\n": # For lax response parsing + line = line.rstrip(b"\r") + self._trailer_lines.append(line) + + # \r\n\r\n found, end of stream + if self._trailer_lines[-1] == b"": + # Headers and trailers are defined the same way, + # so we reuse the HeadersParser here. + try: + trailers, raw_trailers = self._headers_parser.parse_headers( + self._trailer_lines + ) + finally: + self._trailer_lines.clear() + self.payload.feed_eof() + return True, chunk + + # Read all bytes until eof + elif self._type == ParseState.PARSE_UNTIL_EOF: + self.payload.feed_data(chunk, len(chunk)) + + return False, b"" + + +class DeflateBuffer: + """DeflateStream decompress stream and feed data into specified stream.""" + + decompressor: Any + + def __init__( + self, + out: StreamReader, + encoding: Optional[str], + max_decompress_size: int = DEFAULT_MAX_DECOMPRESS_SIZE, + ) -> None: + self.out = out + self.size = 0 + out.total_compressed_bytes = self.size + self.encoding = encoding + self._started_decoding = False + + self.decompressor: Union[BrotliDecompressor, ZLibDecompressor, ZSTDDecompressor] + if encoding == "br": + if not HAS_BROTLI: # pragma: no cover + raise ContentEncodingError( + "Can not decode content-encoding: brotli (br). " + "Please install `Brotli`" + ) + self.decompressor = BrotliDecompressor() + elif encoding == "zstd": + if not HAS_ZSTD: + raise ContentEncodingError( + "Can not decode content-encoding: zstandard (zstd). " + "Please install `backports.zstd`" + ) + self.decompressor = ZSTDDecompressor() + else: + self.decompressor = ZLibDecompressor(encoding=encoding) + + self._max_decompress_size = max_decompress_size + + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + set_exception(self.out, exc, exc_cause) + + def feed_data(self, chunk: bytes, size: int) -> None: + if not size: + return + + self.size += size + self.out.total_compressed_bytes = self.size + + # RFC1950 + # bits 0..3 = CM = 0b1000 = 8 = "deflate" + # bits 4..7 = CINFO = 1..7 = windows size. + if ( + not self._started_decoding + and self.encoding == "deflate" + and chunk[0] & 0xF != 8 + ): + # Change the decoder to decompress incorrectly compressed data + # Actually we should issue a warning about non-RFC-compliant data. + self.decompressor = ZLibDecompressor( + encoding=self.encoding, suppress_deflate_header=True + ) + + try: + # Decompress with limit + 1 so we can detect if output exceeds limit + chunk = self.decompressor.decompress_sync( + chunk, max_length=self._max_decompress_size + 1 + ) + except Exception: + raise ContentEncodingError( + "Can not decode content-encoding: %s" % self.encoding + ) + + self._started_decoding = True + + # Check if decompression limit was exceeded + if len(chunk) > self._max_decompress_size: + raise DecompressSizeError( + "Decompressed data exceeds the configured limit of %d bytes" + % self._max_decompress_size + ) + + if chunk: + self.out.feed_data(chunk, len(chunk)) + + def feed_eof(self) -> None: + chunk = self.decompressor.flush() + + if chunk or self.size > 0: + self.out.feed_data(chunk, len(chunk)) + if self.encoding == "deflate" and not self.decompressor.eof: + raise ContentEncodingError("deflate") + + self.out.feed_eof() + + def begin_http_chunk_receiving(self) -> None: + self.out.begin_http_chunk_receiving() + + def end_http_chunk_receiving(self) -> None: + self.out.end_http_chunk_receiving() + + +HttpRequestParserPy = HttpRequestParser +HttpResponseParserPy = HttpResponseParser +RawRequestMessagePy = RawRequestMessage +RawResponseMessagePy = RawResponseMessage + +try: + if not NO_EXTENSIONS: + from ._http_parser import ( # type: ignore[import-not-found,no-redef] + HttpRequestParser, + HttpResponseParser, + RawRequestMessage, + RawResponseMessage, + ) + + HttpRequestParserC = HttpRequestParser + HttpResponseParserC = HttpResponseParser + RawRequestMessageC = RawRequestMessage + RawResponseMessageC = RawResponseMessage +except ImportError: # pragma: no cover + pass diff --git a/lib/python3.12/site-packages/aiohttp/http_websocket.py b/lib/python3.12/site-packages/aiohttp/http_websocket.py new file mode 100644 index 0000000000000000000000000000000000000000..6b4b30e02b247e30e0c84d3eb118b749bbe52079 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/http_websocket.py @@ -0,0 +1,36 @@ +"""WebSocket protocol versions 13 and 8.""" + +from ._websocket.helpers import WS_KEY, ws_ext_gen, ws_ext_parse +from ._websocket.models import ( + WS_CLOSED_MESSAGE, + WS_CLOSING_MESSAGE, + WebSocketError, + WSCloseCode, + WSHandshakeError, + WSMessage, + WSMsgType, +) +from ._websocket.reader import WebSocketReader +from ._websocket.writer import WebSocketWriter + +# Messages that the WebSocketResponse.receive needs to handle internally +_INTERNAL_RECEIVE_TYPES = frozenset( + (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.PING, WSMsgType.PONG) +) + + +__all__ = ( + "WS_CLOSED_MESSAGE", + "WS_CLOSING_MESSAGE", + "WS_KEY", + "WebSocketReader", + "WebSocketWriter", + "WSMessage", + "WebSocketError", + "WSMsgType", + "WSCloseCode", + "ws_ext_gen", + "ws_ext_parse", + "WSHandshakeError", + "WSMessage", +) diff --git a/lib/python3.12/site-packages/aiohttp/http_writer.py b/lib/python3.12/site-packages/aiohttp/http_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..a140b218b25fedccb49451c547d4d01326367cfb --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/http_writer.py @@ -0,0 +1,378 @@ +"""Http related parsers and protocol.""" + +import asyncio +import sys +from typing import ( # noqa + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Iterable, + List, + NamedTuple, + Optional, + Union, +) + +from multidict import CIMultiDict + +from .abc import AbstractStreamWriter +from .base_protocol import BaseProtocol +from .client_exceptions import ClientConnectionResetError +from .compression_utils import ZLibCompressor +from .helpers import NO_EXTENSIONS + +__all__ = ("StreamWriter", "HttpVersion", "HttpVersion10", "HttpVersion11") + + +MIN_PAYLOAD_FOR_WRITELINES = 2048 +IS_PY313_BEFORE_313_2 = (3, 13, 0) <= sys.version_info < (3, 13, 2) +IS_PY_BEFORE_312_9 = sys.version_info < (3, 12, 9) +SKIP_WRITELINES = IS_PY313_BEFORE_313_2 or IS_PY_BEFORE_312_9 +# writelines is not safe for use +# on Python 3.12+ until 3.12.9 +# on Python 3.13+ until 3.13.2 +# and on older versions it not any faster than write +# CVE-2024-12254: https://github.com/python/cpython/pull/127656 + + +class HttpVersion(NamedTuple): + major: int + minor: int + + +HttpVersion10 = HttpVersion(1, 0) +HttpVersion11 = HttpVersion(1, 1) + + +_T_OnChunkSent = Optional[Callable[[bytes], Awaitable[None]]] +_T_OnHeadersSent = Optional[Callable[["CIMultiDict[str]"], Awaitable[None]]] + + +class StreamWriter(AbstractStreamWriter): + + length: Optional[int] = None + chunked: bool = False + _eof: bool = False + _compress: Optional[ZLibCompressor] = None + + def __init__( + self, + protocol: BaseProtocol, + loop: asyncio.AbstractEventLoop, + on_chunk_sent: _T_OnChunkSent = None, + on_headers_sent: _T_OnHeadersSent = None, + ) -> None: + self._protocol = protocol + self.loop = loop + self._on_chunk_sent: _T_OnChunkSent = on_chunk_sent + self._on_headers_sent: _T_OnHeadersSent = on_headers_sent + self._headers_buf: Optional[bytes] = None + self._headers_written: bool = False + + @property + def transport(self) -> Optional[asyncio.Transport]: + return self._protocol.transport + + @property + def protocol(self) -> BaseProtocol: + return self._protocol + + def enable_chunking(self) -> None: + self.chunked = True + + def enable_compression( + self, encoding: str = "deflate", strategy: Optional[int] = None + ) -> None: + self._compress = ZLibCompressor(encoding=encoding, strategy=strategy) + + def _write(self, chunk: Union[bytes, bytearray, memoryview]) -> None: + size = len(chunk) + self.buffer_size += size + self.output_size += size + transport = self._protocol.transport + if transport is None or transport.is_closing(): + raise ClientConnectionResetError("Cannot write to closing transport") + transport.write(chunk) + + def _writelines(self, chunks: Iterable[bytes]) -> None: + size = 0 + for chunk in chunks: + size += len(chunk) + self.buffer_size += size + self.output_size += size + transport = self._protocol.transport + if transport is None or transport.is_closing(): + raise ClientConnectionResetError("Cannot write to closing transport") + if SKIP_WRITELINES or size < MIN_PAYLOAD_FOR_WRITELINES: + transport.write(b"".join(chunks)) + else: + transport.writelines(chunks) + + def _write_chunked_payload( + self, chunk: Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"] + ) -> None: + """Write a chunk with proper chunked encoding.""" + chunk_len_pre = f"{len(chunk):x}\r\n".encode("ascii") + self._writelines((chunk_len_pre, chunk, b"\r\n")) + + def _send_headers_with_payload( + self, + chunk: Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"], + is_eof: bool, + ) -> None: + """Send buffered headers with payload, coalescing into single write.""" + # Mark headers as written + self._headers_written = True + headers_buf = self._headers_buf + self._headers_buf = None + + if TYPE_CHECKING: + # Safe because callers (write() and write_eof()) only invoke this method + # after checking that self._headers_buf is truthy + assert headers_buf is not None + + if not self.chunked: + # Non-chunked: coalesce headers with body + if chunk: + self._writelines((headers_buf, chunk)) + else: + self._write(headers_buf) + return + + # Coalesce headers with chunked data + if chunk: + chunk_len_pre = f"{len(chunk):x}\r\n".encode("ascii") + if is_eof: + self._writelines((headers_buf, chunk_len_pre, chunk, b"\r\n0\r\n\r\n")) + else: + self._writelines((headers_buf, chunk_len_pre, chunk, b"\r\n")) + elif is_eof: + self._writelines((headers_buf, b"0\r\n\r\n")) + else: + self._write(headers_buf) + + async def write( + self, + chunk: Union[bytes, bytearray, memoryview], + *, + drain: bool = True, + LIMIT: int = 0x10000, + ) -> None: + """ + Writes chunk of data to a stream. + + write_eof() indicates end of stream. + writer can't be used after write_eof() method being called. + write() return drain future. + """ + if self._on_chunk_sent is not None: + await self._on_chunk_sent(chunk) + + if isinstance(chunk, memoryview): + if chunk.nbytes != len(chunk): + # just reshape it + chunk = chunk.cast("c") + + if self._compress is not None: + chunk = await self._compress.compress(chunk) + if not chunk: + return + + if self.length is not None: + chunk_len = len(chunk) + if self.length >= chunk_len: + self.length = self.length - chunk_len + else: + chunk = chunk[: self.length] + self.length = 0 + if not chunk: + return + + # Handle buffered headers for small payload optimization + if self._headers_buf and not self._headers_written: + self._send_headers_with_payload(chunk, False) + if drain and self.buffer_size > LIMIT: + self.buffer_size = 0 + await self.drain() + return + + if chunk: + if self.chunked: + self._write_chunked_payload(chunk) + else: + self._write(chunk) + + if drain and self.buffer_size > LIMIT: + self.buffer_size = 0 + await self.drain() + + async def write_headers( + self, status_line: str, headers: "CIMultiDict[str]" + ) -> None: + """Write headers to the stream.""" + if self._on_headers_sent is not None: + await self._on_headers_sent(headers) + # status + headers + buf = _serialize_headers(status_line, headers) + self._headers_written = False + self._headers_buf = buf + + def send_headers(self) -> None: + """Force sending buffered headers if not already sent.""" + if not self._headers_buf or self._headers_written: + return + + self._headers_written = True + headers_buf = self._headers_buf + self._headers_buf = None + + if TYPE_CHECKING: + # Safe because we only enter this block when self._headers_buf is truthy + assert headers_buf is not None + + self._write(headers_buf) + + def set_eof(self) -> None: + """Indicate that the message is complete.""" + if self._eof: + return + + # If headers haven't been sent yet, send them now + # This handles the case where there's no body at all + if self._headers_buf and not self._headers_written: + self._headers_written = True + headers_buf = self._headers_buf + self._headers_buf = None + + if TYPE_CHECKING: + # Safe because we only enter this block when self._headers_buf is truthy + assert headers_buf is not None + + # Combine headers and chunked EOF marker in a single write + if self.chunked: + self._writelines((headers_buf, b"0\r\n\r\n")) + else: + self._write(headers_buf) + elif self.chunked and self._headers_written: + # Headers already sent, just send the final chunk marker + self._write(b"0\r\n\r\n") + + self._eof = True + + async def write_eof(self, chunk: bytes = b"") -> None: + if self._eof: + return + + if chunk and self._on_chunk_sent is not None: + await self._on_chunk_sent(chunk) + + # Handle body/compression + if self._compress: + chunks: List[bytes] = [] + chunks_len = 0 + if chunk and (compressed_chunk := await self._compress.compress(chunk)): + chunks_len = len(compressed_chunk) + chunks.append(compressed_chunk) + + flush_chunk = self._compress.flush() + chunks_len += len(flush_chunk) + chunks.append(flush_chunk) + assert chunks_len + + # Send buffered headers with compressed data if not yet sent + if self._headers_buf and not self._headers_written: + self._headers_written = True + headers_buf = self._headers_buf + self._headers_buf = None + + if self.chunked: + # Coalesce headers with compressed chunked data + chunk_len_pre = f"{chunks_len:x}\r\n".encode("ascii") + self._writelines( + (headers_buf, chunk_len_pre, *chunks, b"\r\n0\r\n\r\n") + ) + else: + # Coalesce headers with compressed data + self._writelines((headers_buf, *chunks)) + await self.drain() + self._eof = True + return + + # Headers already sent, just write compressed data + if self.chunked: + chunk_len_pre = f"{chunks_len:x}\r\n".encode("ascii") + self._writelines((chunk_len_pre, *chunks, b"\r\n0\r\n\r\n")) + elif len(chunks) > 1: + self._writelines(chunks) + else: + self._write(chunks[0]) + await self.drain() + self._eof = True + return + + # No compression - send buffered headers if not yet sent + if self._headers_buf and not self._headers_written: + # Use helper to send headers with payload + self._send_headers_with_payload(chunk, True) + await self.drain() + self._eof = True + return + + # Handle remaining body + if self.chunked: + if chunk: + # Write final chunk with EOF marker + self._writelines( + (f"{len(chunk):x}\r\n".encode("ascii"), chunk, b"\r\n0\r\n\r\n") + ) + else: + self._write(b"0\r\n\r\n") + await self.drain() + self._eof = True + return + + if chunk: + self._write(chunk) + await self.drain() + + self._eof = True + + async def drain(self) -> None: + """Flush the write buffer. + + The intended use is to write + + await w.write(data) + await w.drain() + """ + protocol = self._protocol + if protocol.transport is not None and protocol._paused: + await protocol._drain_helper() + + +def _safe_header(string: str) -> str: + if "\r" in string or "\n" in string: + raise ValueError( + "Newline or carriage return detected in headers. " + "Potential header injection attack." + ) + return string + + +def _py_serialize_headers(status_line: str, headers: "CIMultiDict[str]") -> bytes: + headers_gen = (_safe_header(k) + ": " + _safe_header(v) for k, v in headers.items()) + line = status_line + "\r\n" + "\r\n".join(headers_gen) + "\r\n\r\n" + return line.encode("utf-8") + + +_serialize_headers = _py_serialize_headers + +try: + import aiohttp._http_writer as _http_writer # type: ignore[import-not-found] + + _c_serialize_headers = _http_writer._serialize_headers + if not NO_EXTENSIONS: + _serialize_headers = _c_serialize_headers +except ImportError: + pass diff --git a/lib/python3.12/site-packages/aiohttp/log.py b/lib/python3.12/site-packages/aiohttp/log.py new file mode 100644 index 0000000000000000000000000000000000000000..3cecea2bac185df741bccd0a32a5fef9cfe23299 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/log.py @@ -0,0 +1,8 @@ +import logging + +access_logger = logging.getLogger("aiohttp.access") +client_logger = logging.getLogger("aiohttp.client") +internal_logger = logging.getLogger("aiohttp.internal") +server_logger = logging.getLogger("aiohttp.server") +web_logger = logging.getLogger("aiohttp.web") +ws_logger = logging.getLogger("aiohttp.websocket") diff --git a/lib/python3.12/site-packages/aiohttp/multipart.py b/lib/python3.12/site-packages/aiohttp/multipart.py new file mode 100644 index 0000000000000000000000000000000000000000..9c37f0bb71699d9001f4baf0fcd64bbeb94410ef --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/multipart.py @@ -0,0 +1,1152 @@ +import base64 +import binascii +import json +import re +import sys +import uuid +import warnings +from collections import deque +from collections.abc import Mapping, Sequence +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Deque, + Dict, + Iterator, + List, + Optional, + Tuple, + Type, + Union, + cast, +) +from urllib.parse import parse_qsl, unquote, urlencode + +from multidict import CIMultiDict, CIMultiDictProxy + +from .abc import AbstractStreamWriter +from .compression_utils import ( + DEFAULT_MAX_DECOMPRESS_SIZE, + ZLibCompressor, + ZLibDecompressor, +) +from .hdrs import ( + CONTENT_DISPOSITION, + CONTENT_ENCODING, + CONTENT_LENGTH, + CONTENT_TRANSFER_ENCODING, + CONTENT_TYPE, +) +from .helpers import CHAR, TOKEN, parse_mimetype, reify +from .http import HeadersParser +from .log import internal_logger +from .payload import ( + JsonPayload, + LookupError, + Order, + Payload, + StringPayload, + get_payload, + payload_type, +) +from .streams import StreamReader + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing import TypeVar + + Self = TypeVar("Self", bound="BodyPartReader") + +__all__ = ( + "MultipartReader", + "MultipartWriter", + "BodyPartReader", + "BadContentDispositionHeader", + "BadContentDispositionParam", + "parse_content_disposition", + "content_disposition_filename", +) + + +if TYPE_CHECKING: + from .client_reqrep import ClientResponse + + +class BadContentDispositionHeader(RuntimeWarning): + pass + + +class BadContentDispositionParam(RuntimeWarning): + pass + + +def parse_content_disposition( + header: Optional[str], +) -> Tuple[Optional[str], Dict[str, str]]: + def is_token(string: str) -> bool: + return bool(string) and TOKEN >= set(string) + + def is_quoted(string: str) -> bool: + return string[0] == string[-1] == '"' + + def is_rfc5987(string: str) -> bool: + return is_token(string) and string.count("'") == 2 + + def is_extended_param(string: str) -> bool: + return string.endswith("*") + + def is_continuous_param(string: str) -> bool: + pos = string.find("*") + 1 + if not pos: + return False + substring = string[pos:-1] if string.endswith("*") else string[pos:] + return substring.isdigit() + + def unescape(text: str, *, chars: str = "".join(map(re.escape, CHAR))) -> str: + return re.sub(f"\\\\([{chars}])", "\\1", text) + + if not header: + return None, {} + + disptype, *parts = header.split(";") + if not is_token(disptype): + warnings.warn(BadContentDispositionHeader(header)) + return None, {} + + params: Dict[str, str] = {} + while parts: + item = parts.pop(0) + + if not item: # To handle trailing semicolons + warnings.warn(BadContentDispositionHeader(header)) + continue + + if "=" not in item: + warnings.warn(BadContentDispositionHeader(header)) + return None, {} + + key, value = item.split("=", 1) + key = key.lower().strip() + value = value.lstrip() + + if key in params: + warnings.warn(BadContentDispositionHeader(header)) + return None, {} + + if not is_token(key): + warnings.warn(BadContentDispositionParam(item)) + continue + + elif is_continuous_param(key): + if is_quoted(value): + value = unescape(value[1:-1]) + elif not is_token(value): + warnings.warn(BadContentDispositionParam(item)) + continue + + elif is_extended_param(key): + if is_rfc5987(value): + encoding, _, value = value.split("'", 2) + encoding = encoding or "utf-8" + else: + warnings.warn(BadContentDispositionParam(item)) + continue + + try: + value = unquote(value, encoding, "strict") + except UnicodeDecodeError: # pragma: nocover + warnings.warn(BadContentDispositionParam(item)) + continue + + else: + failed = True + if is_quoted(value): + failed = False + value = unescape(value[1:-1].lstrip("\\/")) + elif is_token(value): + failed = False + elif parts: + # maybe just ; in filename, in any case this is just + # one case fix, for proper fix we need to redesign parser + _value = f"{value};{parts[0]}" + if is_quoted(_value): + parts.pop(0) + value = unescape(_value[1:-1].lstrip("\\/")) + failed = False + + if failed: + warnings.warn(BadContentDispositionHeader(header)) + return None, {} + + params[key] = value + + return disptype.lower(), params + + +def content_disposition_filename( + params: Mapping[str, str], name: str = "filename" +) -> Optional[str]: + name_suf = "%s*" % name + if not params: + return None + elif name_suf in params: + return params[name_suf] + elif name in params: + return params[name] + else: + parts = [] + fnparams = sorted( + (key, value) for key, value in params.items() if key.startswith(name_suf) + ) + for num, (key, value) in enumerate(fnparams): + _, tail = key.split("*", 1) + if tail.endswith("*"): + tail = tail[:-1] + if tail == str(num): + parts.append(value) + else: + break + if not parts: + return None + value = "".join(parts) + if "'" in value: + encoding, _, value = value.split("'", 2) + encoding = encoding or "utf-8" + return unquote(value, encoding, "strict") + return value + + +class MultipartResponseWrapper: + """Wrapper around the MultipartReader. + + It takes care about + underlying connection and close it when it needs in. + """ + + def __init__( + self, + resp: "ClientResponse", + stream: "MultipartReader", + ) -> None: + self.resp = resp + self.stream = stream + + def __aiter__(self) -> "MultipartResponseWrapper": + return self + + async def __anext__( + self, + ) -> Union["MultipartReader", "BodyPartReader"]: + part = await self.next() + if part is None: + raise StopAsyncIteration + return part + + def at_eof(self) -> bool: + """Returns True when all response data had been read.""" + return self.resp.content.at_eof() + + async def next( + self, + ) -> Optional[Union["MultipartReader", "BodyPartReader"]]: + """Emits next multipart reader object.""" + item = await self.stream.next() + if self.stream.at_eof(): + await self.release() + return item + + async def release(self) -> None: + """Release the connection gracefully. + + All remaining content is read to the void. + """ + await self.resp.release() + + +class BodyPartReader: + """Multipart reader for single body part.""" + + chunk_size = 8192 + + def __init__( + self, + boundary: bytes, + headers: "CIMultiDictProxy[str]", + content: StreamReader, + *, + subtype: str = "mixed", + default_charset: Optional[str] = None, + max_decompress_size: int = DEFAULT_MAX_DECOMPRESS_SIZE, + ) -> None: + self.headers = headers + self._boundary = boundary + self._boundary_len = len(boundary) + 2 # Boundary + \r\n + self._content = content + self._default_charset = default_charset + self._at_eof = False + self._is_form_data = subtype == "form-data" + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.8 + length = None if self._is_form_data else self.headers.get(CONTENT_LENGTH, None) + self._length = int(length) if length is not None else None + self._read_bytes = 0 + self._unread: Deque[bytes] = deque() + self._prev_chunk: Optional[bytes] = None + self._content_eof = 0 + self._cache: Dict[str, Any] = {} + self._max_decompress_size = max_decompress_size + + def __aiter__(self: Self) -> Self: + return self + + async def __anext__(self) -> bytes: + part = await self.next() + if part is None: + raise StopAsyncIteration + return part + + async def next(self) -> Optional[bytes]: + item = await self.read() + if not item: + return None + return item + + async def read(self, *, decode: bool = False) -> bytes: + """Reads body part data. + + decode: Decodes data following by encoding + method from Content-Encoding header. If it missed + data remains untouched + """ + if self._at_eof: + return b"" + data = bytearray() + while not self._at_eof: + data.extend(await self.read_chunk(self.chunk_size)) + if decode: + return await self.decode(data) + return data + + async def read_chunk(self, size: int = chunk_size) -> bytes: + """Reads body part content chunk of the specified size. + + size: chunk size + """ + if self._at_eof: + return b"" + if self._length: + chunk = await self._read_chunk_from_length(size) + else: + chunk = await self._read_chunk_from_stream(size) + + # For the case of base64 data, we must read a fragment of size with a + # remainder of 0 by dividing by 4 for string without symbols \n or \r + encoding = self.headers.get(CONTENT_TRANSFER_ENCODING) + if encoding and encoding.lower() == "base64": + stripped_chunk = b"".join(chunk.split()) + remainder = len(stripped_chunk) % 4 + + while remainder != 0 and not self.at_eof(): + over_chunk_size = 4 - remainder + over_chunk = b"" + + if self._prev_chunk: + over_chunk = self._prev_chunk[:over_chunk_size] + self._prev_chunk = self._prev_chunk[len(over_chunk) :] + + if len(over_chunk) != over_chunk_size: + over_chunk += await self._content.read(4 - len(over_chunk)) + + if not over_chunk: + self._at_eof = True + + stripped_chunk += b"".join(over_chunk.split()) + chunk += over_chunk + remainder = len(stripped_chunk) % 4 + + self._read_bytes += len(chunk) + if self._read_bytes == self._length: + self._at_eof = True + if self._at_eof and await self._content.readline() != b"\r\n": + raise ValueError("Reader did not read all the data or it is malformed") + return chunk + + async def _read_chunk_from_length(self, size: int) -> bytes: + # Reads body part content chunk of the specified size. + # The body part must has Content-Length header with proper value. + assert self._length is not None, "Content-Length required for chunked read" + chunk_size = min(size, self._length - self._read_bytes) + chunk = await self._content.read(chunk_size) + if self._content.at_eof(): + self._at_eof = True + return chunk + + async def _read_chunk_from_stream(self, size: int) -> bytes: + # Reads content chunk of body part with unknown length. + # The Content-Length header for body part is not necessary. + assert ( + size >= self._boundary_len + ), "Chunk size must be greater or equal than boundary length + 2" + first_chunk = self._prev_chunk is None + if first_chunk: + # We need to re-add the CRLF that got removed from headers parsing. + self._prev_chunk = b"\r\n" + await self._content.read(size) + + chunk = b"" + # content.read() may return less than size, so we need to loop to ensure + # we have enough data to detect the boundary. + while len(chunk) < self._boundary_len: + chunk += await self._content.read(size) + self._content_eof += int(self._content.at_eof()) + if self._content_eof > 2: + raise ValueError("Reading after EOF") + if self._content_eof: + break + if len(chunk) > size: + self._content.unread_data(chunk[size:]) + chunk = chunk[:size] + + assert self._prev_chunk is not None + window = self._prev_chunk + chunk + sub = b"\r\n" + self._boundary + if first_chunk: + idx = window.find(sub) + else: + idx = window.find(sub, max(0, len(self._prev_chunk) - len(sub))) + if idx >= 0: + # pushing boundary back to content + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=DeprecationWarning) + self._content.unread_data(window[idx:]) + self._prev_chunk = self._prev_chunk[:idx] + chunk = window[len(self._prev_chunk) : idx] + if not chunk: + self._at_eof = True + result = self._prev_chunk[2 if first_chunk else 0 :] # Strip initial CRLF + self._prev_chunk = chunk + return result + + async def readline(self) -> bytes: + """Reads body part by line by line.""" + if self._at_eof: + return b"" + + if self._unread: + line = self._unread.popleft() + else: + line = await self._content.readline() + + if line.startswith(self._boundary): + # the very last boundary may not come with \r\n, + # so set single rules for everyone + sline = line.rstrip(b"\r\n") + boundary = self._boundary + last_boundary = self._boundary + b"--" + # ensure that we read exactly the boundary, not something alike + if sline == boundary or sline == last_boundary: + self._at_eof = True + self._unread.append(line) + return b"" + else: + next_line = await self._content.readline() + if next_line.startswith(self._boundary): + line = line[:-2] # strip CRLF but only once + self._unread.append(next_line) + + return line + + async def release(self) -> None: + """Like read(), but reads all the data to the void.""" + if self._at_eof: + return + while not self._at_eof: + await self.read_chunk(self.chunk_size) + + async def text(self, *, encoding: Optional[str] = None) -> str: + """Like read(), but assumes that body part contains text data.""" + data = await self.read(decode=True) + # see https://www.w3.org/TR/html5/forms.html#multipart/form-data-encoding-algorithm + # and https://dvcs.w3.org/hg/xhr/raw-file/tip/Overview.html#dom-xmlhttprequest-send + encoding = encoding or self.get_charset(default="utf-8") + return data.decode(encoding) + + async def json(self, *, encoding: Optional[str] = None) -> Optional[Dict[str, Any]]: + """Like read(), but assumes that body parts contains JSON data.""" + data = await self.read(decode=True) + if not data: + return None + encoding = encoding or self.get_charset(default="utf-8") + return cast(Dict[str, Any], json.loads(data.decode(encoding))) + + async def form(self, *, encoding: Optional[str] = None) -> List[Tuple[str, str]]: + """Like read(), but assumes that body parts contain form urlencoded data.""" + data = await self.read(decode=True) + if not data: + return [] + if encoding is not None: + real_encoding = encoding + else: + real_encoding = self.get_charset(default="utf-8") + try: + decoded_data = data.rstrip().decode(real_encoding) + except UnicodeDecodeError: + raise ValueError("data cannot be decoded with %s encoding" % real_encoding) + + return parse_qsl( + decoded_data, + keep_blank_values=True, + encoding=real_encoding, + ) + + def at_eof(self) -> bool: + """Returns True if the boundary was reached or False otherwise.""" + return self._at_eof + + async def decode(self, data: bytes) -> bytes: + """Decodes data. + + Decoding is done according the specified Content-Encoding + or Content-Transfer-Encoding headers value. + """ + if CONTENT_TRANSFER_ENCODING in self.headers: + data = self._decode_content_transfer(data) + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.8 + if not self._is_form_data and CONTENT_ENCODING in self.headers: + return await self._decode_content(data) + return data + + async def _decode_content(self, data: bytes) -> bytes: + encoding = self.headers.get(CONTENT_ENCODING, "").lower() + if encoding == "identity": + return data + if encoding in {"deflate", "gzip"}: + return await ZLibDecompressor( + encoding=encoding, + suppress_deflate_header=True, + ).decompress(data, max_length=self._max_decompress_size) + + raise RuntimeError(f"unknown content encoding: {encoding}") + + def _decode_content_transfer(self, data: bytes) -> bytes: + encoding = self.headers.get(CONTENT_TRANSFER_ENCODING, "").lower() + + if encoding == "base64": + return base64.b64decode(data) + elif encoding == "quoted-printable": + return binascii.a2b_qp(data) + elif encoding in ("binary", "8bit", "7bit"): + return data + else: + raise RuntimeError(f"unknown content transfer encoding: {encoding}") + + def get_charset(self, default: str) -> str: + """Returns charset parameter from Content-Type header or default.""" + ctype = self.headers.get(CONTENT_TYPE, "") + mimetype = parse_mimetype(ctype) + return mimetype.parameters.get("charset", self._default_charset or default) + + @reify + def name(self) -> Optional[str]: + """Returns name specified in Content-Disposition header. + + If the header is missing or malformed, returns None. + """ + _, params = parse_content_disposition(self.headers.get(CONTENT_DISPOSITION)) + return content_disposition_filename(params, "name") + + @reify + def filename(self) -> Optional[str]: + """Returns filename specified in Content-Disposition header. + + Returns None if the header is missing or malformed. + """ + _, params = parse_content_disposition(self.headers.get(CONTENT_DISPOSITION)) + return content_disposition_filename(params, "filename") + + +@payload_type(BodyPartReader, order=Order.try_first) +class BodyPartReaderPayload(Payload): + _value: BodyPartReader + # _autoclose = False (inherited) - Streaming reader that may have resources + + def __init__(self, value: BodyPartReader, *args: Any, **kwargs: Any) -> None: + super().__init__(value, *args, **kwargs) + + params: Dict[str, str] = {} + if value.name is not None: + params["name"] = value.name + if value.filename is not None: + params["filename"] = value.filename + + if params: + self.set_content_disposition("attachment", True, **params) + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + raise TypeError("Unable to decode.") + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """Raises TypeError as body parts should be consumed via write(). + + This is intentional: BodyPartReader payloads are designed for streaming + large data (potentially gigabytes) and must be consumed only once via + the write() method to avoid memory exhaustion. They cannot be buffered + in memory for reuse. + """ + raise TypeError("Unable to read body part as bytes. Use write() to consume.") + + async def write(self, writer: AbstractStreamWriter) -> None: + field = self._value + chunk = await field.read_chunk(size=2**16) + while chunk: + await writer.write(await field.decode(chunk)) + chunk = await field.read_chunk(size=2**16) + + +class MultipartReader: + """Multipart body reader.""" + + #: Response wrapper, used when multipart readers constructs from response. + response_wrapper_cls = MultipartResponseWrapper + #: Multipart reader class, used to handle multipart/* body parts. + #: None points to type(self) + multipart_reader_cls: Optional[Type["MultipartReader"]] = None + #: Body part reader class for non multipart/* content types. + part_reader_cls = BodyPartReader + + def __init__(self, headers: Mapping[str, str], content: StreamReader) -> None: + self._mimetype = parse_mimetype(headers[CONTENT_TYPE]) + assert self._mimetype.type == "multipart", "multipart/* content type expected" + if "boundary" not in self._mimetype.parameters: + raise ValueError( + "boundary missed for Content-Type: %s" % headers[CONTENT_TYPE] + ) + + self.headers = headers + self._boundary = ("--" + self._get_boundary()).encode() + self._content = content + self._default_charset: Optional[str] = None + self._last_part: Optional[Union["MultipartReader", BodyPartReader]] = None + self._at_eof = False + self._at_bof = True + self._unread: List[bytes] = [] + + def __aiter__(self: Self) -> Self: + return self + + async def __anext__( + self, + ) -> Optional[Union["MultipartReader", BodyPartReader]]: + part = await self.next() + if part is None: + raise StopAsyncIteration + return part + + @classmethod + def from_response( + cls, + response: "ClientResponse", + ) -> MultipartResponseWrapper: + """Constructs reader instance from HTTP response. + + :param response: :class:`~aiohttp.client.ClientResponse` instance + """ + obj = cls.response_wrapper_cls( + response, cls(response.headers, response.content) + ) + return obj + + def at_eof(self) -> bool: + """Returns True if the final boundary was reached, false otherwise.""" + return self._at_eof + + async def next( + self, + ) -> Optional[Union["MultipartReader", BodyPartReader]]: + """Emits the next multipart body part.""" + # So, if we're at BOF, we need to skip till the boundary. + if self._at_eof: + return None + await self._maybe_release_last_part() + if self._at_bof: + await self._read_until_first_boundary() + self._at_bof = False + else: + await self._read_boundary() + if self._at_eof: # we just read the last boundary, nothing to do there + return None + + part = await self.fetch_next_part() + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.6 + if ( + self._last_part is None + and self._mimetype.subtype == "form-data" + and isinstance(part, BodyPartReader) + ): + _, params = parse_content_disposition(part.headers.get(CONTENT_DISPOSITION)) + if params.get("name") == "_charset_": + # Longest encoding in https://encoding.spec.whatwg.org/encodings.json + # is 19 characters, so 32 should be more than enough for any valid encoding. + charset = await part.read_chunk(32) + if len(charset) > 31: + raise RuntimeError("Invalid default charset") + self._default_charset = charset.strip().decode() + part = await self.fetch_next_part() + self._last_part = part + return self._last_part + + async def release(self) -> None: + """Reads all the body parts to the void till the final boundary.""" + while not self._at_eof: + item = await self.next() + if item is None: + break + await item.release() + + async def fetch_next_part( + self, + ) -> Union["MultipartReader", BodyPartReader]: + """Returns the next body part reader.""" + headers = await self._read_headers() + return self._get_part_reader(headers) + + def _get_part_reader( + self, + headers: "CIMultiDictProxy[str]", + ) -> Union["MultipartReader", BodyPartReader]: + """Dispatches the response by the `Content-Type` header. + + Returns a suitable reader instance. + + :param dict headers: Response headers + """ + ctype = headers.get(CONTENT_TYPE, "") + mimetype = parse_mimetype(ctype) + + if mimetype.type == "multipart": + if self.multipart_reader_cls is None: + return type(self)(headers, self._content) + return self.multipart_reader_cls(headers, self._content) + else: + return self.part_reader_cls( + self._boundary, + headers, + self._content, + subtype=self._mimetype.subtype, + default_charset=self._default_charset, + ) + + def _get_boundary(self) -> str: + boundary = self._mimetype.parameters["boundary"] + if len(boundary) > 70: + raise ValueError("boundary %r is too long (70 chars max)" % boundary) + + return boundary + + async def _readline(self) -> bytes: + if self._unread: + return self._unread.pop() + return await self._content.readline() + + async def _read_until_first_boundary(self) -> None: + while True: + chunk = await self._readline() + if chunk == b"": + raise ValueError( + "Could not find starting boundary %r" % (self._boundary) + ) + chunk = chunk.rstrip() + if chunk == self._boundary: + return + elif chunk == self._boundary + b"--": + self._at_eof = True + return + + async def _read_boundary(self) -> None: + chunk = (await self._readline()).rstrip() + if chunk == self._boundary: + pass + elif chunk == self._boundary + b"--": + self._at_eof = True + epilogue = await self._readline() + next_line = await self._readline() + + # the epilogue is expected and then either the end of input or the + # parent multipart boundary, if the parent boundary is found then + # it should be marked as unread and handed to the parent for + # processing + if next_line[:2] == b"--": + self._unread.append(next_line) + # otherwise the request is likely missing an epilogue and both + # lines should be passed to the parent for processing + # (this handles the old behavior gracefully) + else: + self._unread.extend([next_line, epilogue]) + else: + raise ValueError(f"Invalid boundary {chunk!r}, expected {self._boundary!r}") + + async def _read_headers(self) -> "CIMultiDictProxy[str]": + lines = [] + while True: + chunk = await self._content.readline() + chunk = chunk.rstrip(b"\r\n") + lines.append(chunk) + if not chunk: + break + parser = HeadersParser() + headers, raw_headers = parser.parse_headers(lines) + return headers + + async def _maybe_release_last_part(self) -> None: + """Ensures that the last read body part is read completely.""" + if self._last_part is not None: + if not self._last_part.at_eof(): + await self._last_part.release() + self._unread.extend(self._last_part._unread) + self._last_part = None + + +_Part = Tuple[Payload, str, str] + + +class MultipartWriter(Payload): + """Multipart body writer.""" + + _value: None + # _consumed = False (inherited) - Can be encoded multiple times + _autoclose = True # No file handles, just collects parts in memory + + def __init__(self, subtype: str = "mixed", boundary: Optional[str] = None) -> None: + boundary = boundary if boundary is not None else uuid.uuid4().hex + # The underlying Payload API demands a str (utf-8), not bytes, + # so we need to ensure we don't lose anything during conversion. + # As a result, require the boundary to be ASCII only. + # In both situations. + + try: + self._boundary = boundary.encode("ascii") + except UnicodeEncodeError: + raise ValueError("boundary should contain ASCII only chars") from None + ctype = f"multipart/{subtype}; boundary={self._boundary_value}" + + super().__init__(None, content_type=ctype) + + self._parts: List[_Part] = [] + self._is_form_data = subtype == "form-data" + + def __enter__(self) -> "MultipartWriter": + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + pass + + def __iter__(self) -> Iterator[_Part]: + return iter(self._parts) + + def __len__(self) -> int: + return len(self._parts) + + def __bool__(self) -> bool: + return True + + _valid_tchar_regex = re.compile(rb"\A[!#$%&'*+\-.^_`|~\w]+\Z") + _invalid_qdtext_char_regex = re.compile(rb"[\x00-\x08\x0A-\x1F\x7F]") + + @property + def _boundary_value(self) -> str: + """Wrap boundary parameter value in quotes, if necessary. + + Reads self.boundary and returns a unicode string. + """ + # Refer to RFCs 7231, 7230, 5234. + # + # parameter = token "=" ( token / quoted-string ) + # token = 1*tchar + # quoted-string = DQUOTE *( qdtext / quoted-pair ) DQUOTE + # qdtext = HTAB / SP / %x21 / %x23-5B / %x5D-7E / obs-text + # obs-text = %x80-FF + # quoted-pair = "\" ( HTAB / SP / VCHAR / obs-text ) + # tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" + # / "+" / "-" / "." / "^" / "_" / "`" / "|" / "~" + # / DIGIT / ALPHA + # ; any VCHAR, except delimiters + # VCHAR = %x21-7E + value = self._boundary + if re.match(self._valid_tchar_regex, value): + return value.decode("ascii") # cannot fail + + if re.search(self._invalid_qdtext_char_regex, value): + raise ValueError("boundary value contains invalid characters") + + # escape %x5C and %x22 + quoted_value_content = value.replace(b"\\", b"\\\\") + quoted_value_content = quoted_value_content.replace(b'"', b'\\"') + + return '"' + quoted_value_content.decode("ascii") + '"' + + @property + def boundary(self) -> str: + return self._boundary.decode("ascii") + + def append(self, obj: Any, headers: Optional[Mapping[str, str]] = None) -> Payload: + if headers is None: + headers = CIMultiDict() + + if isinstance(obj, Payload): + obj.headers.update(headers) + return self.append_payload(obj) + else: + try: + payload = get_payload(obj, headers=headers) + except LookupError: + raise TypeError("Cannot create payload from %r" % obj) + else: + return self.append_payload(payload) + + def append_payload(self, payload: Payload) -> Payload: + """Adds a new body part to multipart writer.""" + encoding: Optional[str] = None + te_encoding: Optional[str] = None + if self._is_form_data: + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.7 + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.8 + assert ( + not {CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TRANSFER_ENCODING} + & payload.headers.keys() + ) + # Set default Content-Disposition in case user doesn't create one + if CONTENT_DISPOSITION not in payload.headers: + name = f"section-{len(self._parts)}" + payload.set_content_disposition("form-data", name=name) + else: + # compression + encoding = payload.headers.get(CONTENT_ENCODING, "").lower() + if encoding and encoding not in ("deflate", "gzip", "identity"): + raise RuntimeError(f"unknown content encoding: {encoding}") + if encoding == "identity": + encoding = None + + # te encoding + te_encoding = payload.headers.get(CONTENT_TRANSFER_ENCODING, "").lower() + if te_encoding not in ("", "base64", "quoted-printable", "binary"): + raise RuntimeError(f"unknown content transfer encoding: {te_encoding}") + if te_encoding == "binary": + te_encoding = None + + # size + size = payload.size + if size is not None and not (encoding or te_encoding): + payload.headers[CONTENT_LENGTH] = str(size) + + self._parts.append((payload, encoding, te_encoding)) # type: ignore[arg-type] + return payload + + def append_json( + self, obj: Any, headers: Optional[Mapping[str, str]] = None + ) -> Payload: + """Helper to append JSON part.""" + if headers is None: + headers = CIMultiDict() + + return self.append_payload(JsonPayload(obj, headers=headers)) + + def append_form( + self, + obj: Union[Sequence[Tuple[str, str]], Mapping[str, str]], + headers: Optional[Mapping[str, str]] = None, + ) -> Payload: + """Helper to append form urlencoded part.""" + assert isinstance(obj, (Sequence, Mapping)) + + if headers is None: + headers = CIMultiDict() + + if isinstance(obj, Mapping): + obj = list(obj.items()) + data = urlencode(obj, doseq=True) + + return self.append_payload( + StringPayload( + data, headers=headers, content_type="application/x-www-form-urlencoded" + ) + ) + + @property + def size(self) -> Optional[int]: + """Size of the payload.""" + total = 0 + for part, encoding, te_encoding in self._parts: + part_size = part.size + if encoding or te_encoding or part_size is None: + return None + + total += int( + 2 + + len(self._boundary) + + 2 + + part_size # b'--'+self._boundary+b'\r\n' + + len(part._binary_headers) + + 2 # b'\r\n' + ) + + total += 2 + len(self._boundary) + 4 # b'--'+self._boundary+b'--\r\n' + return total + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + """Return string representation of the multipart data. + + WARNING: This method may do blocking I/O if parts contain file payloads. + It should not be called in the event loop. Use as_bytes().decode() instead. + """ + return "".join( + "--" + + self.boundary + + "\r\n" + + part._binary_headers.decode(encoding, errors) + + part.decode() + for part, _e, _te in self._parts + ) + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """Return bytes representation of the multipart data. + + This method is async-safe and calls as_bytes on underlying payloads. + """ + parts: List[bytes] = [] + + # Process each part + for part, _e, _te in self._parts: + # Add boundary + parts.append(b"--" + self._boundary + b"\r\n") + + # Add headers + parts.append(part._binary_headers) + + # Add payload content using as_bytes for async safety + part_bytes = await part.as_bytes(encoding, errors) + parts.append(part_bytes) + + # Add trailing CRLF + parts.append(b"\r\n") + + # Add closing boundary + parts.append(b"--" + self._boundary + b"--\r\n") + + return b"".join(parts) + + async def write( + self, writer: AbstractStreamWriter, close_boundary: bool = True + ) -> None: + """Write body.""" + for part, encoding, te_encoding in self._parts: + if self._is_form_data: + # https://datatracker.ietf.org/doc/html/rfc7578#section-4.2 + assert CONTENT_DISPOSITION in part.headers + assert "name=" in part.headers[CONTENT_DISPOSITION] + + await writer.write(b"--" + self._boundary + b"\r\n") + await writer.write(part._binary_headers) + + if encoding or te_encoding: + w = MultipartPayloadWriter(writer) + if encoding: + w.enable_compression(encoding) + if te_encoding: + w.enable_encoding(te_encoding) + await part.write(w) # type: ignore[arg-type] + await w.write_eof() + else: + await part.write(writer) + + await writer.write(b"\r\n") + + if close_boundary: + await writer.write(b"--" + self._boundary + b"--\r\n") + + async def close(self) -> None: + """ + Close all part payloads that need explicit closing. + + IMPORTANT: This method must not await anything that might not finish + immediately, as it may be called during cleanup/cancellation. Schedule + any long-running operations without awaiting them. + """ + if self._consumed: + return + self._consumed = True + + # Close all parts that need explicit closing + # We catch and log exceptions to ensure all parts get a chance to close + # we do not use asyncio.gather() here because we are not allowed + # to suspend given we may be called during cleanup + for idx, (part, _, _) in enumerate(self._parts): + if not part.autoclose and not part.consumed: + try: + await part.close() + except Exception as exc: + internal_logger.error( + "Failed to close multipart part %d: %s", idx, exc, exc_info=True + ) + + +class MultipartPayloadWriter: + def __init__(self, writer: AbstractStreamWriter) -> None: + self._writer = writer + self._encoding: Optional[str] = None + self._compress: Optional[ZLibCompressor] = None + self._encoding_buffer: Optional[bytearray] = None + + def enable_encoding(self, encoding: str) -> None: + if encoding == "base64": + self._encoding = encoding + self._encoding_buffer = bytearray() + elif encoding == "quoted-printable": + self._encoding = "quoted-printable" + + def enable_compression( + self, encoding: str = "deflate", strategy: Optional[int] = None + ) -> None: + self._compress = ZLibCompressor( + encoding=encoding, + suppress_deflate_header=True, + strategy=strategy, + ) + + async def write_eof(self) -> None: + if self._compress is not None: + chunk = self._compress.flush() + if chunk: + self._compress = None + await self.write(chunk) + + if self._encoding == "base64": + if self._encoding_buffer: + await self._writer.write(base64.b64encode(self._encoding_buffer)) + + async def write(self, chunk: bytes) -> None: + if self._compress is not None: + if chunk: + chunk = await self._compress.compress(chunk) + if not chunk: + return + + if self._encoding == "base64": + buf = self._encoding_buffer + assert buf is not None + buf.extend(chunk) + + if buf: + div, mod = divmod(len(buf), 3) + enc_chunk, self._encoding_buffer = (buf[: div * 3], buf[div * 3 :]) + if enc_chunk: + b64chunk = base64.b64encode(enc_chunk) + await self._writer.write(b64chunk) + elif self._encoding == "quoted-printable": + await self._writer.write(binascii.b2a_qp(chunk)) + else: + await self._writer.write(chunk) diff --git a/lib/python3.12/site-packages/aiohttp/payload.py b/lib/python3.12/site-packages/aiohttp/payload.py new file mode 100644 index 0000000000000000000000000000000000000000..5b88fa094050fcef280050b54429f6bd43b166d3 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/payload.py @@ -0,0 +1,1120 @@ +import asyncio +import enum +import io +import json +import mimetypes +import os +import sys +import warnings +from abc import ABC, abstractmethod +from collections.abc import Iterable +from itertools import chain +from typing import ( + IO, + TYPE_CHECKING, + Any, + Dict, + Final, + List, + Optional, + Set, + TextIO, + Tuple, + Type, + Union, +) + +from multidict import CIMultiDict + +from . import hdrs +from .abc import AbstractStreamWriter +from .helpers import ( + _SENTINEL, + content_disposition_header, + guess_filename, + parse_mimetype, + sentinel, +) +from .streams import StreamReader +from .typedefs import JSONEncoder, _CIMultiDict + +__all__ = ( + "PAYLOAD_REGISTRY", + "get_payload", + "payload_type", + "Payload", + "BytesPayload", + "StringPayload", + "IOBasePayload", + "BytesIOPayload", + "BufferedReaderPayload", + "TextIOPayload", + "StringIOPayload", + "JsonPayload", + "AsyncIterablePayload", +) + +TOO_LARGE_BYTES_BODY: Final[int] = 2**20 # 1 MB +READ_SIZE: Final[int] = 2**16 # 64 KB +_CLOSE_FUTURES: Set[asyncio.Future[None]] = set() + + +class LookupError(Exception): + """Raised when no payload factory is found for the given data type.""" + + +class Order(str, enum.Enum): + normal = "normal" + try_first = "try_first" + try_last = "try_last" + + +def get_payload(data: Any, *args: Any, **kwargs: Any) -> "Payload": + return PAYLOAD_REGISTRY.get(data, *args, **kwargs) + + +def register_payload( + factory: Type["Payload"], type: Any, *, order: Order = Order.normal +) -> None: + PAYLOAD_REGISTRY.register(factory, type, order=order) + + +class payload_type: + def __init__(self, type: Any, *, order: Order = Order.normal) -> None: + self.type = type + self.order = order + + def __call__(self, factory: Type["Payload"]) -> Type["Payload"]: + register_payload(factory, self.type, order=self.order) + return factory + + +PayloadType = Type["Payload"] +_PayloadRegistryItem = Tuple[PayloadType, Any] + + +class PayloadRegistry: + """Payload registry. + + note: we need zope.interface for more efficient adapter search + """ + + __slots__ = ("_first", "_normal", "_last", "_normal_lookup") + + def __init__(self) -> None: + self._first: List[_PayloadRegistryItem] = [] + self._normal: List[_PayloadRegistryItem] = [] + self._last: List[_PayloadRegistryItem] = [] + self._normal_lookup: Dict[Any, PayloadType] = {} + + def get( + self, + data: Any, + *args: Any, + _CHAIN: "Type[chain[_PayloadRegistryItem]]" = chain, + **kwargs: Any, + ) -> "Payload": + if self._first: + for factory, type_ in self._first: + if isinstance(data, type_): + return factory(data, *args, **kwargs) + # Try the fast lookup first + if lookup_factory := self._normal_lookup.get(type(data)): + return lookup_factory(data, *args, **kwargs) + # Bail early if its already a Payload + if isinstance(data, Payload): + return data + # Fallback to the slower linear search + for factory, type_ in _CHAIN(self._normal, self._last): + if isinstance(data, type_): + return factory(data, *args, **kwargs) + raise LookupError() + + def register( + self, factory: PayloadType, type: Any, *, order: Order = Order.normal + ) -> None: + if order is Order.try_first: + self._first.append((factory, type)) + elif order is Order.normal: + self._normal.append((factory, type)) + if isinstance(type, Iterable): + for t in type: + self._normal_lookup[t] = factory + else: + self._normal_lookup[type] = factory + elif order is Order.try_last: + self._last.append((factory, type)) + else: + raise ValueError(f"Unsupported order {order!r}") + + +class Payload(ABC): + + _default_content_type: str = "application/octet-stream" + _size: Optional[int] = None + _consumed: bool = False # Default: payload has not been consumed yet + _autoclose: bool = False # Default: assume resource needs explicit closing + + def __init__( + self, + value: Any, + headers: Optional[ + Union[_CIMultiDict, Dict[str, str], Iterable[Tuple[str, str]]] + ] = None, + content_type: Union[str, None, _SENTINEL] = sentinel, + filename: Optional[str] = None, + encoding: Optional[str] = None, + **kwargs: Any, + ) -> None: + self._encoding = encoding + self._filename = filename + self._headers: _CIMultiDict = CIMultiDict() + self._value = value + if content_type is not sentinel and content_type is not None: + self._headers[hdrs.CONTENT_TYPE] = content_type + elif self._filename is not None: + if sys.version_info >= (3, 13): + guesser = mimetypes.guess_file_type + else: + guesser = mimetypes.guess_type + content_type = guesser(self._filename)[0] + if content_type is None: + content_type = self._default_content_type + self._headers[hdrs.CONTENT_TYPE] = content_type + else: + self._headers[hdrs.CONTENT_TYPE] = self._default_content_type + if headers: + self._headers.update(headers) + + @property + def size(self) -> Optional[int]: + """Size of the payload in bytes. + + Returns the number of bytes that will be transmitted when the payload + is written. For string payloads, this is the size after encoding to bytes, + not the length of the string. + """ + return self._size + + @property + def filename(self) -> Optional[str]: + """Filename of the payload.""" + return self._filename + + @property + def headers(self) -> _CIMultiDict: + """Custom item headers""" + return self._headers + + @property + def _binary_headers(self) -> bytes: + return ( + "".join([k + ": " + v + "\r\n" for k, v in self.headers.items()]).encode( + "utf-8" + ) + + b"\r\n" + ) + + @property + def encoding(self) -> Optional[str]: + """Payload encoding""" + return self._encoding + + @property + def content_type(self) -> str: + """Content type""" + return self._headers[hdrs.CONTENT_TYPE] + + @property + def consumed(self) -> bool: + """Whether the payload has been consumed and cannot be reused.""" + return self._consumed + + @property + def autoclose(self) -> bool: + """ + Whether the payload can close itself automatically. + + Returns True if the payload has no file handles or resources that need + explicit closing. If False, callers must await close() to release resources. + """ + return self._autoclose + + def set_content_disposition( + self, + disptype: str, + quote_fields: bool = True, + _charset: str = "utf-8", + **params: Any, + ) -> None: + """Sets ``Content-Disposition`` header.""" + self._headers[hdrs.CONTENT_DISPOSITION] = content_disposition_header( + disptype, quote_fields=quote_fields, _charset=_charset, **params + ) + + @abstractmethod + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + """ + Return string representation of the value. + + This is named decode() to allow compatibility with bytes objects. + """ + + @abstractmethod + async def write(self, writer: AbstractStreamWriter) -> None: + """ + Write payload to the writer stream. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + + This is a legacy method that writes the entire payload without length constraints. + + Important: + For new implementations, use write_with_length() instead of this method. + This method is maintained for backwards compatibility and will eventually + delegate to write_with_length(writer, None) in all implementations. + + All payload subclasses must override this method for backwards compatibility, + but new code should use write_with_length for more flexibility and control. + + """ + + # write_with_length is new in aiohttp 3.12 + # it should be overridden by subclasses + async def write_with_length( + self, writer: AbstractStreamWriter, content_length: Optional[int] + ) -> None: + """ + Write payload with a specific content length constraint. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + content_length: Maximum number of bytes to write (None for unlimited) + + This method allows writing payload content with a specific length constraint, + which is particularly useful for HTTP responses with Content-Length header. + + Note: + This is the base implementation that provides backwards compatibility + for subclasses that don't override this method. Specific payload types + should override this method to implement proper length-constrained writing. + + """ + # Backwards compatibility for subclasses that don't override this method + # and for the default implementation + await self.write(writer) + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This is a convenience method that calls decode() and encodes the result + to bytes using the specified encoding. + """ + # Use instance encoding if available, otherwise use parameter + actual_encoding = self._encoding or encoding + return self.decode(actual_encoding, errors).encode(actual_encoding) + + def _close(self) -> None: + """ + Async safe synchronous close operations for backwards compatibility. + + This method exists only for backwards compatibility with code that + needs to clean up payloads synchronously. In the future, we will + drop this method and only support the async close() method. + + WARNING: This method must be safe to call from within the event loop + without blocking. Subclasses should not perform any blocking I/O here. + + WARNING: This method must be called from within an event loop for + certain payload types (e.g., IOBasePayload). Calling it outside an + event loop may raise RuntimeError. + """ + # This is a no-op by default, but subclasses can override it + # for non-blocking cleanup operations. + + async def close(self) -> None: + """ + Close the payload if it holds any resources. + + IMPORTANT: This method must not await anything that might not finish + immediately, as it may be called during cleanup/cancellation. Schedule + any long-running operations without awaiting them. + + In the future, this will be the only close method supported. + """ + self._close() + + +class BytesPayload(Payload): + _value: bytes + # _consumed = False (inherited) - Bytes are immutable and can be reused + _autoclose = True # No file handle, just bytes in memory + + def __init__( + self, value: Union[bytes, bytearray, memoryview], *args: Any, **kwargs: Any + ) -> None: + if "content_type" not in kwargs: + kwargs["content_type"] = "application/octet-stream" + + super().__init__(value, *args, **kwargs) + + if isinstance(value, memoryview): + self._size = value.nbytes + elif isinstance(value, (bytes, bytearray)): + self._size = len(value) + else: + raise TypeError(f"value argument must be byte-ish, not {type(value)!r}") + + if self._size > TOO_LARGE_BYTES_BODY: + kwargs = {"source": self} + warnings.warn( + "Sending a large body directly with raw bytes might" + " lock the event loop. You should probably pass an " + "io.BytesIO object instead", + ResourceWarning, + **kwargs, + ) + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + return self._value.decode(encoding, errors) + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This method returns the raw bytes content of the payload. + It is equivalent to accessing the _value attribute directly. + """ + return self._value + + async def write(self, writer: AbstractStreamWriter) -> None: + """ + Write the entire bytes payload to the writer stream. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + + This method writes the entire bytes content without any length constraint. + + Note: + For new implementations that need length control, use write_with_length(). + This method is maintained for backwards compatibility and is equivalent + to write_with_length(writer, None). + + """ + await writer.write(self._value) + + async def write_with_length( + self, writer: AbstractStreamWriter, content_length: Optional[int] + ) -> None: + """ + Write bytes payload with a specific content length constraint. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + content_length: Maximum number of bytes to write (None for unlimited) + + This method writes either the entire byte sequence or a slice of it + up to the specified content_length. For BytesPayload, this operation + is performed efficiently using array slicing. + + """ + if content_length is not None: + await writer.write(self._value[:content_length]) + else: + await writer.write(self._value) + + +class StringPayload(BytesPayload): + def __init__( + self, + value: str, + *args: Any, + encoding: Optional[str] = None, + content_type: Optional[str] = None, + **kwargs: Any, + ) -> None: + + if encoding is None: + if content_type is None: + real_encoding = "utf-8" + content_type = "text/plain; charset=utf-8" + else: + mimetype = parse_mimetype(content_type) + real_encoding = mimetype.parameters.get("charset", "utf-8") + else: + if content_type is None: + content_type = "text/plain; charset=%s" % encoding + real_encoding = encoding + + super().__init__( + value.encode(real_encoding), + encoding=real_encoding, + content_type=content_type, + *args, + **kwargs, + ) + + +class StringIOPayload(StringPayload): + def __init__(self, value: IO[str], *args: Any, **kwargs: Any) -> None: + super().__init__(value.read(), *args, **kwargs) + + +class IOBasePayload(Payload): + _value: io.IOBase + # _consumed = False (inherited) - File can be re-read from the same position + _start_position: Optional[int] = None + # _autoclose = False (inherited) - Has file handle that needs explicit closing + + def __init__( + self, value: IO[Any], disposition: str = "attachment", *args: Any, **kwargs: Any + ) -> None: + if "filename" not in kwargs: + kwargs["filename"] = guess_filename(value) + + super().__init__(value, *args, **kwargs) + + if self._filename is not None and disposition is not None: + if hdrs.CONTENT_DISPOSITION not in self.headers: + self.set_content_disposition(disposition, filename=self._filename) + + def _set_or_restore_start_position(self) -> None: + """Set or restore the start position of the file-like object.""" + if self._start_position is None: + try: + self._start_position = self._value.tell() + except (OSError, AttributeError): + self._consumed = True # Cannot seek, mark as consumed + return + try: + self._value.seek(self._start_position) + except (OSError, AttributeError): + # Failed to seek back - mark as consumed since we've already read + self._consumed = True + + def _read_and_available_len( + self, remaining_content_len: Optional[int] + ) -> Tuple[Optional[int], bytes]: + """ + Read the file-like object and return both its total size and the first chunk. + + Args: + remaining_content_len: Optional limit on how many bytes to read in this operation. + If None, READ_SIZE will be used as the default chunk size. + + Returns: + A tuple containing: + - The total size of the remaining unread content (None if size cannot be determined) + - The first chunk of bytes read from the file object + + This method is optimized to perform both size calculation and initial read + in a single operation, which is executed in a single executor job to minimize + context switches and file operations when streaming content. + + """ + self._set_or_restore_start_position() + size = self.size # Call size only once since it does I/O + return size, self._value.read( + min(READ_SIZE, size or READ_SIZE, remaining_content_len or READ_SIZE) + ) + + def _read(self, remaining_content_len: Optional[int]) -> bytes: + """ + Read a chunk of data from the file-like object. + + Args: + remaining_content_len: Optional maximum number of bytes to read. + If None, READ_SIZE will be used as the default chunk size. + + Returns: + A chunk of bytes read from the file object, respecting the + remaining_content_len limit if specified. + + This method is used for subsequent reads during streaming after + the initial _read_and_available_len call has been made. + + """ + return self._value.read(remaining_content_len or READ_SIZE) # type: ignore[no-any-return] + + @property + def size(self) -> Optional[int]: + """ + Size of the payload in bytes. + + Returns the total size of the payload content from the initial position. + This ensures consistent Content-Length for requests, including 307/308 redirects + where the same payload instance is reused. + + Returns None if the size cannot be determined (e.g., for unseekable streams). + """ + try: + # Store the start position on first access. + # This is critical when the same payload instance is reused (e.g., 307/308 + # redirects). Without storing the initial position, after the payload is + # read once, the file position would be at EOF, which would cause the + # size calculation to return 0 (file_size - EOF position). + # By storing the start position, we ensure the size calculation always + # returns the correct total size for any subsequent use. + if self._start_position is None: + self._start_position = self._value.tell() + + # Return the total size from the start position + # This ensures Content-Length is correct even after reading + return os.fstat(self._value.fileno()).st_size - self._start_position + except (AttributeError, OSError): + return None + + async def write(self, writer: AbstractStreamWriter) -> None: + """ + Write the entire file-like payload to the writer stream. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + + This method writes the entire file content without any length constraint. + It delegates to write_with_length() with no length limit for implementation + consistency. + + Note: + For new implementations that need length control, use write_with_length() directly. + This method is maintained for backwards compatibility with existing code. + + """ + await self.write_with_length(writer, None) + + async def write_with_length( + self, writer: AbstractStreamWriter, content_length: Optional[int] + ) -> None: + """ + Write file-like payload with a specific content length constraint. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + content_length: Maximum number of bytes to write (None for unlimited) + + This method implements optimized streaming of file content with length constraints: + + 1. File reading is performed in a thread pool to avoid blocking the event loop + 2. Content is read and written in chunks to maintain memory efficiency + 3. Writing stops when either: + - All available file content has been written (when size is known) + - The specified content_length has been reached + 4. File resources are properly closed even if the operation is cancelled + + The implementation carefully handles both known-size and unknown-size payloads, + as well as constrained and unconstrained content lengths. + + """ + loop = asyncio.get_running_loop() + total_written_len = 0 + remaining_content_len = content_length + + # Get initial data and available length + available_len, chunk = await loop.run_in_executor( + None, self._read_and_available_len, remaining_content_len + ) + # Process data chunks until done + while chunk: + chunk_len = len(chunk) + + # Write data with or without length constraint + if remaining_content_len is None: + await writer.write(chunk) + else: + await writer.write(chunk[:remaining_content_len]) + remaining_content_len -= chunk_len + + total_written_len += chunk_len + + # Check if we're done writing + if self._should_stop_writing( + available_len, total_written_len, remaining_content_len + ): + return + + # Read next chunk + chunk = await loop.run_in_executor( + None, + self._read, + ( + min(READ_SIZE, remaining_content_len) + if remaining_content_len is not None + else READ_SIZE + ), + ) + + def _should_stop_writing( + self, + available_len: Optional[int], + total_written_len: int, + remaining_content_len: Optional[int], + ) -> bool: + """ + Determine if we should stop writing data. + + Args: + available_len: Known size of the payload if available (None if unknown) + total_written_len: Number of bytes already written + remaining_content_len: Remaining bytes to be written for content-length limited responses + + Returns: + True if we should stop writing data, based on either: + - Having written all available data (when size is known) + - Having written all requested content (when content-length is specified) + + """ + return (available_len is not None and total_written_len >= available_len) or ( + remaining_content_len is not None and remaining_content_len <= 0 + ) + + def _close(self) -> None: + """ + Async safe synchronous close operations for backwards compatibility. + + This method exists only for backwards + compatibility. Use the async close() method instead. + + WARNING: This method MUST be called from within an event loop. + Calling it outside an event loop will raise RuntimeError. + """ + # Skip if already consumed + if self._consumed: + return + self._consumed = True # Mark as consumed to prevent further writes + # Schedule file closing without awaiting to prevent cancellation issues + loop = asyncio.get_running_loop() + close_future = loop.run_in_executor(None, self._value.close) + # Hold a strong reference to the future to prevent it from being + # garbage collected before it completes. + _CLOSE_FUTURES.add(close_future) + close_future.add_done_callback(_CLOSE_FUTURES.remove) + + async def close(self) -> None: + """ + Close the payload if it holds any resources. + + IMPORTANT: This method must not await anything that might not finish + immediately, as it may be called during cleanup/cancellation. Schedule + any long-running operations without awaiting them. + """ + self._close() + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + """ + Return string representation of the value. + + WARNING: This method does blocking I/O and should not be called in the event loop. + """ + return self._read_all().decode(encoding, errors) + + def _read_all(self) -> bytes: + """Read the entire file-like object and return its content as bytes.""" + self._set_or_restore_start_position() + # Use readlines() to ensure we get all content + return b"".join(self._value.readlines()) + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This method reads the entire file content and returns it as bytes. + It is equivalent to reading the file-like object directly. + The file reading is performed in an executor to avoid blocking the event loop. + """ + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, self._read_all) + + +class TextIOPayload(IOBasePayload): + _value: io.TextIOBase + # _autoclose = False (inherited) - Has text file handle that needs explicit closing + + def __init__( + self, + value: TextIO, + *args: Any, + encoding: Optional[str] = None, + content_type: Optional[str] = None, + **kwargs: Any, + ) -> None: + + if encoding is None: + if content_type is None: + encoding = "utf-8" + content_type = "text/plain; charset=utf-8" + else: + mimetype = parse_mimetype(content_type) + encoding = mimetype.parameters.get("charset", "utf-8") + else: + if content_type is None: + content_type = "text/plain; charset=%s" % encoding + + super().__init__( + value, + content_type=content_type, + encoding=encoding, + *args, + **kwargs, + ) + + def _read_and_available_len( + self, remaining_content_len: Optional[int] + ) -> Tuple[Optional[int], bytes]: + """ + Read the text file-like object and return both its total size and the first chunk. + + Args: + remaining_content_len: Optional limit on how many bytes to read in this operation. + If None, READ_SIZE will be used as the default chunk size. + + Returns: + A tuple containing: + - The total size of the remaining unread content (None if size cannot be determined) + - The first chunk of bytes read from the file object, encoded using the payload's encoding + + This method is optimized to perform both size calculation and initial read + in a single operation, which is executed in a single executor job to minimize + context switches and file operations when streaming content. + + Note: + TextIOPayload handles encoding of the text content before writing it + to the stream. If no encoding is specified, UTF-8 is used as the default. + + """ + self._set_or_restore_start_position() + size = self.size + chunk = self._value.read( + min(READ_SIZE, size or READ_SIZE, remaining_content_len or READ_SIZE) + ) + return size, chunk.encode(self._encoding) if self._encoding else chunk.encode() + + def _read(self, remaining_content_len: Optional[int]) -> bytes: + """ + Read a chunk of data from the text file-like object. + + Args: + remaining_content_len: Optional maximum number of bytes to read. + If None, READ_SIZE will be used as the default chunk size. + + Returns: + A chunk of bytes read from the file object and encoded using the payload's + encoding. The data is automatically converted from text to bytes. + + This method is used for subsequent reads during streaming after + the initial _read_and_available_len call has been made. It properly + handles text encoding, converting the text content to bytes using + the specified encoding (or UTF-8 if none was provided). + + """ + chunk = self._value.read(remaining_content_len or READ_SIZE) + return chunk.encode(self._encoding) if self._encoding else chunk.encode() + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + """ + Return string representation of the value. + + WARNING: This method does blocking I/O and should not be called in the event loop. + """ + self._set_or_restore_start_position() + return self._value.read() + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This method reads the entire text file content and returns it as bytes. + It encodes the text content using the specified encoding. + The file reading is performed in an executor to avoid blocking the event loop. + """ + loop = asyncio.get_running_loop() + + # Use instance encoding if available, otherwise use parameter + actual_encoding = self._encoding or encoding + + def _read_and_encode() -> bytes: + self._set_or_restore_start_position() + # TextIO read() always returns the full content + return self._value.read().encode(actual_encoding, errors) + + return await loop.run_in_executor(None, _read_and_encode) + + +class BytesIOPayload(IOBasePayload): + _value: io.BytesIO + _size: int # Always initialized in __init__ + _autoclose = True # BytesIO is in-memory, safe to auto-close + + def __init__(self, value: io.BytesIO, *args: Any, **kwargs: Any) -> None: + super().__init__(value, *args, **kwargs) + # Calculate size once during initialization + self._size = len(self._value.getbuffer()) - self._value.tell() + + @property + def size(self) -> int: + """Size of the payload in bytes. + + Returns the number of bytes in the BytesIO buffer that will be transmitted. + This is calculated once during initialization for efficiency. + """ + return self._size + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + self._set_or_restore_start_position() + return self._value.read().decode(encoding, errors) + + async def write(self, writer: AbstractStreamWriter) -> None: + return await self.write_with_length(writer, None) + + async def write_with_length( + self, writer: AbstractStreamWriter, content_length: Optional[int] + ) -> None: + """ + Write BytesIO payload with a specific content length constraint. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + content_length: Maximum number of bytes to write (None for unlimited) + + This implementation is specifically optimized for BytesIO objects: + + 1. Reads content in chunks to maintain memory efficiency + 2. Yields control back to the event loop periodically to prevent blocking + when dealing with large BytesIO objects + 3. Respects content_length constraints when specified + 4. Properly cleans up by closing the BytesIO object when done or on error + + The periodic yielding to the event loop is important for maintaining + responsiveness when processing large in-memory buffers. + + """ + self._set_or_restore_start_position() + loop_count = 0 + remaining_bytes = content_length + while chunk := self._value.read(READ_SIZE): + if loop_count > 0: + # Avoid blocking the event loop + # if they pass a large BytesIO object + # and we are not in the first iteration + # of the loop + await asyncio.sleep(0) + if remaining_bytes is None: + await writer.write(chunk) + else: + await writer.write(chunk[:remaining_bytes]) + remaining_bytes -= len(chunk) + if remaining_bytes <= 0: + return + loop_count += 1 + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This method reads the entire BytesIO content and returns it as bytes. + It is equivalent to accessing the _value attribute directly. + """ + self._set_or_restore_start_position() + return self._value.read() + + async def close(self) -> None: + """ + Close the BytesIO payload. + + This does nothing since BytesIO is in-memory and does not require explicit closing. + """ + + +class BufferedReaderPayload(IOBasePayload): + _value: io.BufferedIOBase + # _autoclose = False (inherited) - Has buffered file handle that needs explicit closing + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + self._set_or_restore_start_position() + return self._value.read().decode(encoding, errors) + + +class JsonPayload(BytesPayload): + def __init__( + self, + value: Any, + encoding: str = "utf-8", + content_type: str = "application/json", + dumps: JSONEncoder = json.dumps, + *args: Any, + **kwargs: Any, + ) -> None: + + super().__init__( + dumps(value).encode(encoding), + content_type=content_type, + encoding=encoding, + *args, + **kwargs, + ) + + +if TYPE_CHECKING: + from typing import AsyncIterable, AsyncIterator + + _AsyncIterator = AsyncIterator[bytes] + _AsyncIterable = AsyncIterable[bytes] +else: + from collections.abc import AsyncIterable, AsyncIterator + + _AsyncIterator = AsyncIterator + _AsyncIterable = AsyncIterable + + +class AsyncIterablePayload(Payload): + + _iter: Optional[_AsyncIterator] = None + _value: _AsyncIterable + _cached_chunks: Optional[List[bytes]] = None + # _consumed stays False to allow reuse with cached content + _autoclose = True # Iterator doesn't need explicit closing + + def __init__(self, value: _AsyncIterable, *args: Any, **kwargs: Any) -> None: + if not isinstance(value, AsyncIterable): + raise TypeError( + "value argument must support " + "collections.abc.AsyncIterable interface, " + "got {!r}".format(type(value)) + ) + + if "content_type" not in kwargs: + kwargs["content_type"] = "application/octet-stream" + + super().__init__(value, *args, **kwargs) + + self._iter = value.__aiter__() + + async def write(self, writer: AbstractStreamWriter) -> None: + """ + Write the entire async iterable payload to the writer stream. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + + This method iterates through the async iterable and writes each chunk + to the writer without any length constraint. + + Note: + For new implementations that need length control, use write_with_length() directly. + This method is maintained for backwards compatibility with existing code. + + """ + await self.write_with_length(writer, None) + + async def write_with_length( + self, writer: AbstractStreamWriter, content_length: Optional[int] + ) -> None: + """ + Write async iterable payload with a specific content length constraint. + + Args: + writer: An AbstractStreamWriter instance that handles the actual writing + content_length: Maximum number of bytes to write (None for unlimited) + + This implementation handles streaming of async iterable content with length constraints: + + 1. If cached chunks are available, writes from them + 2. Otherwise iterates through the async iterable one chunk at a time + 3. Respects content_length constraints when specified + 4. Does NOT generate cache - that's done by as_bytes() + + """ + # If we have cached chunks, use them + if self._cached_chunks is not None: + remaining_bytes = content_length + for chunk in self._cached_chunks: + if remaining_bytes is None: + await writer.write(chunk) + elif remaining_bytes > 0: + await writer.write(chunk[:remaining_bytes]) + remaining_bytes -= len(chunk) + else: + break + return + + # If iterator is exhausted and we don't have cached chunks, nothing to write + if self._iter is None: + return + + # Stream from the iterator + remaining_bytes = content_length + + try: + while True: + if sys.version_info >= (3, 10): + chunk = await anext(self._iter) + else: + chunk = await self._iter.__anext__() + if remaining_bytes is None: + await writer.write(chunk) + # If we have a content length limit + elif remaining_bytes > 0: + await writer.write(chunk[:remaining_bytes]) + remaining_bytes -= len(chunk) + # We still want to exhaust the iterator even + # if we have reached the content length limit + # since the file handle may not get closed by + # the iterator if we don't do this + except StopAsyncIteration: + # Iterator is exhausted + self._iter = None + self._consumed = True # Mark as consumed when streamed without caching + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + """Decode the payload content as a string if cached chunks are available.""" + if self._cached_chunks is not None: + return b"".join(self._cached_chunks).decode(encoding, errors) + raise TypeError("Unable to decode - content not cached. Call as_bytes() first.") + + async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> bytes: + """ + Return bytes representation of the value. + + This method reads the entire async iterable content and returns it as bytes. + It generates and caches the chunks for future reuse. + """ + # If we have cached chunks, return them joined + if self._cached_chunks is not None: + return b"".join(self._cached_chunks) + + # If iterator is exhausted and no cache, return empty + if self._iter is None: + return b"" + + # Read all chunks and cache them + chunks: List[bytes] = [] + async for chunk in self._iter: + chunks.append(chunk) + + # Iterator is exhausted, cache the chunks + self._iter = None + self._cached_chunks = chunks + # Keep _consumed as False to allow reuse with cached chunks + + return b"".join(chunks) + + +class StreamReaderPayload(AsyncIterablePayload): + def __init__(self, value: StreamReader, *args: Any, **kwargs: Any) -> None: + super().__init__(value.iter_any(), *args, **kwargs) + + +PAYLOAD_REGISTRY = PayloadRegistry() +PAYLOAD_REGISTRY.register(BytesPayload, (bytes, bytearray, memoryview)) +PAYLOAD_REGISTRY.register(StringPayload, str) +PAYLOAD_REGISTRY.register(StringIOPayload, io.StringIO) +PAYLOAD_REGISTRY.register(TextIOPayload, io.TextIOBase) +PAYLOAD_REGISTRY.register(BytesIOPayload, io.BytesIO) +PAYLOAD_REGISTRY.register(BufferedReaderPayload, (io.BufferedReader, io.BufferedRandom)) +PAYLOAD_REGISTRY.register(IOBasePayload, io.IOBase) +PAYLOAD_REGISTRY.register(StreamReaderPayload, StreamReader) +# try_last for giving a chance to more specialized async interables like +# multipart.BodyPartReaderPayload override the default +PAYLOAD_REGISTRY.register(AsyncIterablePayload, AsyncIterable, order=Order.try_last) diff --git a/lib/python3.12/site-packages/aiohttp/payload_streamer.py b/lib/python3.12/site-packages/aiohttp/payload_streamer.py new file mode 100644 index 0000000000000000000000000000000000000000..831fdc0a77f302acaf9a000be408fe7c9a9035aa --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/payload_streamer.py @@ -0,0 +1,78 @@ +""" +Payload implementation for coroutines as data provider. + +As a simple case, you can upload data from file:: + + @aiohttp.streamer + async def file_sender(writer, file_name=None): + with open(file_name, 'rb') as f: + chunk = f.read(2**16) + while chunk: + await writer.write(chunk) + + chunk = f.read(2**16) + +Then you can use `file_sender` like this: + + async with session.post('http://httpbin.org/post', + data=file_sender(file_name='huge_file')) as resp: + print(await resp.text()) + +..note:: Coroutine must accept `writer` as first argument + +""" + +import types +import warnings +from typing import Any, Awaitable, Callable, Dict, Tuple + +from .abc import AbstractStreamWriter +from .payload import Payload, payload_type + +__all__ = ("streamer",) + + +class _stream_wrapper: + def __init__( + self, + coro: Callable[..., Awaitable[None]], + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> None: + self.coro = types.coroutine(coro) + self.args = args + self.kwargs = kwargs + + async def __call__(self, writer: AbstractStreamWriter) -> None: + await self.coro(writer, *self.args, **self.kwargs) + + +class streamer: + def __init__(self, coro: Callable[..., Awaitable[None]]) -> None: + warnings.warn( + "@streamer is deprecated, use async generators instead", + DeprecationWarning, + stacklevel=2, + ) + self.coro = coro + + def __call__(self, *args: Any, **kwargs: Any) -> _stream_wrapper: + return _stream_wrapper(self.coro, args, kwargs) + + +@payload_type(_stream_wrapper) +class StreamWrapperPayload(Payload): + async def write(self, writer: AbstractStreamWriter) -> None: + await self._value(writer) + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + raise TypeError("Unable to decode.") + + +@payload_type(streamer) +class StreamPayload(StreamWrapperPayload): + def __init__(self, value: Any, *args: Any, **kwargs: Any) -> None: + super().__init__(value(), *args, **kwargs) + + async def write(self, writer: AbstractStreamWriter) -> None: + await self._value(writer) diff --git a/lib/python3.12/site-packages/aiohttp/py.typed b/lib/python3.12/site-packages/aiohttp/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..f5642f79f21d872f010979dcf6f0c4a415acc19d --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/py.typed @@ -0,0 +1 @@ +Marker diff --git a/lib/python3.12/site-packages/aiohttp/pytest_plugin.py b/lib/python3.12/site-packages/aiohttp/pytest_plugin.py new file mode 100644 index 0000000000000000000000000000000000000000..7d59fe820d697632c9a3311ff96841ebc0ee735b --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/pytest_plugin.py @@ -0,0 +1,444 @@ +import asyncio +import contextlib +import inspect +import warnings +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Iterator, + Optional, + Protocol, + Union, + overload, +) + +import pytest + +from .test_utils import ( + BaseTestServer, + RawTestServer, + TestClient, + TestServer, + loop_context, + setup_test_loop, + teardown_test_loop, + unused_port as _unused_port, +) +from .web import Application, BaseRequest, Request +from .web_protocol import _RequestHandler + +try: + import uvloop +except ImportError: # pragma: no cover + uvloop = None # type: ignore[assignment] + + +class AiohttpClient(Protocol): + @overload + async def __call__( + self, + __param: Application, + *, + server_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> TestClient[Request, Application]: ... + @overload + async def __call__( + self, + __param: BaseTestServer, + *, + server_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> TestClient[BaseRequest, None]: ... + + +class AiohttpServer(Protocol): + def __call__( + self, app: Application, *, port: Optional[int] = None, **kwargs: Any + ) -> Awaitable[TestServer]: ... + + +class AiohttpRawServer(Protocol): + def __call__( + self, handler: _RequestHandler, *, port: Optional[int] = None, **kwargs: Any + ) -> Awaitable[RawTestServer]: ... + + +def pytest_addoption(parser): # type: ignore[no-untyped-def] + parser.addoption( + "--aiohttp-fast", + action="store_true", + default=False, + help="run tests faster by disabling extra checks", + ) + parser.addoption( + "--aiohttp-loop", + action="store", + default="pyloop", + help="run tests with specific loop: pyloop, uvloop or all", + ) + parser.addoption( + "--aiohttp-enable-loop-debug", + action="store_true", + default=False, + help="enable event loop debug mode", + ) + + +def pytest_fixture_setup(fixturedef): # type: ignore[no-untyped-def] + """Set up pytest fixture. + + Allow fixtures to be coroutines. Run coroutine fixtures in an event loop. + """ + func = fixturedef.func + + if inspect.isasyncgenfunction(func): + # async generator fixture + is_async_gen = True + elif inspect.iscoroutinefunction(func): + # regular async fixture + is_async_gen = False + else: + # not an async fixture, nothing to do + return + + strip_request = False + if "request" not in fixturedef.argnames: + fixturedef.argnames += ("request",) + strip_request = True + + def wrapper(*args, **kwargs): # type: ignore[no-untyped-def] + request = kwargs["request"] + if strip_request: + del kwargs["request"] + + # if neither the fixture nor the test use the 'loop' fixture, + # 'getfixturevalue' will fail because the test is not parameterized + # (this can be removed someday if 'loop' is no longer parameterized) + if "loop" not in request.fixturenames: + raise Exception( + "Asynchronous fixtures must depend on the 'loop' fixture or " + "be used in tests depending from it." + ) + + _loop = request.getfixturevalue("loop") + + if is_async_gen: + # for async generators, we need to advance the generator once, + # then advance it again in a finalizer + gen = func(*args, **kwargs) + + def finalizer(): # type: ignore[no-untyped-def] + try: + return _loop.run_until_complete(gen.__anext__()) + except StopAsyncIteration: + pass + + request.addfinalizer(finalizer) + return _loop.run_until_complete(gen.__anext__()) + else: + return _loop.run_until_complete(func(*args, **kwargs)) + + fixturedef.func = wrapper + + +@pytest.fixture +def fast(request): # type: ignore[no-untyped-def] + """--fast config option""" + return request.config.getoption("--aiohttp-fast") + + +@pytest.fixture +def loop_debug(request): # type: ignore[no-untyped-def] + """--enable-loop-debug config option""" + return request.config.getoption("--aiohttp-enable-loop-debug") + + +@contextlib.contextmanager +def _runtime_warning_context(): # type: ignore[no-untyped-def] + """Context manager which checks for RuntimeWarnings. + + This exists specifically to + avoid "coroutine 'X' was never awaited" warnings being missed. + + If RuntimeWarnings occur in the context a RuntimeError is raised. + """ + with warnings.catch_warnings(record=True) as _warnings: + yield + rw = [ + "{w.filename}:{w.lineno}:{w.message}".format(w=w) + for w in _warnings + if w.category == RuntimeWarning + ] + if rw: + raise RuntimeError( + "{} Runtime Warning{},\n{}".format( + len(rw), "" if len(rw) == 1 else "s", "\n".join(rw) + ) + ) + + +@contextlib.contextmanager +def _passthrough_loop_context(loop, fast=False): # type: ignore[no-untyped-def] + """Passthrough loop context. + + Sets up and tears down a loop unless one is passed in via the loop + argument when it's passed straight through. + """ + if loop: + # loop already exists, pass it straight through + yield loop + else: + # this shadows loop_context's standard behavior + loop = setup_test_loop() + yield loop + teardown_test_loop(loop, fast=fast) + + +def pytest_pycollect_makeitem(collector, name, obj): # type: ignore[no-untyped-def] + """Fix pytest collecting for coroutines.""" + if collector.funcnamefilter(name) and inspect.iscoroutinefunction(obj): + return list(collector._genfunctions(name, obj)) + + +def pytest_pyfunc_call(pyfuncitem): # type: ignore[no-untyped-def] + """Run coroutines in an event loop instead of a normal function call.""" + fast = pyfuncitem.config.getoption("--aiohttp-fast") + if inspect.iscoroutinefunction(pyfuncitem.function): + existing_loop = ( + pyfuncitem.funcargs.get("proactor_loop") + or pyfuncitem.funcargs.get("selector_loop") + or pyfuncitem.funcargs.get("uvloop_loop") + or pyfuncitem.funcargs.get("loop", None) + ) + + with _runtime_warning_context(): + with _passthrough_loop_context(existing_loop, fast=fast) as _loop: + testargs = { + arg: pyfuncitem.funcargs[arg] + for arg in pyfuncitem._fixtureinfo.argnames + } + _loop.run_until_complete(pyfuncitem.obj(**testargs)) + + return True + + +def pytest_generate_tests(metafunc): # type: ignore[no-untyped-def] + if "loop_factory" not in metafunc.fixturenames: + return + + loops = metafunc.config.option.aiohttp_loop + avail_factories: dict[str, Callable[[], asyncio.AbstractEventLoop]] + avail_factories = {"pyloop": asyncio.new_event_loop} + + if uvloop is not None: # pragma: no cover + avail_factories["uvloop"] = uvloop.new_event_loop + + if loops == "all": + loops = "pyloop,uvloop?" + + factories = {} # type: ignore[var-annotated] + for name in loops.split(","): + required = not name.endswith("?") + name = name.strip(" ?") + if name not in avail_factories: # pragma: no cover + if required: + raise ValueError( + "Unknown loop '%s', available loops: %s" + % (name, list(factories.keys())) + ) + else: + continue + factories[name] = avail_factories[name] + metafunc.parametrize( + "loop_factory", list(factories.values()), ids=list(factories.keys()) + ) + + +@pytest.fixture +def loop( + loop_factory: Callable[[], asyncio.AbstractEventLoop], + fast: bool, + loop_debug: bool, +) -> Iterator[asyncio.AbstractEventLoop]: + """Return an instance of the event loop.""" + with loop_context(loop_factory, fast=fast) as _loop: + if loop_debug: + _loop.set_debug(True) # pragma: no cover + asyncio.set_event_loop(_loop) + yield _loop + + +@pytest.fixture +def proactor_loop() -> Iterator[asyncio.AbstractEventLoop]: + factory = asyncio.ProactorEventLoop # type: ignore[attr-defined] + + with loop_context(factory) as _loop: + asyncio.set_event_loop(_loop) + yield _loop + + +@pytest.fixture +def unused_port(aiohttp_unused_port: Callable[[], int]) -> Callable[[], int]: + warnings.warn( + "Deprecated, use aiohttp_unused_port fixture instead", + DeprecationWarning, + stacklevel=2, + ) + return aiohttp_unused_port + + +@pytest.fixture +def aiohttp_unused_port() -> Callable[[], int]: + """Return a port that is unused on the current host.""" + return _unused_port + + +@pytest.fixture +def aiohttp_server(loop: asyncio.AbstractEventLoop) -> Iterator[AiohttpServer]: + """Factory to create a TestServer instance, given an app. + + aiohttp_server(app, **kwargs) + """ + servers = [] + + async def go( + app: Application, + *, + host: str = "127.0.0.1", + port: Optional[int] = None, + **kwargs: Any, + ) -> TestServer: + server = TestServer(app, host=host, port=port) + await server.start_server(loop=loop, **kwargs) + servers.append(server) + return server + + yield go + + async def finalize() -> None: + while servers: + await servers.pop().close() + + loop.run_until_complete(finalize()) + + +@pytest.fixture +def test_server(aiohttp_server): # type: ignore[no-untyped-def] # pragma: no cover + warnings.warn( + "Deprecated, use aiohttp_server fixture instead", + DeprecationWarning, + stacklevel=2, + ) + return aiohttp_server + + +@pytest.fixture +def aiohttp_raw_server(loop: asyncio.AbstractEventLoop) -> Iterator[AiohttpRawServer]: + """Factory to create a RawTestServer instance, given a web handler. + + aiohttp_raw_server(handler, **kwargs) + """ + servers = [] + + async def go( + handler: _RequestHandler, *, port: Optional[int] = None, **kwargs: Any + ) -> RawTestServer: + server = RawTestServer(handler, port=port) + await server.start_server(loop=loop, **kwargs) + servers.append(server) + return server + + yield go + + async def finalize() -> None: + while servers: + await servers.pop().close() + + loop.run_until_complete(finalize()) + + +@pytest.fixture +def raw_test_server( # type: ignore[no-untyped-def] # pragma: no cover + aiohttp_raw_server, +): + warnings.warn( + "Deprecated, use aiohttp_raw_server fixture instead", + DeprecationWarning, + stacklevel=2, + ) + return aiohttp_raw_server + + +@pytest.fixture +def aiohttp_client(loop: asyncio.AbstractEventLoop) -> Iterator[AiohttpClient]: + """Factory to create a TestClient instance. + + aiohttp_client(app, **kwargs) + aiohttp_client(server, **kwargs) + aiohttp_client(raw_server, **kwargs) + """ + clients = [] + + @overload + async def go( + __param: Application, + *, + server_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> TestClient[Request, Application]: ... + + @overload + async def go( + __param: BaseTestServer, + *, + server_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> TestClient[BaseRequest, None]: ... + + async def go( + __param: Union[Application, BaseTestServer], + *args: Any, + server_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> TestClient[Any, Any]: + if isinstance(__param, Callable) and not isinstance( # type: ignore[arg-type] + __param, (Application, BaseTestServer) + ): + __param = __param(loop, *args, **kwargs) + kwargs = {} + else: + assert not args, "args should be empty" + + if isinstance(__param, Application): + server_kwargs = server_kwargs or {} + server = TestServer(__param, loop=loop, **server_kwargs) + client = TestClient(server, loop=loop, **kwargs) + elif isinstance(__param, BaseTestServer): + client = TestClient(__param, loop=loop, **kwargs) + else: + raise ValueError("Unknown argument type: %r" % type(__param)) + + await client.start_server() + clients.append(client) + return client + + yield go + + async def finalize() -> None: + while clients: + await clients.pop().close() + + loop.run_until_complete(finalize()) + + +@pytest.fixture +def test_client(aiohttp_client): # type: ignore[no-untyped-def] # pragma: no cover + warnings.warn( + "Deprecated, use aiohttp_client fixture instead", + DeprecationWarning, + stacklevel=2, + ) + return aiohttp_client diff --git a/lib/python3.12/site-packages/aiohttp/resolver.py b/lib/python3.12/site-packages/aiohttp/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..b20e5672ce51a1aadd7768301a3a1c8eb6007bf4 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/resolver.py @@ -0,0 +1,274 @@ +import asyncio +import socket +import weakref +from typing import Any, Dict, Final, List, Optional, Tuple, Type, Union + +from .abc import AbstractResolver, ResolveResult + +__all__ = ("ThreadedResolver", "AsyncResolver", "DefaultResolver") + + +try: + import aiodns + + aiodns_default = hasattr(aiodns.DNSResolver, "getaddrinfo") +except ImportError: # pragma: no cover + aiodns = None # type: ignore[assignment] + aiodns_default = False + + +_NUMERIC_SOCKET_FLAGS = socket.AI_NUMERICHOST | socket.AI_NUMERICSERV +_NAME_SOCKET_FLAGS = socket.NI_NUMERICHOST | socket.NI_NUMERICSERV +_AI_ADDRCONFIG = socket.AI_ADDRCONFIG +if hasattr(socket, "AI_MASK"): + _AI_ADDRCONFIG &= socket.AI_MASK + + +class ThreadedResolver(AbstractResolver): + """Threaded resolver. + + Uses an Executor for synchronous getaddrinfo() calls. + concurrent.futures.ThreadPoolExecutor is used by default. + """ + + def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None: + self._loop = loop or asyncio.get_running_loop() + + async def resolve( + self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET + ) -> List[ResolveResult]: + infos = await self._loop.getaddrinfo( + host, + port, + type=socket.SOCK_STREAM, + family=family, + flags=_AI_ADDRCONFIG, + ) + + hosts: List[ResolveResult] = [] + for family, _, proto, _, address in infos: + if family == socket.AF_INET6: + if len(address) < 3: + # IPv6 is not supported by Python build, + # or IPv6 is not enabled in the host + continue + if address[3]: + # This is essential for link-local IPv6 addresses. + # LL IPv6 is a VERY rare case. Strictly speaking, we should use + # getnameinfo() unconditionally, but performance makes sense. + resolved_host, _port = await self._loop.getnameinfo( + address, _NAME_SOCKET_FLAGS + ) + port = int(_port) + else: + resolved_host, port = address[:2] + else: # IPv4 + assert family == socket.AF_INET + resolved_host, port = address # type: ignore[misc] + hosts.append( + ResolveResult( + hostname=host, + host=resolved_host, + port=port, + family=family, + proto=proto, + flags=_NUMERIC_SOCKET_FLAGS, + ) + ) + + return hosts + + async def close(self) -> None: + pass + + +class AsyncResolver(AbstractResolver): + """Use the `aiodns` package to make asynchronous DNS lookups""" + + def __init__( + self, + loop: Optional[asyncio.AbstractEventLoop] = None, + *args: Any, + **kwargs: Any, + ) -> None: + if aiodns is None: + raise RuntimeError("Resolver requires aiodns library") + + self._loop = loop or asyncio.get_running_loop() + self._manager: Optional[_DNSResolverManager] = None + # If custom args are provided, create a dedicated resolver instance + # This means each AsyncResolver with custom args gets its own + # aiodns.DNSResolver instance + if args or kwargs: + self._resolver = aiodns.DNSResolver(*args, **kwargs) + return + # Use the shared resolver from the manager for default arguments + self._manager = _DNSResolverManager() + self._resolver = self._manager.get_resolver(self, self._loop) + + if not hasattr(self._resolver, "gethostbyname"): + # aiodns 1.1 is not available, fallback to DNSResolver.query + self.resolve = self._resolve_with_query # type: ignore + + async def resolve( + self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET + ) -> List[ResolveResult]: + try: + resp = await self._resolver.getaddrinfo( + host, + port=port, + type=socket.SOCK_STREAM, + family=family, + flags=_AI_ADDRCONFIG, + ) + except aiodns.error.DNSError as exc: + msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed" + raise OSError(None, msg) from exc + hosts: List[ResolveResult] = [] + for node in resp.nodes: + address: Union[Tuple[bytes, int], Tuple[bytes, int, int, int]] = node.addr + family = node.family + if family == socket.AF_INET6: + if len(address) > 3 and address[3]: + # This is essential for link-local IPv6 addresses. + # LL IPv6 is a VERY rare case. Strictly speaking, we should use + # getnameinfo() unconditionally, but performance makes sense. + result = await self._resolver.getnameinfo( + (address[0].decode("ascii"), *address[1:]), + _NAME_SOCKET_FLAGS, + ) + resolved_host = result.node + else: + resolved_host = address[0].decode("ascii") + port = address[1] + else: # IPv4 + assert family == socket.AF_INET + resolved_host = address[0].decode("ascii") + port = address[1] + hosts.append( + ResolveResult( + hostname=host, + host=resolved_host, + port=port, + family=family, + proto=0, + flags=_NUMERIC_SOCKET_FLAGS, + ) + ) + + if not hosts: + raise OSError(None, "DNS lookup failed") + + return hosts + + async def _resolve_with_query( + self, host: str, port: int = 0, family: int = socket.AF_INET + ) -> List[Dict[str, Any]]: + qtype: Final = "AAAA" if family == socket.AF_INET6 else "A" + + try: + resp = await self._resolver.query(host, qtype) + except aiodns.error.DNSError as exc: + msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed" + raise OSError(None, msg) from exc + + hosts = [] + for rr in resp: + hosts.append( + { + "hostname": host, + "host": rr.host, + "port": port, + "family": family, + "proto": 0, + "flags": socket.AI_NUMERICHOST, + } + ) + + if not hosts: + raise OSError(None, "DNS lookup failed") + + return hosts + + async def close(self) -> None: + if self._manager: + # Release the resolver from the manager if using the shared resolver + self._manager.release_resolver(self, self._loop) + self._manager = None # Clear reference to manager + self._resolver = None # type: ignore[assignment] # Clear reference to resolver + return + # Otherwise cancel our dedicated resolver + if self._resolver is not None: + self._resolver.cancel() + self._resolver = None # type: ignore[assignment] # Clear reference + + +class _DNSResolverManager: + """Manager for aiodns.DNSResolver objects. + + This class manages shared aiodns.DNSResolver instances + with no custom arguments across different event loops. + """ + + _instance: Optional["_DNSResolverManager"] = None + + def __new__(cls) -> "_DNSResolverManager": + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._init() + return cls._instance + + def _init(self) -> None: + # Use WeakKeyDictionary to allow event loops to be garbage collected + self._loop_data: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, + tuple["aiodns.DNSResolver", weakref.WeakSet["AsyncResolver"]], + ] = weakref.WeakKeyDictionary() + + def get_resolver( + self, client: "AsyncResolver", loop: asyncio.AbstractEventLoop + ) -> "aiodns.DNSResolver": + """Get or create the shared aiodns.DNSResolver instance for a specific event loop. + + Args: + client: The AsyncResolver instance requesting the resolver. + This is required to track resolver usage. + loop: The event loop to use for the resolver. + """ + # Create a new resolver and client set for this loop if it doesn't exist + if loop not in self._loop_data: + resolver = aiodns.DNSResolver(loop=loop) + client_set: weakref.WeakSet["AsyncResolver"] = weakref.WeakSet() + self._loop_data[loop] = (resolver, client_set) + else: + # Get the existing resolver and client set + resolver, client_set = self._loop_data[loop] + + # Register this client with the loop + client_set.add(client) + return resolver + + def release_resolver( + self, client: "AsyncResolver", loop: asyncio.AbstractEventLoop + ) -> None: + """Release the resolver for an AsyncResolver client when it's closed. + + Args: + client: The AsyncResolver instance to release. + loop: The event loop the resolver was using. + """ + # Remove client from its loop's tracking + current_loop_data = self._loop_data.get(loop) + if current_loop_data is None: + return + resolver, client_set = current_loop_data + client_set.discard(client) + # If no more clients for this loop, cancel and remove its resolver + if not client_set: + if resolver is not None: + resolver.cancel() + del self._loop_data[loop] + + +_DefaultType = Type[Union[AsyncResolver, ThreadedResolver]] +DefaultResolver: _DefaultType = AsyncResolver if aiodns_default else ThreadedResolver diff --git a/lib/python3.12/site-packages/aiohttp/streams.py b/lib/python3.12/site-packages/aiohttp/streams.py new file mode 100644 index 0000000000000000000000000000000000000000..6cc74fc9cbd5bcbc7e213b79cadd8d42b701a812 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/streams.py @@ -0,0 +1,758 @@ +import asyncio +import collections +import warnings +from typing import ( + Awaitable, + Callable, + Deque, + Final, + Generic, + List, + Optional, + Tuple, + TypeVar, +) + +from .base_protocol import BaseProtocol +from .helpers import ( + _EXC_SENTINEL, + BaseTimerContext, + TimerNoop, + set_exception, + set_result, +) +from .log import internal_logger + +__all__ = ( + "EMPTY_PAYLOAD", + "EofStream", + "StreamReader", + "DataQueue", +) + +_T = TypeVar("_T") + + +class EofStream(Exception): + """eof stream indication.""" + + +class AsyncStreamIterator(Generic[_T]): + + __slots__ = ("read_func",) + + def __init__(self, read_func: Callable[[], Awaitable[_T]]) -> None: + self.read_func = read_func + + def __aiter__(self) -> "AsyncStreamIterator[_T]": + return self + + async def __anext__(self) -> _T: + try: + rv = await self.read_func() + except EofStream: + raise StopAsyncIteration + if rv == b"": + raise StopAsyncIteration + return rv + + +class ChunkTupleAsyncStreamIterator: + + __slots__ = ("_stream",) + + def __init__(self, stream: "StreamReader") -> None: + self._stream = stream + + def __aiter__(self) -> "ChunkTupleAsyncStreamIterator": + return self + + async def __anext__(self) -> Tuple[bytes, bool]: + rv = await self._stream.readchunk() + if rv == (b"", False): + raise StopAsyncIteration + return rv + + +class AsyncStreamReaderMixin: + + __slots__ = () + + def __aiter__(self) -> AsyncStreamIterator[bytes]: + return AsyncStreamIterator(self.readline) # type: ignore[attr-defined] + + def iter_chunked(self, n: int) -> AsyncStreamIterator[bytes]: + """Returns an asynchronous iterator that yields chunks of size n.""" + return AsyncStreamIterator(lambda: self.read(n)) # type: ignore[attr-defined] + + def iter_any(self) -> AsyncStreamIterator[bytes]: + """Yield all available data as soon as it is received.""" + return AsyncStreamIterator(self.readany) # type: ignore[attr-defined] + + def iter_chunks(self) -> ChunkTupleAsyncStreamIterator: + """Yield chunks of data as they are received by the server. + + The yielded objects are tuples + of (bytes, bool) as returned by the StreamReader.readchunk method. + """ + return ChunkTupleAsyncStreamIterator(self) # type: ignore[arg-type] + + +class StreamReader(AsyncStreamReaderMixin): + """An enhancement of asyncio.StreamReader. + + Supports asynchronous iteration by line, chunk or as available:: + + async for line in reader: + ... + async for chunk in reader.iter_chunked(1024): + ... + async for slice in reader.iter_any(): + ... + + """ + + __slots__ = ( + "_protocol", + "_low_water", + "_high_water", + "_low_water_chunks", + "_high_water_chunks", + "_loop", + "_size", + "_cursor", + "_http_chunk_splits", + "_buffer", + "_buffer_offset", + "_eof", + "_waiter", + "_eof_waiter", + "_exception", + "_timer", + "_eof_callbacks", + "_eof_counter", + "total_bytes", + "total_compressed_bytes", + ) + + def __init__( + self, + protocol: BaseProtocol, + limit: int, + *, + timer: Optional[BaseTimerContext] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> None: + self._protocol = protocol + self._low_water = limit + self._high_water = limit * 2 + if loop is None: + loop = asyncio.get_event_loop() + # Ensure high_water_chunks >= 3 so it's always > low_water_chunks. + self._high_water_chunks = max(3, limit // 4) + # Use max(2, ...) because there's always at least 1 chunk split remaining + # (the current position), so we need low_water >= 2 to allow resume. + self._low_water_chunks = max(2, self._high_water_chunks // 2) + self._loop = loop + self._size = 0 + self._cursor = 0 + self._http_chunk_splits: Optional[Deque[int]] = None + self._buffer: Deque[bytes] = collections.deque() + self._buffer_offset = 0 + self._eof = False + self._waiter: Optional[asyncio.Future[None]] = None + self._eof_waiter: Optional[asyncio.Future[None]] = None + self._exception: Optional[BaseException] = None + self._timer = TimerNoop() if timer is None else timer + self._eof_callbacks: List[Callable[[], None]] = [] + self._eof_counter = 0 + self.total_bytes = 0 + self.total_compressed_bytes: Optional[int] = None + + def __repr__(self) -> str: + info = [self.__class__.__name__] + if self._size: + info.append("%d bytes" % self._size) + if self._eof: + info.append("eof") + if self._low_water != 2**16: # default limit + info.append("low=%d high=%d" % (self._low_water, self._high_water)) + if self._waiter: + info.append("w=%r" % self._waiter) + if self._exception: + info.append("e=%r" % self._exception) + return "<%s>" % " ".join(info) + + def get_read_buffer_limits(self) -> Tuple[int, int]: + return (self._low_water, self._high_water) + + def exception(self) -> Optional[BaseException]: + return self._exception + + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + self._exception = exc + self._eof_callbacks.clear() + + waiter = self._waiter + if waiter is not None: + self._waiter = None + set_exception(waiter, exc, exc_cause) + + waiter = self._eof_waiter + if waiter is not None: + self._eof_waiter = None + set_exception(waiter, exc, exc_cause) + + def on_eof(self, callback: Callable[[], None]) -> None: + if self._eof: + try: + callback() + except Exception: + internal_logger.exception("Exception in eof callback") + else: + self._eof_callbacks.append(callback) + + def feed_eof(self) -> None: + self._eof = True + + waiter = self._waiter + if waiter is not None: + self._waiter = None + set_result(waiter, None) + + waiter = self._eof_waiter + if waiter is not None: + self._eof_waiter = None + set_result(waiter, None) + + if self._protocol._reading_paused: + self._protocol.resume_reading() + + for cb in self._eof_callbacks: + try: + cb() + except Exception: + internal_logger.exception("Exception in eof callback") + + self._eof_callbacks.clear() + + def is_eof(self) -> bool: + """Return True if 'feed_eof' was called.""" + return self._eof + + def at_eof(self) -> bool: + """Return True if the buffer is empty and 'feed_eof' was called.""" + return self._eof and not self._buffer + + async def wait_eof(self) -> None: + if self._eof: + return + + assert self._eof_waiter is None + self._eof_waiter = self._loop.create_future() + try: + await self._eof_waiter + finally: + self._eof_waiter = None + + @property + def total_raw_bytes(self) -> int: + if self.total_compressed_bytes is None: + return self.total_bytes + return self.total_compressed_bytes + + def unread_data(self, data: bytes) -> None: + """rollback reading some data from stream, inserting it to buffer head.""" + warnings.warn( + "unread_data() is deprecated " + "and will be removed in future releases (#3260)", + DeprecationWarning, + stacklevel=2, + ) + if not data: + return + + if self._buffer_offset: + self._buffer[0] = self._buffer[0][self._buffer_offset :] + self._buffer_offset = 0 + self._size += len(data) + self._cursor -= len(data) + self._buffer.appendleft(data) + self._eof_counter = 0 + + # TODO: size is ignored, remove the param later + def feed_data(self, data: bytes, size: int = 0) -> None: + assert not self._eof, "feed_data after feed_eof" + + if not data: + return + + data_len = len(data) + self._size += data_len + self._buffer.append(data) + self.total_bytes += data_len + + waiter = self._waiter + if waiter is not None: + self._waiter = None + set_result(waiter, None) + + if self._size > self._high_water and not self._protocol._reading_paused: + self._protocol.pause_reading() + + def begin_http_chunk_receiving(self) -> None: + if self._http_chunk_splits is None: + if self.total_bytes: + raise RuntimeError( + "Called begin_http_chunk_receiving when some data was already fed" + ) + self._http_chunk_splits = collections.deque() + + def end_http_chunk_receiving(self) -> None: + if self._http_chunk_splits is None: + raise RuntimeError( + "Called end_chunk_receiving without calling " + "begin_chunk_receiving first" + ) + + # self._http_chunk_splits contains logical byte offsets from start of + # the body transfer. Each offset is the offset of the end of a chunk. + # "Logical" means bytes, accessible for a user. + # If no chunks containing logical data were received, current position + # is difinitely zero. + pos = self._http_chunk_splits[-1] if self._http_chunk_splits else 0 + + if self.total_bytes == pos: + # We should not add empty chunks here. So we check for that. + # Note, when chunked + gzip is used, we can receive a chunk + # of compressed data, but that data may not be enough for gzip FSM + # to yield any uncompressed data. That's why current position may + # not change after receiving a chunk. + return + + self._http_chunk_splits.append(self.total_bytes) + + # If we get too many small chunks before self._high_water is reached, then any + # .read() call becomes computationally expensive, and could block the event loop + # for too long, hence an additional self._high_water_chunks here. + if ( + len(self._http_chunk_splits) > self._high_water_chunks + and not self._protocol._reading_paused + ): + self._protocol.pause_reading() + + # wake up readchunk when end of http chunk received + waiter = self._waiter + if waiter is not None: + self._waiter = None + set_result(waiter, None) + + async def _wait(self, func_name: str) -> None: + if not self._protocol.connected: + raise RuntimeError("Connection closed.") + + # StreamReader uses a future to link the protocol feed_data() method + # to a read coroutine. Running two read coroutines at the same time + # would have an unexpected behaviour. It would not possible to know + # which coroutine would get the next data. + if self._waiter is not None: + raise RuntimeError( + "%s() called while another coroutine is " + "already waiting for incoming data" % func_name + ) + + waiter = self._waiter = self._loop.create_future() + try: + with self._timer: + await waiter + finally: + self._waiter = None + + async def readline(self) -> bytes: + return await self.readuntil() + + async def readuntil(self, separator: bytes = b"\n") -> bytes: + seplen = len(separator) + if seplen == 0: + raise ValueError("Separator should be at least one-byte string") + + if self._exception is not None: + raise self._exception + + chunk = b"" + chunk_size = 0 + not_enough = True + + while not_enough: + while self._buffer and not_enough: + offset = self._buffer_offset + ichar = self._buffer[0].find(separator, offset) + 1 + # Read from current offset to found separator or to the end. + data = self._read_nowait_chunk( + ichar - offset + seplen - 1 if ichar else -1 + ) + chunk += data + chunk_size += len(data) + if ichar: + not_enough = False + + if chunk_size > self._high_water: + raise ValueError("Chunk too big") + + if self._eof: + break + + if not_enough: + await self._wait("readuntil") + + return chunk + + async def read(self, n: int = -1) -> bytes: + if self._exception is not None: + raise self._exception + + # migration problem; with DataQueue you have to catch + # EofStream exception, so common way is to run payload.read() inside + # infinite loop. what can cause real infinite loop with StreamReader + # lets keep this code one major release. + if __debug__: + if self._eof and not self._buffer: + self._eof_counter = getattr(self, "_eof_counter", 0) + 1 + if self._eof_counter > 5: + internal_logger.warning( + "Multiple access to StreamReader in eof state, " + "might be infinite loop.", + stack_info=True, + ) + + if not n: + return b"" + + if n < 0: + # This used to just loop creating a new waiter hoping to + # collect everything in self._buffer, but that would + # deadlock if the subprocess sends more than self.limit + # bytes. So just call self.readany() until EOF. + blocks = [] + while True: + block = await self.readany() + if not block: + break + blocks.append(block) + return b"".join(blocks) + + # TODO: should be `if` instead of `while` + # because waiter maybe triggered on chunk end, + # without feeding any data + while not self._buffer and not self._eof: + await self._wait("read") + + return self._read_nowait(n) + + async def readany(self) -> bytes: + if self._exception is not None: + raise self._exception + + # TODO: should be `if` instead of `while` + # because waiter maybe triggered on chunk end, + # without feeding any data + while not self._buffer and not self._eof: + await self._wait("readany") + + return self._read_nowait(-1) + + async def readchunk(self) -> Tuple[bytes, bool]: + """Returns a tuple of (data, end_of_http_chunk). + + When chunked transfer + encoding is used, end_of_http_chunk is a boolean indicating if the end + of the data corresponds to the end of a HTTP chunk , otherwise it is + always False. + """ + while True: + if self._exception is not None: + raise self._exception + + while self._http_chunk_splits: + pos = self._http_chunk_splits.popleft() + if pos == self._cursor: + return (b"", True) + if pos > self._cursor: + return (self._read_nowait(pos - self._cursor), True) + internal_logger.warning( + "Skipping HTTP chunk end due to data " + "consumption beyond chunk boundary" + ) + + if self._buffer: + return (self._read_nowait_chunk(-1), False) + # return (self._read_nowait(-1), False) + + if self._eof: + # Special case for signifying EOF. + # (b'', True) is not a final return value actually. + return (b"", False) + + await self._wait("readchunk") + + async def readexactly(self, n: int) -> bytes: + if self._exception is not None: + raise self._exception + + blocks: List[bytes] = [] + while n > 0: + block = await self.read(n) + if not block: + partial = b"".join(blocks) + raise asyncio.IncompleteReadError(partial, len(partial) + n) + blocks.append(block) + n -= len(block) + + return b"".join(blocks) + + def read_nowait(self, n: int = -1) -> bytes: + # default was changed to be consistent with .read(-1) + # + # I believe the most users don't know about the method and + # they are not affected. + if self._exception is not None: + raise self._exception + + if self._waiter and not self._waiter.done(): + raise RuntimeError( + "Called while some coroutine is waiting for incoming data." + ) + + return self._read_nowait(n) + + def _read_nowait_chunk(self, n: int) -> bytes: + first_buffer = self._buffer[0] + offset = self._buffer_offset + if n != -1 and len(first_buffer) - offset > n: + data = first_buffer[offset : offset + n] + self._buffer_offset += n + + elif offset: + self._buffer.popleft() + data = first_buffer[offset:] + self._buffer_offset = 0 + + else: + data = self._buffer.popleft() + + data_len = len(data) + self._size -= data_len + self._cursor += data_len + + chunk_splits = self._http_chunk_splits + # Prevent memory leak: drop useless chunk splits + while chunk_splits and chunk_splits[0] < self._cursor: + chunk_splits.popleft() + + if ( + self._protocol._reading_paused + and self._size < self._low_water + and ( + self._http_chunk_splits is None + or len(self._http_chunk_splits) < self._low_water_chunks + ) + ): + self._protocol.resume_reading() + return data + + def _read_nowait(self, n: int) -> bytes: + """Read not more than n bytes, or whole buffer if n == -1""" + self._timer.assert_timeout() + + chunks = [] + while self._buffer: + chunk = self._read_nowait_chunk(n) + chunks.append(chunk) + if n != -1: + n -= len(chunk) + if n == 0: + break + + return b"".join(chunks) if chunks else b"" + + +class EmptyStreamReader(StreamReader): # lgtm [py/missing-call-to-init] + + __slots__ = ("_read_eof_chunk",) + + def __init__(self) -> None: + self._read_eof_chunk = False + self.total_bytes = 0 + + def __repr__(self) -> str: + return "<%s>" % self.__class__.__name__ + + def exception(self) -> Optional[BaseException]: + return None + + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + pass + + def on_eof(self, callback: Callable[[], None]) -> None: + try: + callback() + except Exception: + internal_logger.exception("Exception in eof callback") + + def feed_eof(self) -> None: + pass + + def is_eof(self) -> bool: + return True + + def at_eof(self) -> bool: + return True + + async def wait_eof(self) -> None: + return + + def feed_data(self, data: bytes, n: int = 0) -> None: + pass + + async def readline(self) -> bytes: + return b"" + + async def read(self, n: int = -1) -> bytes: + return b"" + + # TODO add async def readuntil + + async def readany(self) -> bytes: + return b"" + + async def readchunk(self) -> Tuple[bytes, bool]: + if not self._read_eof_chunk: + self._read_eof_chunk = True + return (b"", False) + + return (b"", True) + + async def readexactly(self, n: int) -> bytes: + raise asyncio.IncompleteReadError(b"", n) + + def read_nowait(self, n: int = -1) -> bytes: + return b"" + + +EMPTY_PAYLOAD: Final[StreamReader] = EmptyStreamReader() + + +class DataQueue(Generic[_T]): + """DataQueue is a general-purpose blocking queue with one reader.""" + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop = loop + self._eof = False + self._waiter: Optional[asyncio.Future[None]] = None + self._exception: Optional[BaseException] = None + self._buffer: Deque[Tuple[_T, int]] = collections.deque() + + def __len__(self) -> int: + return len(self._buffer) + + def is_eof(self) -> bool: + return self._eof + + def at_eof(self) -> bool: + return self._eof and not self._buffer + + def exception(self) -> Optional[BaseException]: + return self._exception + + def set_exception( + self, + exc: BaseException, + exc_cause: BaseException = _EXC_SENTINEL, + ) -> None: + self._eof = True + self._exception = exc + if (waiter := self._waiter) is not None: + self._waiter = None + set_exception(waiter, exc, exc_cause) + + def feed_data(self, data: _T, size: int = 0) -> None: + self._buffer.append((data, size)) + if (waiter := self._waiter) is not None: + self._waiter = None + set_result(waiter, None) + + def feed_eof(self) -> None: + self._eof = True + if (waiter := self._waiter) is not None: + self._waiter = None + set_result(waiter, None) + + async def read(self) -> _T: + if not self._buffer and not self._eof: + assert not self._waiter + self._waiter = self._loop.create_future() + try: + await self._waiter + except (asyncio.CancelledError, asyncio.TimeoutError): + self._waiter = None + raise + if self._buffer: + data, _ = self._buffer.popleft() + return data + if self._exception is not None: + raise self._exception + raise EofStream + + def __aiter__(self) -> AsyncStreamIterator[_T]: + return AsyncStreamIterator(self.read) + + +class FlowControlDataQueue(DataQueue[_T]): + """FlowControlDataQueue resumes and pauses an underlying stream. + + It is a destination for parsed data. + + This class is deprecated and will be removed in version 4.0. + """ + + def __init__( + self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop + ) -> None: + super().__init__(loop=loop) + self._size = 0 + self._protocol = protocol + self._limit = limit * 2 + + def feed_data(self, data: _T, size: int = 0) -> None: + super().feed_data(data, size) + self._size += size + + if self._size > self._limit and not self._protocol._reading_paused: + self._protocol.pause_reading() + + async def read(self) -> _T: + if not self._buffer and not self._eof: + assert not self._waiter + self._waiter = self._loop.create_future() + try: + await self._waiter + except (asyncio.CancelledError, asyncio.TimeoutError): + self._waiter = None + raise + if self._buffer: + data, size = self._buffer.popleft() + self._size -= size + if self._size < self._limit and self._protocol._reading_paused: + self._protocol.resume_reading() + return data + if self._exception is not None: + raise self._exception + raise EofStream diff --git a/lib/python3.12/site-packages/aiohttp/tcp_helpers.py b/lib/python3.12/site-packages/aiohttp/tcp_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..88b244223741ad2decb6cb612eae644fae88b2b2 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/tcp_helpers.py @@ -0,0 +1,37 @@ +"""Helper methods to tune a TCP connection""" + +import asyncio +import socket +from contextlib import suppress +from typing import Optional # noqa + +__all__ = ("tcp_keepalive", "tcp_nodelay") + + +if hasattr(socket, "SO_KEEPALIVE"): + + def tcp_keepalive(transport: asyncio.Transport) -> None: + sock = transport.get_extra_info("socket") + if sock is not None: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + +else: + + def tcp_keepalive(transport: asyncio.Transport) -> None: # pragma: no cover + pass + + +def tcp_nodelay(transport: asyncio.Transport, value: bool) -> None: + sock = transport.get_extra_info("socket") + + if sock is None: + return + + if sock.family not in (socket.AF_INET, socket.AF_INET6): + return + + value = bool(value) + + # socket may be closed already, on windows OSError get raised + with suppress(OSError): + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, value) diff --git a/lib/python3.12/site-packages/aiohttp/test_utils.py b/lib/python3.12/site-packages/aiohttp/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..87c31427867f90244c11c8440a5e1f47fc5a079f --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/test_utils.py @@ -0,0 +1,774 @@ +"""Utilities shared by tests.""" + +import asyncio +import contextlib +import gc +import inspect +import ipaddress +import os +import socket +import sys +import warnings +from abc import ABC, abstractmethod +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Generic, + Iterator, + List, + Optional, + Type, + TypeVar, + cast, + overload, +) +from unittest import IsolatedAsyncioTestCase, mock + +from aiosignal import Signal +from multidict import CIMultiDict, CIMultiDictProxy +from yarl import URL + +import aiohttp +from aiohttp.client import ( + _RequestContextManager, + _RequestOptions, + _WSRequestContextManager, +) + +from . import ClientSession, hdrs +from .abc import AbstractCookieJar +from .client_reqrep import ClientResponse +from .client_ws import ClientWebSocketResponse +from .helpers import sentinel +from .http import HttpVersion, RawRequestMessage +from .streams import EMPTY_PAYLOAD, StreamReader +from .typedefs import StrOrURL +from .web import ( + Application, + AppRunner, + BaseRequest, + BaseRunner, + Request, + Server, + ServerRunner, + SockSite, + UrlMappingMatchInfo, +) +from .web_protocol import _RequestHandler + +if TYPE_CHECKING: + from ssl import SSLContext +else: + SSLContext = None + +if sys.version_info >= (3, 11) and TYPE_CHECKING: + from typing import Unpack + +if sys.version_info >= (3, 11): + from typing import Self +else: + Self = Any + +_ApplicationNone = TypeVar("_ApplicationNone", Application, None) +_Request = TypeVar("_Request", bound=BaseRequest) + +REUSE_ADDRESS = os.name == "posix" and sys.platform != "cygwin" + + +def get_unused_port_socket( + host: str, family: socket.AddressFamily = socket.AF_INET +) -> socket.socket: + return get_port_socket(host, 0, family) + + +def get_port_socket( + host: str, port: int, family: socket.AddressFamily +) -> socket.socket: + s = socket.socket(family, socket.SOCK_STREAM) + if REUSE_ADDRESS: + # Windows has different semantics for SO_REUSEADDR, + # so don't set it. Ref: + # https://docs.microsoft.com/en-us/windows/win32/winsock/using-so-reuseaddr-and-so-exclusiveaddruse + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind((host, port)) + return s + + +def unused_port() -> int: + """Return a port that is unused on the current host.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return cast(int, s.getsockname()[1]) + + +class BaseTestServer(ABC): + __test__ = False + + def __init__( + self, + *, + scheme: str = "", + loop: Optional[asyncio.AbstractEventLoop] = None, + host: str = "127.0.0.1", + port: Optional[int] = None, + skip_url_asserts: bool = False, + socket_factory: Callable[ + [str, int, socket.AddressFamily], socket.socket + ] = get_port_socket, + **kwargs: Any, + ) -> None: + self._loop = loop + self.runner: Optional[BaseRunner] = None + self._root: Optional[URL] = None + self.host = host + self.port = port + self._closed = False + self.scheme = scheme + self.skip_url_asserts = skip_url_asserts + self.socket_factory = socket_factory + + async def start_server( + self, loop: Optional[asyncio.AbstractEventLoop] = None, **kwargs: Any + ) -> None: + if self.runner: + return + self._loop = loop + self._ssl = kwargs.pop("ssl", None) + self.runner = await self._make_runner(handler_cancellation=True, **kwargs) + await self.runner.setup() + if not self.port: + self.port = 0 + absolute_host = self.host + try: + version = ipaddress.ip_address(self.host).version + except ValueError: + version = 4 + if version == 6: + absolute_host = f"[{self.host}]" + family = socket.AF_INET6 if version == 6 else socket.AF_INET + _sock = self.socket_factory(self.host, self.port, family) + self.host, self.port = _sock.getsockname()[:2] + site = SockSite(self.runner, sock=_sock, ssl_context=self._ssl) + await site.start() + server = site._server + assert server is not None + sockets = server.sockets # type: ignore[attr-defined] + assert sockets is not None + self.port = sockets[0].getsockname()[1] + if not self.scheme: + self.scheme = "https" if self._ssl else "http" + self._root = URL(f"{self.scheme}://{absolute_host}:{self.port}") + + @abstractmethod # pragma: no cover + async def _make_runner(self, **kwargs: Any) -> BaseRunner: + pass + + def make_url(self, path: StrOrURL) -> URL: + assert self._root is not None + url = URL(path) + if not self.skip_url_asserts: + assert not url.absolute + return self._root.join(url) + else: + return URL(str(self._root) + str(path)) + + @property + def started(self) -> bool: + return self.runner is not None + + @property + def closed(self) -> bool: + return self._closed + + @property + def handler(self) -> Server: + # for backward compatibility + # web.Server instance + runner = self.runner + assert runner is not None + assert runner.server is not None + return runner.server + + async def close(self) -> None: + """Close all fixtures created by the test client. + + After that point, the TestClient is no longer usable. + + This is an idempotent function: running close multiple times + will not have any additional effects. + + close is also run when the object is garbage collected, and on + exit when used as a context manager. + + """ + if self.started and not self.closed: + assert self.runner is not None + await self.runner.cleanup() + self._root = None + self.port = None + self._closed = True + + def __enter__(self) -> None: + raise TypeError("Use async with instead") + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + # __exit__ should exist in pair with __enter__ but never executed + pass # pragma: no cover + + async def __aenter__(self) -> "BaseTestServer": + await self.start_server(loop=self._loop) + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await self.close() + + +class TestServer(BaseTestServer): + def __init__( + self, + app: Application, + *, + scheme: str = "", + host: str = "127.0.0.1", + port: Optional[int] = None, + **kwargs: Any, + ): + self.app = app + super().__init__(scheme=scheme, host=host, port=port, **kwargs) + + async def _make_runner(self, **kwargs: Any) -> BaseRunner: + return AppRunner(self.app, **kwargs) + + +class RawTestServer(BaseTestServer): + def __init__( + self, + handler: _RequestHandler, + *, + scheme: str = "", + host: str = "127.0.0.1", + port: Optional[int] = None, + **kwargs: Any, + ) -> None: + self._handler = handler + super().__init__(scheme=scheme, host=host, port=port, **kwargs) + + async def _make_runner(self, debug: bool = True, **kwargs: Any) -> ServerRunner: + srv = Server(self._handler, loop=self._loop, debug=debug, **kwargs) + return ServerRunner(srv, debug=debug, **kwargs) + + +class TestClient(Generic[_Request, _ApplicationNone]): + """ + A test client implementation. + + To write functional tests for aiohttp based servers. + + """ + + __test__ = False + + @overload + def __init__( + self: "TestClient[Request, Application]", + server: TestServer, + *, + cookie_jar: Optional[AbstractCookieJar] = None, + **kwargs: Any, + ) -> None: ... + @overload + def __init__( + self: "TestClient[_Request, None]", + server: BaseTestServer, + *, + cookie_jar: Optional[AbstractCookieJar] = None, + **kwargs: Any, + ) -> None: ... + def __init__( + self, + server: BaseTestServer, + *, + cookie_jar: Optional[AbstractCookieJar] = None, + loop: Optional[asyncio.AbstractEventLoop] = None, + **kwargs: Any, + ) -> None: + if not isinstance(server, BaseTestServer): + raise TypeError( + "server must be TestServer instance, found type: %r" % type(server) + ) + self._server = server + self._loop = loop + if cookie_jar is None: + cookie_jar = aiohttp.CookieJar(unsafe=True, loop=loop) + self._session = ClientSession(loop=loop, cookie_jar=cookie_jar, **kwargs) + self._session._retry_connection = False + self._closed = False + self._responses: List[ClientResponse] = [] + self._websockets: List[ClientWebSocketResponse] = [] + + async def start_server(self) -> None: + await self._server.start_server(loop=self._loop) + + @property + def host(self) -> str: + return self._server.host + + @property + def port(self) -> Optional[int]: + return self._server.port + + @property + def server(self) -> BaseTestServer: + return self._server + + @property + def app(self) -> _ApplicationNone: + return getattr(self._server, "app", None) # type: ignore[return-value] + + @property + def session(self) -> ClientSession: + """An internal aiohttp.ClientSession. + + Unlike the methods on the TestClient, client session requests + do not automatically include the host in the url queried, and + will require an absolute path to the resource. + + """ + return self._session + + def make_url(self, path: StrOrURL) -> URL: + return self._server.make_url(path) + + async def _request( + self, method: str, path: StrOrURL, **kwargs: Any + ) -> ClientResponse: + resp = await self._session.request(method, self.make_url(path), **kwargs) + # save it to close later + self._responses.append(resp) + return resp + + if sys.version_info >= (3, 11) and TYPE_CHECKING: + + def request( + self, method: str, path: StrOrURL, **kwargs: Unpack[_RequestOptions] + ) -> _RequestContextManager: ... + + def get( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def options( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def head( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def post( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def put( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def patch( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + def delete( + self, + path: StrOrURL, + **kwargs: Unpack[_RequestOptions], + ) -> _RequestContextManager: ... + + else: + + def request( + self, method: str, path: StrOrURL, **kwargs: Any + ) -> _RequestContextManager: + """Routes a request to tested http server. + + The interface is identical to aiohttp.ClientSession.request, + except the loop kwarg is overridden by the instance used by the + test server. + + """ + return _RequestContextManager(self._request(method, path, **kwargs)) + + def get(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP GET request.""" + return _RequestContextManager(self._request(hdrs.METH_GET, path, **kwargs)) + + def post(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP POST request.""" + return _RequestContextManager(self._request(hdrs.METH_POST, path, **kwargs)) + + def options(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP OPTIONS request.""" + return _RequestContextManager( + self._request(hdrs.METH_OPTIONS, path, **kwargs) + ) + + def head(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP HEAD request.""" + return _RequestContextManager(self._request(hdrs.METH_HEAD, path, **kwargs)) + + def put(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP PUT request.""" + return _RequestContextManager(self._request(hdrs.METH_PUT, path, **kwargs)) + + def patch(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP PATCH request.""" + return _RequestContextManager( + self._request(hdrs.METH_PATCH, path, **kwargs) + ) + + def delete(self, path: StrOrURL, **kwargs: Any) -> _RequestContextManager: + """Perform an HTTP PATCH request.""" + return _RequestContextManager( + self._request(hdrs.METH_DELETE, path, **kwargs) + ) + + def ws_connect(self, path: StrOrURL, **kwargs: Any) -> _WSRequestContextManager: + """Initiate websocket connection. + + The api corresponds to aiohttp.ClientSession.ws_connect. + + """ + return _WSRequestContextManager(self._ws_connect(path, **kwargs)) + + async def _ws_connect( + self, path: StrOrURL, **kwargs: Any + ) -> ClientWebSocketResponse: + ws = await self._session.ws_connect(self.make_url(path), **kwargs) + self._websockets.append(ws) + return ws + + async def close(self) -> None: + """Close all fixtures created by the test client. + + After that point, the TestClient is no longer usable. + + This is an idempotent function: running close multiple times + will not have any additional effects. + + close is also run on exit when used as a(n) (asynchronous) + context manager. + + """ + if not self._closed: + for resp in self._responses: + resp.close() + for ws in self._websockets: + await ws.close() + await self._session.close() + await self._server.close() + self._closed = True + + def __enter__(self) -> None: + raise TypeError("Use async with instead") + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> None: + # __exit__ should exist in pair with __enter__ but never executed + pass # pragma: no cover + + async def __aenter__(self) -> Self: + await self.start_server() + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> None: + await self.close() + + +class AioHTTPTestCase(IsolatedAsyncioTestCase): + """A base class to allow for unittest web applications using aiohttp. + + Provides the following: + + * self.client (aiohttp.test_utils.TestClient): an aiohttp test client. + * self.loop (asyncio.BaseEventLoop): the event loop in which the + application and server are running. + * self.app (aiohttp.web.Application): the application returned by + self.get_application() + + Note that the TestClient's methods are asynchronous: you have to + execute function on the test client using asynchronous methods. + """ + + async def get_application(self) -> Application: + """Get application. + + This method should be overridden + to return the aiohttp.web.Application + object to test. + """ + return self.get_app() + + def get_app(self) -> Application: + """Obsolete method used to constructing web application. + + Use .get_application() coroutine instead. + """ + raise RuntimeError("Did you forget to define get_application()?") + + async def asyncSetUp(self) -> None: + self.loop = asyncio.get_running_loop() + return await self.setUpAsync() + + async def setUpAsync(self) -> None: + self.app = await self.get_application() + self.server = await self.get_server(self.app) + self.client = await self.get_client(self.server) + + await self.client.start_server() + + async def asyncTearDown(self) -> None: + return await self.tearDownAsync() + + async def tearDownAsync(self) -> None: + await self.client.close() + + async def get_server(self, app: Application) -> TestServer: + """Return a TestServer instance.""" + return TestServer(app, loop=self.loop) + + async def get_client(self, server: TestServer) -> TestClient[Request, Application]: + """Return a TestClient instance.""" + return TestClient(server, loop=self.loop) + + +def unittest_run_loop(func: Any, *args: Any, **kwargs: Any) -> Any: + """ + A decorator dedicated to use with asynchronous AioHTTPTestCase test methods. + + In 3.8+, this does nothing. + """ + warnings.warn( + "Decorator `@unittest_run_loop` is no longer needed in aiohttp 3.8+", + DeprecationWarning, + stacklevel=2, + ) + return func + + +_LOOP_FACTORY = Callable[[], asyncio.AbstractEventLoop] + + +@contextlib.contextmanager +def loop_context( + loop_factory: _LOOP_FACTORY = asyncio.new_event_loop, fast: bool = False +) -> Iterator[asyncio.AbstractEventLoop]: + """A contextmanager that creates an event_loop, for test purposes. + + Handles the creation and cleanup of a test loop. + """ + loop = setup_test_loop(loop_factory) + yield loop + teardown_test_loop(loop, fast=fast) + + +def setup_test_loop( + loop_factory: _LOOP_FACTORY = asyncio.new_event_loop, +) -> asyncio.AbstractEventLoop: + """Create and return an asyncio.BaseEventLoop instance. + + The caller should also call teardown_test_loop, + once they are done with the loop. + """ + loop = loop_factory() + asyncio.set_event_loop(loop) + return loop + + +def teardown_test_loop(loop: asyncio.AbstractEventLoop, fast: bool = False) -> None: + """Teardown and cleanup an event_loop created by setup_test_loop.""" + closed = loop.is_closed() + if not closed: + loop.call_soon(loop.stop) + loop.run_forever() + loop.close() + + if not fast: + gc.collect() + + asyncio.set_event_loop(None) + + +def _create_app_mock() -> mock.MagicMock: + def get_dict(app: Any, key: str) -> Any: + return app.__app_dict[key] + + def set_dict(app: Any, key: str, value: Any) -> None: + app.__app_dict[key] = value + + app = mock.MagicMock(spec=Application) + app.__app_dict = {} + app.__getitem__ = get_dict + app.__setitem__ = set_dict + + app._debug = False + app.on_response_prepare = Signal(app) + app.on_response_prepare.freeze() + return app + + +def _create_transport(sslcontext: Optional[SSLContext] = None) -> mock.Mock: + transport = mock.Mock() + + def get_extra_info(key: str) -> Optional[SSLContext]: + if key == "sslcontext": + return sslcontext + else: + return None + + transport.get_extra_info.side_effect = get_extra_info + return transport + + +def make_mocked_request( + method: str, + path: str, + headers: Any = None, + *, + match_info: Any = sentinel, + version: HttpVersion = HttpVersion(1, 1), + closing: bool = False, + app: Any = None, + writer: Any = sentinel, + protocol: Any = sentinel, + transport: Any = sentinel, + payload: StreamReader = EMPTY_PAYLOAD, + sslcontext: Optional[SSLContext] = None, + client_max_size: int = 1024**2, + loop: Any = ..., +) -> Request: + """Creates mocked web.Request testing purposes. + + Useful in unit tests, when spinning full web server is overkill or + specific conditions and errors are hard to trigger. + """ + task = mock.Mock() + if loop is ...: + # no loop passed, try to get the current one if + # its is running as we need a real loop to create + # executor jobs to be able to do testing + # with a real executor + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = mock.Mock() + loop.create_future.return_value = () + + if version < HttpVersion(1, 1): + closing = True + + if headers: + headers = CIMultiDictProxy(CIMultiDict(headers)) + raw_hdrs = tuple( + (k.encode("utf-8"), v.encode("utf-8")) for k, v in headers.items() + ) + else: + headers = CIMultiDictProxy(CIMultiDict()) + raw_hdrs = () + + chunked = "chunked" in headers.get(hdrs.TRANSFER_ENCODING, "").lower() + + message = RawRequestMessage( + method, + path, + version, + headers, + raw_hdrs, + closing, + None, + False, + chunked, + URL(path), + ) + if app is None: + app = _create_app_mock() + + if transport is sentinel: + transport = _create_transport(sslcontext) + + if protocol is sentinel: + protocol = mock.Mock() + protocol.transport = transport + type(protocol).peername = mock.PropertyMock( + return_value=transport.get_extra_info("peername") + ) + type(protocol).ssl_context = mock.PropertyMock(return_value=sslcontext) + + if writer is sentinel: + writer = mock.Mock() + writer.write_headers = make_mocked_coro(None) + writer.write = make_mocked_coro(None) + writer.write_eof = make_mocked_coro(None) + writer.drain = make_mocked_coro(None) + writer.transport = transport + + protocol.transport = transport + protocol.writer = writer + + req = Request( + message, payload, protocol, writer, task, loop, client_max_size=client_max_size + ) + + match_info = UrlMappingMatchInfo( + {} if match_info is sentinel else match_info, mock.Mock() + ) + match_info.add_app(app) + req._match_info = match_info + + return req + + +def make_mocked_coro( + return_value: Any = sentinel, raise_exception: Any = sentinel +) -> Any: + """Creates a coroutine mock.""" + + async def mock_coro(*args: Any, **kwargs: Any) -> Any: + if raise_exception is not sentinel: + raise raise_exception + if not inspect.isawaitable(return_value): + return return_value + await return_value + + return mock.Mock(wraps=mock_coro) diff --git a/lib/python3.12/site-packages/aiohttp/tracing.py b/lib/python3.12/site-packages/aiohttp/tracing.py new file mode 100644 index 0000000000000000000000000000000000000000..568fa7f9e38090e0c0a4738db4d51656ce31b99a --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/tracing.py @@ -0,0 +1,455 @@ +from types import SimpleNamespace +from typing import TYPE_CHECKING, Mapping, Optional, Type, TypeVar + +import attr +from aiosignal import Signal +from multidict import CIMultiDict +from yarl import URL + +from .client_reqrep import ClientResponse + +if TYPE_CHECKING: + from .client import ClientSession + + _ParamT_contra = TypeVar("_ParamT_contra", contravariant=True) + _TracingSignal = Signal[ClientSession, SimpleNamespace, _ParamT_contra] + + +__all__ = ( + "TraceConfig", + "TraceRequestStartParams", + "TraceRequestEndParams", + "TraceRequestExceptionParams", + "TraceConnectionQueuedStartParams", + "TraceConnectionQueuedEndParams", + "TraceConnectionCreateStartParams", + "TraceConnectionCreateEndParams", + "TraceConnectionReuseconnParams", + "TraceDnsResolveHostStartParams", + "TraceDnsResolveHostEndParams", + "TraceDnsCacheHitParams", + "TraceDnsCacheMissParams", + "TraceRequestRedirectParams", + "TraceRequestChunkSentParams", + "TraceResponseChunkReceivedParams", + "TraceRequestHeadersSentParams", +) + + +class TraceConfig: + """First-class used to trace requests launched via ClientSession objects.""" + + def __init__( + self, trace_config_ctx_factory: Type[SimpleNamespace] = SimpleNamespace + ) -> None: + self._on_request_start: _TracingSignal[TraceRequestStartParams] = Signal(self) + self._on_request_chunk_sent: _TracingSignal[TraceRequestChunkSentParams] = ( + Signal(self) + ) + self._on_response_chunk_received: _TracingSignal[ + TraceResponseChunkReceivedParams + ] = Signal(self) + self._on_request_end: _TracingSignal[TraceRequestEndParams] = Signal(self) + self._on_request_exception: _TracingSignal[TraceRequestExceptionParams] = ( + Signal(self) + ) + self._on_request_redirect: _TracingSignal[TraceRequestRedirectParams] = Signal( + self + ) + self._on_connection_queued_start: _TracingSignal[ + TraceConnectionQueuedStartParams + ] = Signal(self) + self._on_connection_queued_end: _TracingSignal[ + TraceConnectionQueuedEndParams + ] = Signal(self) + self._on_connection_create_start: _TracingSignal[ + TraceConnectionCreateStartParams + ] = Signal(self) + self._on_connection_create_end: _TracingSignal[ + TraceConnectionCreateEndParams + ] = Signal(self) + self._on_connection_reuseconn: _TracingSignal[ + TraceConnectionReuseconnParams + ] = Signal(self) + self._on_dns_resolvehost_start: _TracingSignal[ + TraceDnsResolveHostStartParams + ] = Signal(self) + self._on_dns_resolvehost_end: _TracingSignal[TraceDnsResolveHostEndParams] = ( + Signal(self) + ) + self._on_dns_cache_hit: _TracingSignal[TraceDnsCacheHitParams] = Signal(self) + self._on_dns_cache_miss: _TracingSignal[TraceDnsCacheMissParams] = Signal(self) + self._on_request_headers_sent: _TracingSignal[TraceRequestHeadersSentParams] = ( + Signal(self) + ) + + self._trace_config_ctx_factory = trace_config_ctx_factory + + def trace_config_ctx( + self, trace_request_ctx: Optional[Mapping[str, str]] = None + ) -> SimpleNamespace: + """Return a new trace_config_ctx instance""" + return self._trace_config_ctx_factory(trace_request_ctx=trace_request_ctx) + + def freeze(self) -> None: + self._on_request_start.freeze() + self._on_request_chunk_sent.freeze() + self._on_response_chunk_received.freeze() + self._on_request_end.freeze() + self._on_request_exception.freeze() + self._on_request_redirect.freeze() + self._on_connection_queued_start.freeze() + self._on_connection_queued_end.freeze() + self._on_connection_create_start.freeze() + self._on_connection_create_end.freeze() + self._on_connection_reuseconn.freeze() + self._on_dns_resolvehost_start.freeze() + self._on_dns_resolvehost_end.freeze() + self._on_dns_cache_hit.freeze() + self._on_dns_cache_miss.freeze() + self._on_request_headers_sent.freeze() + + @property + def on_request_start(self) -> "_TracingSignal[TraceRequestStartParams]": + return self._on_request_start + + @property + def on_request_chunk_sent( + self, + ) -> "_TracingSignal[TraceRequestChunkSentParams]": + return self._on_request_chunk_sent + + @property + def on_response_chunk_received( + self, + ) -> "_TracingSignal[TraceResponseChunkReceivedParams]": + return self._on_response_chunk_received + + @property + def on_request_end(self) -> "_TracingSignal[TraceRequestEndParams]": + return self._on_request_end + + @property + def on_request_exception( + self, + ) -> "_TracingSignal[TraceRequestExceptionParams]": + return self._on_request_exception + + @property + def on_request_redirect( + self, + ) -> "_TracingSignal[TraceRequestRedirectParams]": + return self._on_request_redirect + + @property + def on_connection_queued_start( + self, + ) -> "_TracingSignal[TraceConnectionQueuedStartParams]": + return self._on_connection_queued_start + + @property + def on_connection_queued_end( + self, + ) -> "_TracingSignal[TraceConnectionQueuedEndParams]": + return self._on_connection_queued_end + + @property + def on_connection_create_start( + self, + ) -> "_TracingSignal[TraceConnectionCreateStartParams]": + return self._on_connection_create_start + + @property + def on_connection_create_end( + self, + ) -> "_TracingSignal[TraceConnectionCreateEndParams]": + return self._on_connection_create_end + + @property + def on_connection_reuseconn( + self, + ) -> "_TracingSignal[TraceConnectionReuseconnParams]": + return self._on_connection_reuseconn + + @property + def on_dns_resolvehost_start( + self, + ) -> "_TracingSignal[TraceDnsResolveHostStartParams]": + return self._on_dns_resolvehost_start + + @property + def on_dns_resolvehost_end( + self, + ) -> "_TracingSignal[TraceDnsResolveHostEndParams]": + return self._on_dns_resolvehost_end + + @property + def on_dns_cache_hit(self) -> "_TracingSignal[TraceDnsCacheHitParams]": + return self._on_dns_cache_hit + + @property + def on_dns_cache_miss(self) -> "_TracingSignal[TraceDnsCacheMissParams]": + return self._on_dns_cache_miss + + @property + def on_request_headers_sent( + self, + ) -> "_TracingSignal[TraceRequestHeadersSentParams]": + return self._on_request_headers_sent + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestStartParams: + """Parameters sent by the `on_request_start` signal""" + + method: str + url: URL + headers: "CIMultiDict[str]" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestChunkSentParams: + """Parameters sent by the `on_request_chunk_sent` signal""" + + method: str + url: URL + chunk: bytes + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceResponseChunkReceivedParams: + """Parameters sent by the `on_response_chunk_received` signal""" + + method: str + url: URL + chunk: bytes + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestEndParams: + """Parameters sent by the `on_request_end` signal""" + + method: str + url: URL + headers: "CIMultiDict[str]" + response: ClientResponse + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestExceptionParams: + """Parameters sent by the `on_request_exception` signal""" + + method: str + url: URL + headers: "CIMultiDict[str]" + exception: BaseException + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestRedirectParams: + """Parameters sent by the `on_request_redirect` signal""" + + method: str + url: URL + headers: "CIMultiDict[str]" + response: ClientResponse + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceConnectionQueuedStartParams: + """Parameters sent by the `on_connection_queued_start` signal""" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceConnectionQueuedEndParams: + """Parameters sent by the `on_connection_queued_end` signal""" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceConnectionCreateStartParams: + """Parameters sent by the `on_connection_create_start` signal""" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceConnectionCreateEndParams: + """Parameters sent by the `on_connection_create_end` signal""" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceConnectionReuseconnParams: + """Parameters sent by the `on_connection_reuseconn` signal""" + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceDnsResolveHostStartParams: + """Parameters sent by the `on_dns_resolvehost_start` signal""" + + host: str + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceDnsResolveHostEndParams: + """Parameters sent by the `on_dns_resolvehost_end` signal""" + + host: str + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceDnsCacheHitParams: + """Parameters sent by the `on_dns_cache_hit` signal""" + + host: str + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceDnsCacheMissParams: + """Parameters sent by the `on_dns_cache_miss` signal""" + + host: str + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class TraceRequestHeadersSentParams: + """Parameters sent by the `on_request_headers_sent` signal""" + + method: str + url: URL + headers: "CIMultiDict[str]" + + +class Trace: + """Internal dependency holder class. + + Used to keep together the main dependencies used + at the moment of send a signal. + """ + + def __init__( + self, + session: "ClientSession", + trace_config: TraceConfig, + trace_config_ctx: SimpleNamespace, + ) -> None: + self._trace_config = trace_config + self._trace_config_ctx = trace_config_ctx + self._session = session + + async def send_request_start( + self, method: str, url: URL, headers: "CIMultiDict[str]" + ) -> None: + return await self._trace_config.on_request_start.send( + self._session, + self._trace_config_ctx, + TraceRequestStartParams(method, url, headers), + ) + + async def send_request_chunk_sent( + self, method: str, url: URL, chunk: bytes + ) -> None: + return await self._trace_config.on_request_chunk_sent.send( + self._session, + self._trace_config_ctx, + TraceRequestChunkSentParams(method, url, chunk), + ) + + async def send_response_chunk_received( + self, method: str, url: URL, chunk: bytes + ) -> None: + return await self._trace_config.on_response_chunk_received.send( + self._session, + self._trace_config_ctx, + TraceResponseChunkReceivedParams(method, url, chunk), + ) + + async def send_request_end( + self, + method: str, + url: URL, + headers: "CIMultiDict[str]", + response: ClientResponse, + ) -> None: + return await self._trace_config.on_request_end.send( + self._session, + self._trace_config_ctx, + TraceRequestEndParams(method, url, headers, response), + ) + + async def send_request_exception( + self, + method: str, + url: URL, + headers: "CIMultiDict[str]", + exception: BaseException, + ) -> None: + return await self._trace_config.on_request_exception.send( + self._session, + self._trace_config_ctx, + TraceRequestExceptionParams(method, url, headers, exception), + ) + + async def send_request_redirect( + self, + method: str, + url: URL, + headers: "CIMultiDict[str]", + response: ClientResponse, + ) -> None: + return await self._trace_config._on_request_redirect.send( + self._session, + self._trace_config_ctx, + TraceRequestRedirectParams(method, url, headers, response), + ) + + async def send_connection_queued_start(self) -> None: + return await self._trace_config.on_connection_queued_start.send( + self._session, self._trace_config_ctx, TraceConnectionQueuedStartParams() + ) + + async def send_connection_queued_end(self) -> None: + return await self._trace_config.on_connection_queued_end.send( + self._session, self._trace_config_ctx, TraceConnectionQueuedEndParams() + ) + + async def send_connection_create_start(self) -> None: + return await self._trace_config.on_connection_create_start.send( + self._session, self._trace_config_ctx, TraceConnectionCreateStartParams() + ) + + async def send_connection_create_end(self) -> None: + return await self._trace_config.on_connection_create_end.send( + self._session, self._trace_config_ctx, TraceConnectionCreateEndParams() + ) + + async def send_connection_reuseconn(self) -> None: + return await self._trace_config.on_connection_reuseconn.send( + self._session, self._trace_config_ctx, TraceConnectionReuseconnParams() + ) + + async def send_dns_resolvehost_start(self, host: str) -> None: + return await self._trace_config.on_dns_resolvehost_start.send( + self._session, self._trace_config_ctx, TraceDnsResolveHostStartParams(host) + ) + + async def send_dns_resolvehost_end(self, host: str) -> None: + return await self._trace_config.on_dns_resolvehost_end.send( + self._session, self._trace_config_ctx, TraceDnsResolveHostEndParams(host) + ) + + async def send_dns_cache_hit(self, host: str) -> None: + return await self._trace_config.on_dns_cache_hit.send( + self._session, self._trace_config_ctx, TraceDnsCacheHitParams(host) + ) + + async def send_dns_cache_miss(self, host: str) -> None: + return await self._trace_config.on_dns_cache_miss.send( + self._session, self._trace_config_ctx, TraceDnsCacheMissParams(host) + ) + + async def send_request_headers( + self, method: str, url: URL, headers: "CIMultiDict[str]" + ) -> None: + return await self._trace_config._on_request_headers_sent.send( + self._session, + self._trace_config_ctx, + TraceRequestHeadersSentParams(method, url, headers), + ) diff --git a/lib/python3.12/site-packages/aiohttp/typedefs.py b/lib/python3.12/site-packages/aiohttp/typedefs.py new file mode 100644 index 0000000000000000000000000000000000000000..cc8c0825b4e522f7d1b6cf0058564f322fb5a905 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/typedefs.py @@ -0,0 +1,69 @@ +import json +import os +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Iterable, + Mapping, + Protocol, + Tuple, + Union, +) + +from multidict import CIMultiDict, CIMultiDictProxy, MultiDict, MultiDictProxy, istr +from yarl import URL, Query as _Query + +Query = _Query + +DEFAULT_JSON_ENCODER = json.dumps +DEFAULT_JSON_DECODER = json.loads + +if TYPE_CHECKING: + _CIMultiDict = CIMultiDict[str] + _CIMultiDictProxy = CIMultiDictProxy[str] + _MultiDict = MultiDict[str] + _MultiDictProxy = MultiDictProxy[str] + from http.cookies import BaseCookie, Morsel + + from .web import Request, StreamResponse +else: + _CIMultiDict = CIMultiDict + _CIMultiDictProxy = CIMultiDictProxy + _MultiDict = MultiDict + _MultiDictProxy = MultiDictProxy + +Byteish = Union[bytes, bytearray, memoryview] +JSONEncoder = Callable[[Any], str] +JSONDecoder = Callable[[str], Any] +LooseHeaders = Union[ + Mapping[str, str], + Mapping[istr, str], + _CIMultiDict, + _CIMultiDictProxy, + Iterable[Tuple[Union[str, istr], str]], +] +RawHeaders = Tuple[Tuple[bytes, bytes], ...] +StrOrURL = Union[str, URL] + +LooseCookiesMappings = Mapping[str, Union[str, "BaseCookie[str]", "Morsel[Any]"]] +LooseCookiesIterables = Iterable[ + Tuple[str, Union[str, "BaseCookie[str]", "Morsel[Any]"]] +] +LooseCookies = Union[ + LooseCookiesMappings, + LooseCookiesIterables, + "BaseCookie[str]", +] + +Handler = Callable[["Request"], Awaitable["StreamResponse"]] + + +class Middleware(Protocol): + def __call__( + self, request: "Request", handler: Handler + ) -> Awaitable["StreamResponse"]: ... + + +PathLike = Union[str, "os.PathLike[str]"] diff --git a/lib/python3.12/site-packages/aiohttp/web.py b/lib/python3.12/site-packages/aiohttp/web.py new file mode 100644 index 0000000000000000000000000000000000000000..5a1fc964172be24b03ffaf3982ba6fccea7b6515 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web.py @@ -0,0 +1,592 @@ +import asyncio +import logging +import os +import socket +import sys +import warnings +from argparse import ArgumentParser +from collections.abc import Iterable +from contextlib import suppress +from importlib import import_module +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Iterable as TypingIterable, + List, + Optional, + Set, + Type, + Union, + cast, +) + +from .abc import AbstractAccessLogger +from .helpers import AppKey as AppKey +from .log import access_logger +from .typedefs import PathLike +from .web_app import Application as Application, CleanupError as CleanupError +from .web_exceptions import ( + HTTPAccepted as HTTPAccepted, + HTTPBadGateway as HTTPBadGateway, + HTTPBadRequest as HTTPBadRequest, + HTTPClientError as HTTPClientError, + HTTPConflict as HTTPConflict, + HTTPCreated as HTTPCreated, + HTTPError as HTTPError, + HTTPException as HTTPException, + HTTPExpectationFailed as HTTPExpectationFailed, + HTTPFailedDependency as HTTPFailedDependency, + HTTPForbidden as HTTPForbidden, + HTTPFound as HTTPFound, + HTTPGatewayTimeout as HTTPGatewayTimeout, + HTTPGone as HTTPGone, + HTTPInsufficientStorage as HTTPInsufficientStorage, + HTTPInternalServerError as HTTPInternalServerError, + HTTPLengthRequired as HTTPLengthRequired, + HTTPMethodNotAllowed as HTTPMethodNotAllowed, + HTTPMisdirectedRequest as HTTPMisdirectedRequest, + HTTPMove as HTTPMove, + HTTPMovedPermanently as HTTPMovedPermanently, + HTTPMultipleChoices as HTTPMultipleChoices, + HTTPNetworkAuthenticationRequired as HTTPNetworkAuthenticationRequired, + HTTPNoContent as HTTPNoContent, + HTTPNonAuthoritativeInformation as HTTPNonAuthoritativeInformation, + HTTPNotAcceptable as HTTPNotAcceptable, + HTTPNotExtended as HTTPNotExtended, + HTTPNotFound as HTTPNotFound, + HTTPNotImplemented as HTTPNotImplemented, + HTTPNotModified as HTTPNotModified, + HTTPOk as HTTPOk, + HTTPPartialContent as HTTPPartialContent, + HTTPPaymentRequired as HTTPPaymentRequired, + HTTPPermanentRedirect as HTTPPermanentRedirect, + HTTPPreconditionFailed as HTTPPreconditionFailed, + HTTPPreconditionRequired as HTTPPreconditionRequired, + HTTPProxyAuthenticationRequired as HTTPProxyAuthenticationRequired, + HTTPRedirection as HTTPRedirection, + HTTPRequestEntityTooLarge as HTTPRequestEntityTooLarge, + HTTPRequestHeaderFieldsTooLarge as HTTPRequestHeaderFieldsTooLarge, + HTTPRequestRangeNotSatisfiable as HTTPRequestRangeNotSatisfiable, + HTTPRequestTimeout as HTTPRequestTimeout, + HTTPRequestURITooLong as HTTPRequestURITooLong, + HTTPResetContent as HTTPResetContent, + HTTPSeeOther as HTTPSeeOther, + HTTPServerError as HTTPServerError, + HTTPServiceUnavailable as HTTPServiceUnavailable, + HTTPSuccessful as HTTPSuccessful, + HTTPTemporaryRedirect as HTTPTemporaryRedirect, + HTTPTooManyRequests as HTTPTooManyRequests, + HTTPUnauthorized as HTTPUnauthorized, + HTTPUnavailableForLegalReasons as HTTPUnavailableForLegalReasons, + HTTPUnprocessableEntity as HTTPUnprocessableEntity, + HTTPUnsupportedMediaType as HTTPUnsupportedMediaType, + HTTPUpgradeRequired as HTTPUpgradeRequired, + HTTPUseProxy as HTTPUseProxy, + HTTPVariantAlsoNegotiates as HTTPVariantAlsoNegotiates, + HTTPVersionNotSupported as HTTPVersionNotSupported, + NotAppKeyWarning as NotAppKeyWarning, +) +from .web_fileresponse import FileResponse as FileResponse +from .web_log import AccessLogger +from .web_middlewares import ( + middleware as middleware, + normalize_path_middleware as normalize_path_middleware, +) +from .web_protocol import ( + PayloadAccessError as PayloadAccessError, + RequestHandler as RequestHandler, + RequestPayloadError as RequestPayloadError, +) +from .web_request import ( + BaseRequest as BaseRequest, + FileField as FileField, + Request as Request, +) +from .web_response import ( + ContentCoding as ContentCoding, + Response as Response, + StreamResponse as StreamResponse, + json_response as json_response, +) +from .web_routedef import ( + AbstractRouteDef as AbstractRouteDef, + RouteDef as RouteDef, + RouteTableDef as RouteTableDef, + StaticDef as StaticDef, + delete as delete, + get as get, + head as head, + options as options, + patch as patch, + post as post, + put as put, + route as route, + static as static, + view as view, +) +from .web_runner import ( + AppRunner as AppRunner, + BaseRunner as BaseRunner, + BaseSite as BaseSite, + GracefulExit as GracefulExit, + NamedPipeSite as NamedPipeSite, + ServerRunner as ServerRunner, + SockSite as SockSite, + TCPSite as TCPSite, + UnixSite as UnixSite, +) +from .web_server import Server as Server +from .web_urldispatcher import ( + AbstractResource as AbstractResource, + AbstractRoute as AbstractRoute, + DynamicResource as DynamicResource, + PlainResource as PlainResource, + PrefixedSubAppResource as PrefixedSubAppResource, + Resource as Resource, + ResourceRoute as ResourceRoute, + StaticResource as StaticResource, + UrlDispatcher as UrlDispatcher, + UrlMappingMatchInfo as UrlMappingMatchInfo, + View as View, +) +from .web_ws import ( + WebSocketReady as WebSocketReady, + WebSocketResponse as WebSocketResponse, + WSMsgType as WSMsgType, +) + +__all__ = ( + # web_app + "AppKey", + "Application", + "CleanupError", + # web_exceptions + "NotAppKeyWarning", + "HTTPAccepted", + "HTTPBadGateway", + "HTTPBadRequest", + "HTTPClientError", + "HTTPConflict", + "HTTPCreated", + "HTTPError", + "HTTPException", + "HTTPExpectationFailed", + "HTTPFailedDependency", + "HTTPForbidden", + "HTTPFound", + "HTTPGatewayTimeout", + "HTTPGone", + "HTTPInsufficientStorage", + "HTTPInternalServerError", + "HTTPLengthRequired", + "HTTPMethodNotAllowed", + "HTTPMisdirectedRequest", + "HTTPMove", + "HTTPMovedPermanently", + "HTTPMultipleChoices", + "HTTPNetworkAuthenticationRequired", + "HTTPNoContent", + "HTTPNonAuthoritativeInformation", + "HTTPNotAcceptable", + "HTTPNotExtended", + "HTTPNotFound", + "HTTPNotImplemented", + "HTTPNotModified", + "HTTPOk", + "HTTPPartialContent", + "HTTPPaymentRequired", + "HTTPPermanentRedirect", + "HTTPPreconditionFailed", + "HTTPPreconditionRequired", + "HTTPProxyAuthenticationRequired", + "HTTPRedirection", + "HTTPRequestEntityTooLarge", + "HTTPRequestHeaderFieldsTooLarge", + "HTTPRequestRangeNotSatisfiable", + "HTTPRequestTimeout", + "HTTPRequestURITooLong", + "HTTPResetContent", + "HTTPSeeOther", + "HTTPServerError", + "HTTPServiceUnavailable", + "HTTPSuccessful", + "HTTPTemporaryRedirect", + "HTTPTooManyRequests", + "HTTPUnauthorized", + "HTTPUnavailableForLegalReasons", + "HTTPUnprocessableEntity", + "HTTPUnsupportedMediaType", + "HTTPUpgradeRequired", + "HTTPUseProxy", + "HTTPVariantAlsoNegotiates", + "HTTPVersionNotSupported", + # web_fileresponse + "FileResponse", + # web_middlewares + "middleware", + "normalize_path_middleware", + # web_protocol + "PayloadAccessError", + "RequestHandler", + "RequestPayloadError", + # web_request + "BaseRequest", + "FileField", + "Request", + # web_response + "ContentCoding", + "Response", + "StreamResponse", + "json_response", + # web_routedef + "AbstractRouteDef", + "RouteDef", + "RouteTableDef", + "StaticDef", + "delete", + "get", + "head", + "options", + "patch", + "post", + "put", + "route", + "static", + "view", + # web_runner + "AppRunner", + "BaseRunner", + "BaseSite", + "GracefulExit", + "ServerRunner", + "SockSite", + "TCPSite", + "UnixSite", + "NamedPipeSite", + # web_server + "Server", + # web_urldispatcher + "AbstractResource", + "AbstractRoute", + "DynamicResource", + "PlainResource", + "PrefixedSubAppResource", + "Resource", + "ResourceRoute", + "StaticResource", + "UrlDispatcher", + "UrlMappingMatchInfo", + "View", + # web_ws + "WebSocketReady", + "WebSocketResponse", + "WSMsgType", + # web + "run_app", +) + + +if TYPE_CHECKING: + from ssl import SSLContext +else: + try: + from ssl import SSLContext + except ImportError: # pragma: no cover + SSLContext = object # type: ignore[misc,assignment] + +# Only display warning when using -Wdefault, -We, -X dev or similar. +warnings.filterwarnings("ignore", category=NotAppKeyWarning, append=True) + +HostSequence = TypingIterable[str] + + +async def _run_app( + app: Union[Application, Awaitable[Application]], + *, + host: Optional[Union[str, HostSequence]] = None, + port: Optional[int] = None, + path: Union[PathLike, TypingIterable[PathLike], None] = None, + sock: Optional[Union[socket.socket, TypingIterable[socket.socket]]] = None, + ssl_context: Optional[SSLContext] = None, + print: Optional[Callable[..., None]] = print, + backlog: int = 128, + reuse_address: Optional[bool] = None, + reuse_port: Optional[bool] = None, + **kwargs: Any, # TODO(PY311): Use Unpack +) -> None: + # An internal function to actually do all dirty job for application running + if asyncio.iscoroutine(app): + app = await app + + app = cast(Application, app) + + runner = AppRunner(app, **kwargs) + + await runner.setup() + + sites: List[BaseSite] = [] + + try: + if host is not None: + if isinstance(host, str): + sites.append( + TCPSite( + runner, + host, + port, + ssl_context=ssl_context, + backlog=backlog, + reuse_address=reuse_address, + reuse_port=reuse_port, + ) + ) + else: + for h in host: + sites.append( + TCPSite( + runner, + h, + port, + ssl_context=ssl_context, + backlog=backlog, + reuse_address=reuse_address, + reuse_port=reuse_port, + ) + ) + elif path is None and sock is None or port is not None: + sites.append( + TCPSite( + runner, + port=port, + ssl_context=ssl_context, + backlog=backlog, + reuse_address=reuse_address, + reuse_port=reuse_port, + ) + ) + + if path is not None: + if isinstance(path, (str, os.PathLike)): + sites.append( + UnixSite( + runner, + path, + ssl_context=ssl_context, + backlog=backlog, + ) + ) + else: + for p in path: + sites.append( + UnixSite( + runner, + p, + ssl_context=ssl_context, + backlog=backlog, + ) + ) + + if sock is not None: + if not isinstance(sock, Iterable): + sites.append( + SockSite( + runner, + sock, + ssl_context=ssl_context, + backlog=backlog, + ) + ) + else: + for s in sock: + sites.append( + SockSite( + runner, + s, + ssl_context=ssl_context, + backlog=backlog, + ) + ) + for site in sites: + await site.start() + + if print: # pragma: no branch + names = sorted(str(s.name) for s in runner.sites) + print( + "======== Running on {} ========\n" + "(Press CTRL+C to quit)".format(", ".join(names)) + ) + + # sleep forever by 1 hour intervals, + while True: + await asyncio.sleep(3600) + finally: + await runner.cleanup() + + +def _cancel_tasks( + to_cancel: Set["asyncio.Task[Any]"], loop: asyncio.AbstractEventLoop +) -> None: + if not to_cancel: + return + + for task in to_cancel: + task.cancel() + + loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True)) + + for task in to_cancel: + if task.cancelled(): + continue + if task.exception() is not None: + loop.call_exception_handler( + { + "message": "unhandled exception during asyncio.run() shutdown", + "exception": task.exception(), + "task": task, + } + ) + + +def run_app( + app: Union[Application, Awaitable[Application]], + *, + host: Optional[Union[str, HostSequence]] = None, + port: Optional[int] = None, + path: Union[PathLike, TypingIterable[PathLike], None] = None, + sock: Optional[Union[socket.socket, TypingIterable[socket.socket]]] = None, + shutdown_timeout: float = 60.0, + keepalive_timeout: float = 75.0, + ssl_context: Optional[SSLContext] = None, + print: Optional[Callable[..., None]] = print, + backlog: int = 128, + access_log_class: Type[AbstractAccessLogger] = AccessLogger, + access_log_format: str = AccessLogger.LOG_FORMAT, + access_log: Optional[logging.Logger] = access_logger, + handle_signals: bool = True, + reuse_address: Optional[bool] = None, + reuse_port: Optional[bool] = None, + handler_cancellation: bool = False, + loop: Optional[asyncio.AbstractEventLoop] = None, + **kwargs: Any, +) -> None: + """Run an app locally""" + if loop is None: + loop = asyncio.new_event_loop() + + # Configure if and only if in debugging mode and using the default logger + if loop.get_debug() and access_log and access_log.name == "aiohttp.access": + if access_log.level == logging.NOTSET: + access_log.setLevel(logging.DEBUG) + if not access_log.hasHandlers(): + access_log.addHandler(logging.StreamHandler()) + + main_task = loop.create_task( + _run_app( + app, + host=host, + port=port, + path=path, + sock=sock, + shutdown_timeout=shutdown_timeout, + keepalive_timeout=keepalive_timeout, + ssl_context=ssl_context, + print=print, + backlog=backlog, + access_log_class=access_log_class, + access_log_format=access_log_format, + access_log=access_log, + handle_signals=handle_signals, + reuse_address=reuse_address, + reuse_port=reuse_port, + handler_cancellation=handler_cancellation, + **kwargs, + ) + ) + + try: + asyncio.set_event_loop(loop) + loop.run_until_complete(main_task) + except (GracefulExit, KeyboardInterrupt): # pragma: no cover + pass + finally: + try: + main_task.cancel() + with suppress(asyncio.CancelledError): + loop.run_until_complete(main_task) + finally: + _cancel_tasks(asyncio.all_tasks(loop), loop) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +def main(argv: List[str]) -> None: + arg_parser = ArgumentParser( + description="aiohttp.web Application server", prog="aiohttp.web" + ) + arg_parser.add_argument( + "entry_func", + help=( + "Callable returning the `aiohttp.web.Application` instance to " + "run. Should be specified in the 'module:function' syntax." + ), + metavar="entry-func", + ) + arg_parser.add_argument( + "-H", + "--hostname", + help="TCP/IP hostname to serve on (default: localhost)", + default=None, + ) + arg_parser.add_argument( + "-P", + "--port", + help="TCP/IP port to serve on (default: %(default)r)", + type=int, + default=8080, + ) + arg_parser.add_argument( + "-U", + "--path", + help="Unix file system path to serve on. Can be combined with hostname " + "to serve on both Unix and TCP.", + ) + args, extra_argv = arg_parser.parse_known_args(argv) + + # Import logic + mod_str, _, func_str = args.entry_func.partition(":") + if not func_str or not mod_str: + arg_parser.error("'entry-func' not in 'module:function' syntax") + if mod_str.startswith("."): + arg_parser.error("relative module names not supported") + try: + module = import_module(mod_str) + except ImportError as ex: + arg_parser.error(f"unable to import {mod_str}: {ex}") + try: + func = getattr(module, func_str) + except AttributeError: + arg_parser.error(f"module {mod_str!r} has no attribute {func_str!r}") + + # Compatibility logic + if args.path is not None and not hasattr(socket, "AF_UNIX"): + arg_parser.error( + "file system paths not supported by your operating environment" + ) + + logging.basicConfig(level=logging.DEBUG) + + if args.path and args.hostname is None: + host = port = None + else: + host = args.hostname or "localhost" + port = args.port + + app = func(extra_argv) + run_app(app, host=host, port=port, path=args.path) + arg_parser.exit(message="Stopped\n") + + +if __name__ == "__main__": # pragma: no branch + main(sys.argv[1:]) # pragma: no cover diff --git a/lib/python3.12/site-packages/aiohttp/web_app.py b/lib/python3.12/site-packages/aiohttp/web_app.py new file mode 100644 index 0000000000000000000000000000000000000000..619c0085da1985b97f7e222d30c1870cb010a128 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_app.py @@ -0,0 +1,620 @@ +import asyncio +import logging +import warnings +from functools import lru_cache, partial, update_wrapper +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Awaitable, + Callable, + Dict, + Iterable, + Iterator, + List, + Mapping, + MutableMapping, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, + cast, + overload, +) + +from aiosignal import Signal +from frozenlist import FrozenList + +from . import hdrs +from .abc import ( + AbstractAccessLogger, + AbstractMatchInfo, + AbstractRouter, + AbstractStreamWriter, +) +from .helpers import DEBUG, AppKey +from .http_parser import RawRequestMessage +from .log import web_logger +from .streams import StreamReader +from .typedefs import Handler, Middleware +from .web_exceptions import NotAppKeyWarning +from .web_log import AccessLogger +from .web_middlewares import _fix_request_current_app +from .web_protocol import RequestHandler +from .web_request import Request +from .web_response import StreamResponse +from .web_routedef import AbstractRouteDef +from .web_server import Server +from .web_urldispatcher import ( + AbstractResource, + AbstractRoute, + Domain, + MaskDomain, + MatchedSubAppResource, + PrefixedSubAppResource, + SystemRoute, + UrlDispatcher, +) + +__all__ = ("Application", "CleanupError") + + +if TYPE_CHECKING: + _AppSignal = Signal["Application"] + _RespPrepareSignal = Signal[Request, StreamResponse] + _Middlewares = FrozenList[Middleware] + _MiddlewaresHandlers = Optional[Sequence[Tuple[Middleware, bool]]] + _Subapps = List["Application"] +else: + # No type checker mode, skip types + _AppSignal = Signal + _RespPrepareSignal = Signal + _Middlewares = FrozenList + _MiddlewaresHandlers = Optional[Sequence] + _Subapps = List + +_T = TypeVar("_T") +_U = TypeVar("_U") +_Resource = TypeVar("_Resource", bound=AbstractResource) + + +def _build_middlewares( + handler: Handler, apps: Tuple["Application", ...] +) -> Callable[[Request], Awaitable[StreamResponse]]: + """Apply middlewares to handler.""" + for app in apps[::-1]: + for m, _ in app._middlewares_handlers: # type: ignore[union-attr] + handler = update_wrapper(partial(m, handler=handler), handler) + return handler + + +_cached_build_middleware = lru_cache(maxsize=1024)(_build_middlewares) + + +class Application(MutableMapping[Union[str, AppKey[Any]], Any]): + ATTRS = frozenset( + [ + "logger", + "_debug", + "_router", + "_loop", + "_handler_args", + "_middlewares", + "_middlewares_handlers", + "_has_legacy_middlewares", + "_run_middlewares", + "_state", + "_frozen", + "_pre_frozen", + "_subapps", + "_on_response_prepare", + "_on_startup", + "_on_shutdown", + "_on_cleanup", + "_client_max_size", + "_cleanup_ctx", + ] + ) + + def __init__( + self, + *, + logger: logging.Logger = web_logger, + router: Optional[UrlDispatcher] = None, + middlewares: Iterable[Middleware] = (), + handler_args: Optional[Mapping[str, Any]] = None, + client_max_size: int = 1024**2, + loop: Optional[asyncio.AbstractEventLoop] = None, + debug: Any = ..., # mypy doesn't support ellipsis + ) -> None: + if router is None: + router = UrlDispatcher() + else: + warnings.warn( + "router argument is deprecated", DeprecationWarning, stacklevel=2 + ) + assert isinstance(router, AbstractRouter), router + + if loop is not None: + warnings.warn( + "loop argument is deprecated", DeprecationWarning, stacklevel=2 + ) + + if debug is not ...: + warnings.warn( + "debug argument is deprecated", DeprecationWarning, stacklevel=2 + ) + self._debug = debug + self._router: UrlDispatcher = router + self._loop = loop + self._handler_args = handler_args + self.logger = logger + + self._middlewares: _Middlewares = FrozenList(middlewares) + + # initialized on freezing + self._middlewares_handlers: _MiddlewaresHandlers = None + # initialized on freezing + self._run_middlewares: Optional[bool] = None + self._has_legacy_middlewares: bool = True + + self._state: Dict[Union[AppKey[Any], str], object] = {} + self._frozen = False + self._pre_frozen = False + self._subapps: _Subapps = [] + + self._on_response_prepare: _RespPrepareSignal = Signal(self) + self._on_startup: _AppSignal = Signal(self) + self._on_shutdown: _AppSignal = Signal(self) + self._on_cleanup: _AppSignal = Signal(self) + self._cleanup_ctx = CleanupContext() + self._on_startup.append(self._cleanup_ctx._on_startup) + self._on_cleanup.append(self._cleanup_ctx._on_cleanup) + self._client_max_size = client_max_size + + def __init_subclass__(cls: Type["Application"]) -> None: + warnings.warn( + "Inheritance class {} from web.Application " + "is discouraged".format(cls.__name__), + DeprecationWarning, + stacklevel=3, + ) + + if DEBUG: # pragma: no cover + + def __setattr__(self, name: str, val: Any) -> None: + if name not in self.ATTRS: + warnings.warn( + "Setting custom web.Application.{} attribute " + "is discouraged".format(name), + DeprecationWarning, + stacklevel=2, + ) + super().__setattr__(name, val) + + # MutableMapping API + + def __eq__(self, other: object) -> bool: + return self is other + + @overload # type: ignore[override] + def __getitem__(self, key: AppKey[_T]) -> _T: ... + + @overload + def __getitem__(self, key: str) -> Any: ... + + def __getitem__(self, key: Union[str, AppKey[_T]]) -> Any: + return self._state[key] + + def _check_frozen(self) -> None: + if self._frozen: + warnings.warn( + "Changing state of started or joined application is deprecated", + DeprecationWarning, + stacklevel=3, + ) + + @overload # type: ignore[override] + def __setitem__(self, key: AppKey[_T], value: _T) -> None: ... + + @overload + def __setitem__(self, key: str, value: Any) -> None: ... + + def __setitem__(self, key: Union[str, AppKey[_T]], value: Any) -> None: + self._check_frozen() + if not isinstance(key, AppKey): + warnings.warn( + "It is recommended to use web.AppKey instances for keys.\n" + + "https://docs.aiohttp.org/en/stable/web_advanced.html" + + "#application-s-config", + category=NotAppKeyWarning, + stacklevel=2, + ) + self._state[key] = value + + def __delitem__(self, key: Union[str, AppKey[_T]]) -> None: + self._check_frozen() + del self._state[key] + + def __len__(self) -> int: + return len(self._state) + + def __iter__(self) -> Iterator[Union[str, AppKey[Any]]]: + return iter(self._state) + + def __hash__(self) -> int: + return id(self) + + @overload # type: ignore[override] + def get(self, key: AppKey[_T], default: None = ...) -> Optional[_T]: ... + + @overload + def get(self, key: AppKey[_T], default: _U) -> Union[_T, _U]: ... + + @overload + def get(self, key: str, default: Any = ...) -> Any: ... + + def get(self, key: Union[str, AppKey[_T]], default: Any = None) -> Any: + return self._state.get(key, default) + + ######## + @property + def loop(self) -> asyncio.AbstractEventLoop: + # Technically the loop can be None + # but we mask it by explicit type cast + # to provide more convenient type annotation + warnings.warn("loop property is deprecated", DeprecationWarning, stacklevel=2) + return cast(asyncio.AbstractEventLoop, self._loop) + + def _set_loop(self, loop: Optional[asyncio.AbstractEventLoop]) -> None: + if loop is None: + loop = asyncio.get_event_loop() + if self._loop is not None and self._loop is not loop: + raise RuntimeError( + "web.Application instance initialized with different loop" + ) + + self._loop = loop + + # set loop debug + if self._debug is ...: + self._debug = loop.get_debug() + + # set loop to sub applications + for subapp in self._subapps: + subapp._set_loop(loop) + + @property + def pre_frozen(self) -> bool: + return self._pre_frozen + + def pre_freeze(self) -> None: + if self._pre_frozen: + return + + self._pre_frozen = True + self._middlewares.freeze() + self._router.freeze() + self._on_response_prepare.freeze() + self._cleanup_ctx.freeze() + self._on_startup.freeze() + self._on_shutdown.freeze() + self._on_cleanup.freeze() + self._middlewares_handlers = tuple(self._prepare_middleware()) + self._has_legacy_middlewares = any( + not new_style for _, new_style in self._middlewares_handlers + ) + + # If current app and any subapp do not have middlewares avoid run all + # of the code footprint that it implies, which have a middleware + # hardcoded per app that sets up the current_app attribute. If no + # middlewares are configured the handler will receive the proper + # current_app without needing all of this code. + self._run_middlewares = True if self.middlewares else False + + for subapp in self._subapps: + subapp.pre_freeze() + self._run_middlewares = self._run_middlewares or subapp._run_middlewares + + @property + def frozen(self) -> bool: + return self._frozen + + def freeze(self) -> None: + if self._frozen: + return + + self.pre_freeze() + self._frozen = True + for subapp in self._subapps: + subapp.freeze() + + @property + def debug(self) -> bool: + warnings.warn("debug property is deprecated", DeprecationWarning, stacklevel=2) + return self._debug # type: ignore[no-any-return] + + def _reg_subapp_signals(self, subapp: "Application") -> None: + def reg_handler(signame: str) -> None: + subsig = getattr(subapp, signame) + + async def handler(app: "Application") -> None: + await subsig.send(subapp) + + appsig = getattr(self, signame) + appsig.append(handler) + + reg_handler("on_startup") + reg_handler("on_shutdown") + reg_handler("on_cleanup") + + def add_subapp(self, prefix: str, subapp: "Application") -> PrefixedSubAppResource: + if not isinstance(prefix, str): + raise TypeError("Prefix must be str") + prefix = prefix.rstrip("/") + if not prefix: + raise ValueError("Prefix cannot be empty") + factory = partial(PrefixedSubAppResource, prefix, subapp) + return self._add_subapp(factory, subapp) + + def _add_subapp( + self, resource_factory: Callable[[], _Resource], subapp: "Application" + ) -> _Resource: + if self.frozen: + raise RuntimeError("Cannot add sub application to frozen application") + if subapp.frozen: + raise RuntimeError("Cannot add frozen application") + resource = resource_factory() + self.router.register_resource(resource) + self._reg_subapp_signals(subapp) + self._subapps.append(subapp) + subapp.pre_freeze() + if self._loop is not None: + subapp._set_loop(self._loop) + return resource + + def add_domain(self, domain: str, subapp: "Application") -> MatchedSubAppResource: + if not isinstance(domain, str): + raise TypeError("Domain must be str") + elif "*" in domain: + rule: Domain = MaskDomain(domain) + else: + rule = Domain(domain) + factory = partial(MatchedSubAppResource, rule, subapp) + return self._add_subapp(factory, subapp) + + def add_routes(self, routes: Iterable[AbstractRouteDef]) -> List[AbstractRoute]: + return self.router.add_routes(routes) + + @property + def on_response_prepare(self) -> _RespPrepareSignal: + return self._on_response_prepare + + @property + def on_startup(self) -> _AppSignal: + return self._on_startup + + @property + def on_shutdown(self) -> _AppSignal: + return self._on_shutdown + + @property + def on_cleanup(self) -> _AppSignal: + return self._on_cleanup + + @property + def cleanup_ctx(self) -> "CleanupContext": + return self._cleanup_ctx + + @property + def router(self) -> UrlDispatcher: + return self._router + + @property + def middlewares(self) -> _Middlewares: + return self._middlewares + + def _make_handler( + self, + *, + loop: Optional[asyncio.AbstractEventLoop] = None, + access_log_class: Type[AbstractAccessLogger] = AccessLogger, + **kwargs: Any, + ) -> Server: + + if not issubclass(access_log_class, AbstractAccessLogger): + raise TypeError( + "access_log_class must be subclass of " + "aiohttp.abc.AbstractAccessLogger, got {}".format(access_log_class) + ) + + self._set_loop(loop) + self.freeze() + + kwargs["debug"] = self._debug + kwargs["access_log_class"] = access_log_class + if self._handler_args: + for k, v in self._handler_args.items(): + kwargs[k] = v + + return Server( + self._handle, # type: ignore[arg-type] + request_factory=self._make_request, + loop=self._loop, + **kwargs, + ) + + def make_handler( + self, + *, + loop: Optional[asyncio.AbstractEventLoop] = None, + access_log_class: Type[AbstractAccessLogger] = AccessLogger, + **kwargs: Any, + ) -> Server: + + warnings.warn( + "Application.make_handler(...) is deprecated, use AppRunner API instead", + DeprecationWarning, + stacklevel=2, + ) + + return self._make_handler( + loop=loop, access_log_class=access_log_class, **kwargs + ) + + async def startup(self) -> None: + """Causes on_startup signal + + Should be called in the event loop along with the request handler. + """ + await self.on_startup.send(self) + + async def shutdown(self) -> None: + """Causes on_shutdown signal + + Should be called before cleanup() + """ + await self.on_shutdown.send(self) + + async def cleanup(self) -> None: + """Causes on_cleanup signal + + Should be called after shutdown() + """ + if self.on_cleanup.frozen: + await self.on_cleanup.send(self) + else: + # If an exception occurs in startup, ensure cleanup contexts are completed. + await self._cleanup_ctx._on_cleanup(self) + + def _make_request( + self, + message: RawRequestMessage, + payload: StreamReader, + protocol: RequestHandler, + writer: AbstractStreamWriter, + task: "asyncio.Task[None]", + _cls: Type[Request] = Request, + ) -> Request: + if TYPE_CHECKING: + assert self._loop is not None + return _cls( + message, + payload, + protocol, + writer, + task, + self._loop, + client_max_size=self._client_max_size, + ) + + def _prepare_middleware(self) -> Iterator[Tuple[Middleware, bool]]: + for m in reversed(self._middlewares): + if getattr(m, "__middleware_version__", None) == 1: + yield m, True + else: + warnings.warn( + f'old-style middleware "{m!r}" deprecated, see #2252', + DeprecationWarning, + stacklevel=2, + ) + yield m, False + + yield _fix_request_current_app(self), True + + async def _handle(self, request: Request) -> StreamResponse: + loop = asyncio.get_event_loop() + debug = loop.get_debug() + match_info = await self._router.resolve(request) + if debug: # pragma: no cover + if not isinstance(match_info, AbstractMatchInfo): + raise TypeError( + "match_info should be AbstractMatchInfo " + "instance, not {!r}".format(match_info) + ) + match_info.add_app(self) + + match_info.freeze() + + request._match_info = match_info + + if request.headers.get(hdrs.EXPECT): + resp = await match_info.expect_handler(request) + await request.writer.drain() + if resp is not None: + return resp + + handler = match_info.handler + + if self._run_middlewares: + # If its a SystemRoute, don't cache building the middlewares since + # they are constructed for every MatchInfoError as a new handler + # is made each time. + if not self._has_legacy_middlewares and not isinstance( + match_info.route, SystemRoute + ): + handler = _cached_build_middleware(handler, match_info.apps) + else: + for app in match_info.apps[::-1]: + for m, new_style in app._middlewares_handlers: # type: ignore[union-attr] + if new_style: + handler = update_wrapper( + partial(m, handler=handler), handler + ) + else: + handler = await m(app, handler) # type: ignore[arg-type,assignment] + + return await handler(request) + + def __call__(self) -> "Application": + """gunicorn compatibility""" + return self + + def __repr__(self) -> str: + return f"" + + def __bool__(self) -> bool: + return True + + +class CleanupError(RuntimeError): + @property + def exceptions(self) -> List[BaseException]: + return cast(List[BaseException], self.args[1]) + + +if TYPE_CHECKING: + _CleanupContextBase = FrozenList[Callable[[Application], AsyncIterator[None]]] +else: + _CleanupContextBase = FrozenList + + +class CleanupContext(_CleanupContextBase): + def __init__(self) -> None: + super().__init__() + self._exits: List[AsyncIterator[None]] = [] + + async def _on_startup(self, app: Application) -> None: + for cb in self: + it = cb(app).__aiter__() + await it.__anext__() + self._exits.append(it) + + async def _on_cleanup(self, app: Application) -> None: + errors = [] + for it in reversed(self._exits): + try: + await it.__anext__() + except StopAsyncIteration: + pass + except (Exception, asyncio.CancelledError) as exc: + errors.append(exc) + else: + errors.append(RuntimeError(f"{it!r} has more than one 'yield'")) + if errors: + if len(errors) == 1: + raise errors[0] + else: + raise CleanupError("Multiple errors on cleanup stage", errors) diff --git a/lib/python3.12/site-packages/aiohttp/web_exceptions.py b/lib/python3.12/site-packages/aiohttp/web_exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2c1e72d40ac93c00cbfcef6c84888a336855a3 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_exceptions.py @@ -0,0 +1,452 @@ +import warnings +from typing import Any, Dict, Iterable, List, Optional, Set # noqa + +from yarl import URL + +from .typedefs import LooseHeaders, StrOrURL +from .web_response import Response + +__all__ = ( + "HTTPException", + "HTTPError", + "HTTPRedirection", + "HTTPSuccessful", + "HTTPOk", + "HTTPCreated", + "HTTPAccepted", + "HTTPNonAuthoritativeInformation", + "HTTPNoContent", + "HTTPResetContent", + "HTTPPartialContent", + "HTTPMove", + "HTTPMultipleChoices", + "HTTPMovedPermanently", + "HTTPFound", + "HTTPSeeOther", + "HTTPNotModified", + "HTTPUseProxy", + "HTTPTemporaryRedirect", + "HTTPPermanentRedirect", + "HTTPClientError", + "HTTPBadRequest", + "HTTPUnauthorized", + "HTTPPaymentRequired", + "HTTPForbidden", + "HTTPNotFound", + "HTTPMethodNotAllowed", + "HTTPNotAcceptable", + "HTTPProxyAuthenticationRequired", + "HTTPRequestTimeout", + "HTTPConflict", + "HTTPGone", + "HTTPLengthRequired", + "HTTPPreconditionFailed", + "HTTPRequestEntityTooLarge", + "HTTPRequestURITooLong", + "HTTPUnsupportedMediaType", + "HTTPRequestRangeNotSatisfiable", + "HTTPExpectationFailed", + "HTTPMisdirectedRequest", + "HTTPUnprocessableEntity", + "HTTPFailedDependency", + "HTTPUpgradeRequired", + "HTTPPreconditionRequired", + "HTTPTooManyRequests", + "HTTPRequestHeaderFieldsTooLarge", + "HTTPUnavailableForLegalReasons", + "HTTPServerError", + "HTTPInternalServerError", + "HTTPNotImplemented", + "HTTPBadGateway", + "HTTPServiceUnavailable", + "HTTPGatewayTimeout", + "HTTPVersionNotSupported", + "HTTPVariantAlsoNegotiates", + "HTTPInsufficientStorage", + "HTTPNotExtended", + "HTTPNetworkAuthenticationRequired", +) + + +class NotAppKeyWarning(UserWarning): + """Warning when not using AppKey in Application.""" + + +############################################################ +# HTTP Exceptions +############################################################ + + +class HTTPException(Response, Exception): + + # You should set in subclasses: + # status = 200 + + status_code = -1 + empty_body = False + + __http_exception__ = True + + def __init__( + self, + *, + headers: Optional[LooseHeaders] = None, + reason: Optional[str] = None, + body: Any = None, + text: Optional[str] = None, + content_type: Optional[str] = None, + ) -> None: + if body is not None: + warnings.warn( + "body argument is deprecated for http web exceptions", + DeprecationWarning, + ) + Response.__init__( + self, + status=self.status_code, + headers=headers, + reason=reason, + body=body, + text=text, + content_type=content_type, + ) + Exception.__init__(self, self.reason) + if self.body is None and not self.empty_body: + self.text = f"{self.status}: {self.reason}" + + def __bool__(self) -> bool: + return True + + +class HTTPError(HTTPException): + """Base class for exceptions with status codes in the 400s and 500s.""" + + +class HTTPRedirection(HTTPException): + """Base class for exceptions with status codes in the 300s.""" + + +class HTTPSuccessful(HTTPException): + """Base class for exceptions with status codes in the 200s.""" + + +class HTTPOk(HTTPSuccessful): + status_code = 200 + + +class HTTPCreated(HTTPSuccessful): + status_code = 201 + + +class HTTPAccepted(HTTPSuccessful): + status_code = 202 + + +class HTTPNonAuthoritativeInformation(HTTPSuccessful): + status_code = 203 + + +class HTTPNoContent(HTTPSuccessful): + status_code = 204 + empty_body = True + + +class HTTPResetContent(HTTPSuccessful): + status_code = 205 + empty_body = True + + +class HTTPPartialContent(HTTPSuccessful): + status_code = 206 + + +############################################################ +# 3xx redirection +############################################################ + + +class HTTPMove(HTTPRedirection): + def __init__( + self, + location: StrOrURL, + *, + headers: Optional[LooseHeaders] = None, + reason: Optional[str] = None, + body: Any = None, + text: Optional[str] = None, + content_type: Optional[str] = None, + ) -> None: + if not location: + raise ValueError("HTTP redirects need a location to redirect to.") + super().__init__( + headers=headers, + reason=reason, + body=body, + text=text, + content_type=content_type, + ) + self.headers["Location"] = str(URL(location)) + self.location = location + + +class HTTPMultipleChoices(HTTPMove): + status_code = 300 + + +class HTTPMovedPermanently(HTTPMove): + status_code = 301 + + +class HTTPFound(HTTPMove): + status_code = 302 + + +# This one is safe after a POST (the redirected location will be +# retrieved with GET): +class HTTPSeeOther(HTTPMove): + status_code = 303 + + +class HTTPNotModified(HTTPRedirection): + # FIXME: this should include a date or etag header + status_code = 304 + empty_body = True + + +class HTTPUseProxy(HTTPMove): + # Not a move, but looks a little like one + status_code = 305 + + +class HTTPTemporaryRedirect(HTTPMove): + status_code = 307 + + +class HTTPPermanentRedirect(HTTPMove): + status_code = 308 + + +############################################################ +# 4xx client error +############################################################ + + +class HTTPClientError(HTTPError): + pass + + +class HTTPBadRequest(HTTPClientError): + status_code = 400 + + +class HTTPUnauthorized(HTTPClientError): + status_code = 401 + + +class HTTPPaymentRequired(HTTPClientError): + status_code = 402 + + +class HTTPForbidden(HTTPClientError): + status_code = 403 + + +class HTTPNotFound(HTTPClientError): + status_code = 404 + + +class HTTPMethodNotAllowed(HTTPClientError): + status_code = 405 + + def __init__( + self, + method: str, + allowed_methods: Iterable[str], + *, + headers: Optional[LooseHeaders] = None, + reason: Optional[str] = None, + body: Any = None, + text: Optional[str] = None, + content_type: Optional[str] = None, + ) -> None: + allow = ",".join(sorted(allowed_methods)) + super().__init__( + headers=headers, + reason=reason, + body=body, + text=text, + content_type=content_type, + ) + self.headers["Allow"] = allow + self.allowed_methods: Set[str] = set(allowed_methods) + self.method = method.upper() + + +class HTTPNotAcceptable(HTTPClientError): + status_code = 406 + + +class HTTPProxyAuthenticationRequired(HTTPClientError): + status_code = 407 + + +class HTTPRequestTimeout(HTTPClientError): + status_code = 408 + + +class HTTPConflict(HTTPClientError): + status_code = 409 + + +class HTTPGone(HTTPClientError): + status_code = 410 + + +class HTTPLengthRequired(HTTPClientError): + status_code = 411 + + +class HTTPPreconditionFailed(HTTPClientError): + status_code = 412 + + +class HTTPRequestEntityTooLarge(HTTPClientError): + status_code = 413 + + def __init__(self, max_size: float, actual_size: float, **kwargs: Any) -> None: + kwargs.setdefault( + "text", + "Maximum request body size {} exceeded, " + "actual body size {}".format(max_size, actual_size), + ) + super().__init__(**kwargs) + + +class HTTPRequestURITooLong(HTTPClientError): + status_code = 414 + + +class HTTPUnsupportedMediaType(HTTPClientError): + status_code = 415 + + +class HTTPRequestRangeNotSatisfiable(HTTPClientError): + status_code = 416 + + +class HTTPExpectationFailed(HTTPClientError): + status_code = 417 + + +class HTTPMisdirectedRequest(HTTPClientError): + status_code = 421 + + +class HTTPUnprocessableEntity(HTTPClientError): + status_code = 422 + + +class HTTPFailedDependency(HTTPClientError): + status_code = 424 + + +class HTTPUpgradeRequired(HTTPClientError): + status_code = 426 + + +class HTTPPreconditionRequired(HTTPClientError): + status_code = 428 + + +class HTTPTooManyRequests(HTTPClientError): + status_code = 429 + + +class HTTPRequestHeaderFieldsTooLarge(HTTPClientError): + status_code = 431 + + +class HTTPUnavailableForLegalReasons(HTTPClientError): + status_code = 451 + + def __init__( + self, + link: Optional[StrOrURL], + *, + headers: Optional[LooseHeaders] = None, + reason: Optional[str] = None, + body: Any = None, + text: Optional[str] = None, + content_type: Optional[str] = None, + ) -> None: + super().__init__( + headers=headers, + reason=reason, + body=body, + text=text, + content_type=content_type, + ) + self._link = None + if link: + self._link = URL(link) + self.headers["Link"] = f'<{str(self._link)}>; rel="blocked-by"' + + @property + def link(self) -> Optional[URL]: + return self._link + + +############################################################ +# 5xx Server Error +############################################################ +# Response status codes beginning with the digit "5" indicate cases in +# which the server is aware that it has erred or is incapable of +# performing the request. Except when responding to a HEAD request, the +# server SHOULD include an entity containing an explanation of the error +# situation, and whether it is a temporary or permanent condition. User +# agents SHOULD display any included entity to the user. These response +# codes are applicable to any request method. + + +class HTTPServerError(HTTPError): + pass + + +class HTTPInternalServerError(HTTPServerError): + status_code = 500 + + +class HTTPNotImplemented(HTTPServerError): + status_code = 501 + + +class HTTPBadGateway(HTTPServerError): + status_code = 502 + + +class HTTPServiceUnavailable(HTTPServerError): + status_code = 503 + + +class HTTPGatewayTimeout(HTTPServerError): + status_code = 504 + + +class HTTPVersionNotSupported(HTTPServerError): + status_code = 505 + + +class HTTPVariantAlsoNegotiates(HTTPServerError): + status_code = 506 + + +class HTTPInsufficientStorage(HTTPServerError): + status_code = 507 + + +class HTTPNotExtended(HTTPServerError): + status_code = 510 + + +class HTTPNetworkAuthenticationRequired(HTTPServerError): + status_code = 511 diff --git a/lib/python3.12/site-packages/aiohttp/web_fileresponse.py b/lib/python3.12/site-packages/aiohttp/web_fileresponse.py new file mode 100644 index 0000000000000000000000000000000000000000..26484b9483a9e7e49e1975a2002535009755e133 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_fileresponse.py @@ -0,0 +1,418 @@ +import asyncio +import io +import os +import pathlib +import sys +from contextlib import suppress +from enum import Enum, auto +from mimetypes import MimeTypes +from stat import S_ISREG +from types import MappingProxyType +from typing import ( # noqa + IO, + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Final, + Iterator, + List, + Optional, + Set, + Tuple, + Union, + cast, +) + +from . import hdrs +from .abc import AbstractStreamWriter +from .helpers import ETAG_ANY, ETag, must_be_empty_body +from .typedefs import LooseHeaders, PathLike +from .web_exceptions import ( + HTTPForbidden, + HTTPNotFound, + HTTPNotModified, + HTTPPartialContent, + HTTPPreconditionFailed, + HTTPRequestRangeNotSatisfiable, +) +from .web_response import StreamResponse + +__all__ = ("FileResponse",) + +if TYPE_CHECKING: + from .web_request import BaseRequest + + +_T_OnChunkSent = Optional[Callable[[bytes], Awaitable[None]]] + + +NOSENDFILE: Final[bool] = bool(os.environ.get("AIOHTTP_NOSENDFILE")) + +CONTENT_TYPES: Final[MimeTypes] = MimeTypes() + +# File extension to IANA encodings map that will be checked in the order defined. +ENCODING_EXTENSIONS = MappingProxyType( + {ext: CONTENT_TYPES.encodings_map[ext] for ext in (".br", ".gz")} +) + +FALLBACK_CONTENT_TYPE = "application/octet-stream" + +# Provide additional MIME type/extension pairs to be recognized. +# https://en.wikipedia.org/wiki/List_of_archive_formats#Compression_only +ADDITIONAL_CONTENT_TYPES = MappingProxyType( + { + "application/gzip": ".gz", + "application/x-brotli": ".br", + "application/x-bzip2": ".bz2", + "application/x-compress": ".Z", + "application/x-xz": ".xz", + } +) + + +class _FileResponseResult(Enum): + """The result of the file response.""" + + SEND_FILE = auto() # Ie a regular file to send + NOT_ACCEPTABLE = auto() # Ie a socket, or non-regular file + PRE_CONDITION_FAILED = auto() # Ie If-Match or If-None-Match failed + NOT_MODIFIED = auto() # 304 Not Modified + + +# Add custom pairs and clear the encodings map so guess_type ignores them. +CONTENT_TYPES.encodings_map.clear() +for content_type, extension in ADDITIONAL_CONTENT_TYPES.items(): + CONTENT_TYPES.add_type(content_type, extension) + + +_CLOSE_FUTURES: Set[asyncio.Future[None]] = set() + + +class FileResponse(StreamResponse): + """A response object can be used to send files.""" + + def __init__( + self, + path: PathLike, + chunk_size: int = 256 * 1024, + status: int = 200, + reason: Optional[str] = None, + headers: Optional[LooseHeaders] = None, + ) -> None: + super().__init__(status=status, reason=reason, headers=headers) + + self._path = pathlib.Path(path) + self._chunk_size = chunk_size + + def _seek_and_read(self, fobj: IO[Any], offset: int, chunk_size: int) -> bytes: + fobj.seek(offset) + return fobj.read(chunk_size) # type: ignore[no-any-return] + + async def _sendfile_fallback( + self, writer: AbstractStreamWriter, fobj: IO[Any], offset: int, count: int + ) -> AbstractStreamWriter: + # To keep memory usage low,fobj is transferred in chunks + # controlled by the constructor's chunk_size argument. + + chunk_size = self._chunk_size + loop = asyncio.get_event_loop() + chunk = await loop.run_in_executor( + None, self._seek_and_read, fobj, offset, chunk_size + ) + while chunk: + await writer.write(chunk) + count = count - chunk_size + if count <= 0: + break + chunk = await loop.run_in_executor(None, fobj.read, min(chunk_size, count)) + + await writer.drain() + return writer + + async def _sendfile( + self, request: "BaseRequest", fobj: IO[Any], offset: int, count: int + ) -> AbstractStreamWriter: + writer = await super().prepare(request) + assert writer is not None + + if NOSENDFILE or self.compression: + return await self._sendfile_fallback(writer, fobj, offset, count) + + loop = request._loop + transport = request.transport + assert transport is not None + + try: + await loop.sendfile(transport, fobj, offset, count) + except NotImplementedError: + return await self._sendfile_fallback(writer, fobj, offset, count) + + await super().write_eof() + return writer + + @staticmethod + def _etag_match(etag_value: str, etags: Tuple[ETag, ...], *, weak: bool) -> bool: + if len(etags) == 1 and etags[0].value == ETAG_ANY: + return True + return any( + etag.value == etag_value for etag in etags if weak or not etag.is_weak + ) + + async def _not_modified( + self, request: "BaseRequest", etag_value: str, last_modified: float + ) -> Optional[AbstractStreamWriter]: + self.set_status(HTTPNotModified.status_code) + self._length_check = False + self.etag = etag_value + self.last_modified = last_modified + # Delete any Content-Length headers provided by user. HTTP 304 + # should always have empty response body + return await super().prepare(request) + + async def _precondition_failed( + self, request: "BaseRequest" + ) -> Optional[AbstractStreamWriter]: + self.set_status(HTTPPreconditionFailed.status_code) + self.content_length = 0 + return await super().prepare(request) + + def _make_response( + self, request: "BaseRequest", accept_encoding: str + ) -> Tuple[ + _FileResponseResult, Optional[io.BufferedReader], os.stat_result, Optional[str] + ]: + """Return the response result, io object, stat result, and encoding. + + If an uncompressed file is returned, the encoding is set to + :py:data:`None`. + + This method should be called from a thread executor + since it calls os.stat which may block. + """ + file_path, st, file_encoding = self._get_file_path_stat_encoding( + accept_encoding + ) + if not file_path: + return _FileResponseResult.NOT_ACCEPTABLE, None, st, None + + etag_value = f"{st.st_mtime_ns:x}-{st.st_size:x}" + + # https://www.rfc-editor.org/rfc/rfc9110#section-13.1.1-2 + if (ifmatch := request.if_match) is not None and not self._etag_match( + etag_value, ifmatch, weak=False + ): + return _FileResponseResult.PRE_CONDITION_FAILED, None, st, file_encoding + + if ( + (unmodsince := request.if_unmodified_since) is not None + and ifmatch is None + and st.st_mtime > unmodsince.timestamp() + ): + return _FileResponseResult.PRE_CONDITION_FAILED, None, st, file_encoding + + # https://www.rfc-editor.org/rfc/rfc9110#section-13.1.2-2 + if (ifnonematch := request.if_none_match) is not None and self._etag_match( + etag_value, ifnonematch, weak=True + ): + return _FileResponseResult.NOT_MODIFIED, None, st, file_encoding + + if ( + (modsince := request.if_modified_since) is not None + and ifnonematch is None + and st.st_mtime <= modsince.timestamp() + ): + return _FileResponseResult.NOT_MODIFIED, None, st, file_encoding + + fobj = file_path.open("rb") + with suppress(OSError): + # fstat() may not be available on all platforms + # Once we open the file, we want the fstat() to ensure + # the file has not changed between the first stat() + # and the open(). + st = os.stat(fobj.fileno()) + return _FileResponseResult.SEND_FILE, fobj, st, file_encoding + + def _get_file_path_stat_encoding( + self, accept_encoding: str + ) -> Tuple[Optional[pathlib.Path], os.stat_result, Optional[str]]: + file_path = self._path + for file_extension, file_encoding in ENCODING_EXTENSIONS.items(): + if file_encoding not in accept_encoding: + continue + + compressed_path = file_path.with_suffix(file_path.suffix + file_extension) + with suppress(OSError): + # Do not follow symlinks and ignore any non-regular files. + st = compressed_path.lstat() + if S_ISREG(st.st_mode): + return compressed_path, st, file_encoding + + # Fallback to the uncompressed file + st = file_path.stat() + return file_path if S_ISREG(st.st_mode) else None, st, None + + async def prepare(self, request: "BaseRequest") -> Optional[AbstractStreamWriter]: + loop = asyncio.get_running_loop() + # Encoding comparisons should be case-insensitive + # https://www.rfc-editor.org/rfc/rfc9110#section-8.4.1 + accept_encoding = request.headers.get(hdrs.ACCEPT_ENCODING, "").lower() + try: + response_result, fobj, st, file_encoding = await loop.run_in_executor( + None, self._make_response, request, accept_encoding + ) + except PermissionError: + self.set_status(HTTPForbidden.status_code) + return await super().prepare(request) + except OSError: + # Most likely to be FileNotFoundError or OSError for circular + # symlinks in python >= 3.13, so respond with 404. + self.set_status(HTTPNotFound.status_code) + return await super().prepare(request) + + # Forbid special files like sockets, pipes, devices, etc. + if response_result is _FileResponseResult.NOT_ACCEPTABLE: + self.set_status(HTTPForbidden.status_code) + return await super().prepare(request) + + if response_result is _FileResponseResult.PRE_CONDITION_FAILED: + return await self._precondition_failed(request) + + if response_result is _FileResponseResult.NOT_MODIFIED: + etag_value = f"{st.st_mtime_ns:x}-{st.st_size:x}" + last_modified = st.st_mtime + return await self._not_modified(request, etag_value, last_modified) + + assert fobj is not None + try: + return await self._prepare_open_file(request, fobj, st, file_encoding) + finally: + # We do not await here because we do not want to wait + # for the executor to finish before returning the response + # so the connection can begin servicing another request + # as soon as possible. + close_future = loop.run_in_executor(None, fobj.close) + # Hold a strong reference to the future to prevent it from being + # garbage collected before it completes. + _CLOSE_FUTURES.add(close_future) + close_future.add_done_callback(_CLOSE_FUTURES.remove) + + async def _prepare_open_file( + self, + request: "BaseRequest", + fobj: io.BufferedReader, + st: os.stat_result, + file_encoding: Optional[str], + ) -> Optional[AbstractStreamWriter]: + status = self._status + file_size: int = st.st_size + file_mtime: float = st.st_mtime + count: int = file_size + start: Optional[int] = None + + if (ifrange := request.if_range) is None or file_mtime <= ifrange.timestamp(): + # If-Range header check: + # condition = cached date >= last modification date + # return 206 if True else 200. + # if False: + # Range header would not be processed, return 200 + # if True but Range header missing + # return 200 + try: + rng = request.http_range + start = rng.start + end: Optional[int] = rng.stop + except ValueError: + # https://tools.ietf.org/html/rfc7233: + # A server generating a 416 (Range Not Satisfiable) response to + # a byte-range request SHOULD send a Content-Range header field + # with an unsatisfied-range value. + # The complete-length in a 416 response indicates the current + # length of the selected representation. + # + # Will do the same below. Many servers ignore this and do not + # send a Content-Range header with HTTP 416 + self._headers[hdrs.CONTENT_RANGE] = f"bytes */{file_size}" + self.set_status(HTTPRequestRangeNotSatisfiable.status_code) + return await super().prepare(request) + + # If a range request has been made, convert start, end slice + # notation into file pointer offset and count + if start is not None: + if start < 0 and end is None: # return tail of file + start += file_size + if start < 0: + # if Range:bytes=-1000 in request header but file size + # is only 200, there would be trouble without this + start = 0 + count = file_size - start + else: + # rfc7233:If the last-byte-pos value is + # absent, or if the value is greater than or equal to + # the current length of the representation data, + # the byte range is interpreted as the remainder + # of the representation (i.e., the server replaces the + # value of last-byte-pos with a value that is one less than + # the current length of the selected representation). + count = ( + min(end if end is not None else file_size, file_size) - start + ) + + if start >= file_size: + # HTTP 416 should be returned in this case. + # + # According to https://tools.ietf.org/html/rfc7233: + # If a valid byte-range-set includes at least one + # byte-range-spec with a first-byte-pos that is less than + # the current length of the representation, or at least one + # suffix-byte-range-spec with a non-zero suffix-length, + # then the byte-range-set is satisfiable. Otherwise, the + # byte-range-set is unsatisfiable. + self._headers[hdrs.CONTENT_RANGE] = f"bytes */{file_size}" + self.set_status(HTTPRequestRangeNotSatisfiable.status_code) + return await super().prepare(request) + + status = HTTPPartialContent.status_code + # Even though you are sending the whole file, you should still + # return a HTTP 206 for a Range request. + self.set_status(status) + + # If the Content-Type header is not already set, guess it based on the + # extension of the request path. The encoding returned by guess_type + # can be ignored since the map was cleared above. + if hdrs.CONTENT_TYPE not in self._headers: + if sys.version_info >= (3, 13): + guesser = CONTENT_TYPES.guess_file_type + else: + guesser = CONTENT_TYPES.guess_type + self.content_type = guesser(self._path)[0] or FALLBACK_CONTENT_TYPE + + if file_encoding: + self._headers[hdrs.CONTENT_ENCODING] = file_encoding + self._headers[hdrs.VARY] = hdrs.ACCEPT_ENCODING + # Disable compression if we are already sending + # a compressed file since we don't want to double + # compress. + self._compression = False + + self.etag = f"{st.st_mtime_ns:x}-{st.st_size:x}" + self.last_modified = file_mtime + self.content_length = count + + self._headers[hdrs.ACCEPT_RANGES] = "bytes" + + if status == HTTPPartialContent.status_code: + real_start = start + assert real_start is not None + self._headers[hdrs.CONTENT_RANGE] = "bytes {}-{}/{}".format( + real_start, real_start + count - 1, file_size + ) + + # If we are sending 0 bytes calling sendfile() will throw a ValueError + if count == 0 or must_be_empty_body(request.method, status): + return await super().prepare(request) + + # be aware that start could be None or int=0 here. + offset = start or 0 + + return await self._sendfile(request, fobj, offset, count) diff --git a/lib/python3.12/site-packages/aiohttp/web_log.py b/lib/python3.12/site-packages/aiohttp/web_log.py new file mode 100644 index 0000000000000000000000000000000000000000..d5ea2beeb152974ce5dd9f3e7990133ce04f7980 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_log.py @@ -0,0 +1,216 @@ +import datetime +import functools +import logging +import os +import re +import time as time_mod +from collections import namedtuple +from typing import Any, Callable, Dict, Iterable, List, Tuple # noqa + +from .abc import AbstractAccessLogger +from .web_request import BaseRequest +from .web_response import StreamResponse + +KeyMethod = namedtuple("KeyMethod", "key method") + + +class AccessLogger(AbstractAccessLogger): + """Helper object to log access. + + Usage: + log = logging.getLogger("spam") + log_format = "%a %{User-Agent}i" + access_logger = AccessLogger(log, log_format) + access_logger.log(request, response, time) + + Format: + %% The percent sign + %a Remote IP-address (IP-address of proxy if using reverse proxy) + %t Time when the request was started to process + %P The process ID of the child that serviced the request + %r First line of request + %s Response status code + %b Size of response in bytes, including HTTP headers + %T Time taken to serve the request, in seconds + %Tf Time taken to serve the request, in seconds with floating fraction + in .06f format + %D Time taken to serve the request, in microseconds + %{FOO}i request.headers['FOO'] + %{FOO}o response.headers['FOO'] + %{FOO}e os.environ['FOO'] + + """ + + LOG_FORMAT_MAP = { + "a": "remote_address", + "t": "request_start_time", + "P": "process_id", + "r": "first_request_line", + "s": "response_status", + "b": "response_size", + "T": "request_time", + "Tf": "request_time_frac", + "D": "request_time_micro", + "i": "request_header", + "o": "response_header", + } + + LOG_FORMAT = '%a %t "%r" %s %b "%{Referer}i" "%{User-Agent}i"' + FORMAT_RE = re.compile(r"%(\{([A-Za-z0-9\-_]+)\}([ioe])|[atPrsbOD]|Tf?)") + CLEANUP_RE = re.compile(r"(%[^s])") + _FORMAT_CACHE: Dict[str, Tuple[str, List[KeyMethod]]] = {} + + def __init__(self, logger: logging.Logger, log_format: str = LOG_FORMAT) -> None: + """Initialise the logger. + + logger is a logger object to be used for logging. + log_format is a string with apache compatible log format description. + + """ + super().__init__(logger, log_format=log_format) + + _compiled_format = AccessLogger._FORMAT_CACHE.get(log_format) + if not _compiled_format: + _compiled_format = self.compile_format(log_format) + AccessLogger._FORMAT_CACHE[log_format] = _compiled_format + + self._log_format, self._methods = _compiled_format + + def compile_format(self, log_format: str) -> Tuple[str, List[KeyMethod]]: + """Translate log_format into form usable by modulo formatting + + All known atoms will be replaced with %s + Also methods for formatting of those atoms will be added to + _methods in appropriate order + + For example we have log_format = "%a %t" + This format will be translated to "%s %s" + Also contents of _methods will be + [self._format_a, self._format_t] + These method will be called and results will be passed + to translated string format. + + Each _format_* method receive 'args' which is list of arguments + given to self.log + + Exceptions are _format_e, _format_i and _format_o methods which + also receive key name (by functools.partial) + + """ + # list of (key, method) tuples, we don't use an OrderedDict as users + # can repeat the same key more than once + methods = list() + + for atom in self.FORMAT_RE.findall(log_format): + if atom[1] == "": + format_key1 = self.LOG_FORMAT_MAP[atom[0]] + m = getattr(AccessLogger, "_format_%s" % atom[0]) + key_method = KeyMethod(format_key1, m) + else: + format_key2 = (self.LOG_FORMAT_MAP[atom[2]], atom[1]) + m = getattr(AccessLogger, "_format_%s" % atom[2]) + key_method = KeyMethod(format_key2, functools.partial(m, atom[1])) + + methods.append(key_method) + + log_format = self.FORMAT_RE.sub(r"%s", log_format) + log_format = self.CLEANUP_RE.sub(r"%\1", log_format) + return log_format, methods + + @staticmethod + def _format_i( + key: str, request: BaseRequest, response: StreamResponse, time: float + ) -> str: + if request is None: + return "(no headers)" + + # suboptimal, make istr(key) once + return request.headers.get(key, "-") + + @staticmethod + def _format_o( + key: str, request: BaseRequest, response: StreamResponse, time: float + ) -> str: + # suboptimal, make istr(key) once + return response.headers.get(key, "-") + + @staticmethod + def _format_a(request: BaseRequest, response: StreamResponse, time: float) -> str: + if request is None: + return "-" + ip = request.remote + return ip if ip is not None else "-" + + @staticmethod + def _format_t(request: BaseRequest, response: StreamResponse, time: float) -> str: + tz = datetime.timezone(datetime.timedelta(seconds=-time_mod.timezone)) + now = datetime.datetime.now(tz) + start_time = now - datetime.timedelta(seconds=time) + return start_time.strftime("[%d/%b/%Y:%H:%M:%S %z]") + + @staticmethod + def _format_P(request: BaseRequest, response: StreamResponse, time: float) -> str: + return "<%s>" % os.getpid() + + @staticmethod + def _format_r(request: BaseRequest, response: StreamResponse, time: float) -> str: + if request is None: + return "-" + return "{} {} HTTP/{}.{}".format( + request.method, + request.path_qs, + request.version.major, + request.version.minor, + ) + + @staticmethod + def _format_s(request: BaseRequest, response: StreamResponse, time: float) -> int: + return response.status + + @staticmethod + def _format_b(request: BaseRequest, response: StreamResponse, time: float) -> int: + return response.body_length + + @staticmethod + def _format_T(request: BaseRequest, response: StreamResponse, time: float) -> str: + return str(round(time)) + + @staticmethod + def _format_Tf(request: BaseRequest, response: StreamResponse, time: float) -> str: + return "%06f" % time + + @staticmethod + def _format_D(request: BaseRequest, response: StreamResponse, time: float) -> str: + return str(round(time * 1000000)) + + def _format_line( + self, request: BaseRequest, response: StreamResponse, time: float + ) -> Iterable[Tuple[str, Callable[[BaseRequest, StreamResponse, float], str]]]: + return [(key, method(request, response, time)) for key, method in self._methods] + + @property + def enabled(self) -> bool: + """Check if logger is enabled.""" + # Avoid formatting the log line if it will not be emitted. + return self.logger.isEnabledFor(logging.INFO) + + def log(self, request: BaseRequest, response: StreamResponse, time: float) -> None: + try: + fmt_info = self._format_line(request, response, time) + + values = list() + extra = dict() + for key, value in fmt_info: + values.append(value) + + if key.__class__ is str: + extra[key] = value + else: + k1, k2 = key # type: ignore[misc] + dct = extra.get(k1, {}) # type: ignore[var-annotated,has-type] + dct[k2] = value # type: ignore[index,has-type] + extra[k1] = dct # type: ignore[has-type,assignment] + + self.logger.info(self._log_format % tuple(values), extra=extra) + except Exception: + self.logger.exception("Error in logging") diff --git a/lib/python3.12/site-packages/aiohttp/web_middlewares.py b/lib/python3.12/site-packages/aiohttp/web_middlewares.py new file mode 100644 index 0000000000000000000000000000000000000000..2f1f5f58e6e38845d4d2d4ffdd2748fc519fa5bf --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_middlewares.py @@ -0,0 +1,121 @@ +import re +from typing import TYPE_CHECKING, Tuple, Type, TypeVar + +from .typedefs import Handler, Middleware +from .web_exceptions import HTTPMove, HTTPPermanentRedirect +from .web_request import Request +from .web_response import StreamResponse +from .web_urldispatcher import SystemRoute + +__all__ = ( + "middleware", + "normalize_path_middleware", +) + +if TYPE_CHECKING: + from .web_app import Application + +_Func = TypeVar("_Func") + + +async def _check_request_resolves(request: Request, path: str) -> Tuple[bool, Request]: + alt_request = request.clone(rel_url=path) + + match_info = await request.app.router.resolve(alt_request) + alt_request._match_info = match_info + + if match_info.http_exception is None: + return True, alt_request + + return False, request + + +def middleware(f: _Func) -> _Func: + f.__middleware_version__ = 1 # type: ignore[attr-defined] + return f + + +def normalize_path_middleware( + *, + append_slash: bool = True, + remove_slash: bool = False, + merge_slashes: bool = True, + redirect_class: Type[HTTPMove] = HTTPPermanentRedirect, +) -> Middleware: + """Factory for producing a middleware that normalizes the path of a request. + + Normalizing means: + - Add or remove a trailing slash to the path. + - Double slashes are replaced by one. + + The middleware returns as soon as it finds a path that resolves + correctly. The order if both merge and append/remove are enabled is + 1) merge slashes + 2) append/remove slash + 3) both merge slashes and append/remove slash. + If the path resolves with at least one of those conditions, it will + redirect to the new path. + + Only one of `append_slash` and `remove_slash` can be enabled. If both + are `True` the factory will raise an assertion error + + If `append_slash` is `True` the middleware will append a slash when + needed. If a resource is defined with trailing slash and the request + comes without it, it will append it automatically. + + If `remove_slash` is `True`, `append_slash` must be `False`. When enabled + the middleware will remove trailing slashes and redirect if the resource + is defined + + If merge_slashes is True, merge multiple consecutive slashes in the + path into one. + """ + correct_configuration = not (append_slash and remove_slash) + assert correct_configuration, "Cannot both remove and append slash" + + @middleware + async def impl(request: Request, handler: Handler) -> StreamResponse: + if isinstance(request.match_info.route, SystemRoute): + paths_to_check = [] + if "?" in request.raw_path: + path, query = request.raw_path.split("?", 1) + query = "?" + query + else: + query = "" + path = request.raw_path + + if merge_slashes: + paths_to_check.append(re.sub("//+", "/", path)) + if append_slash and not request.path.endswith("/"): + paths_to_check.append(path + "/") + if remove_slash and request.path.endswith("/"): + paths_to_check.append(path[:-1]) + if merge_slashes and append_slash: + paths_to_check.append(re.sub("//+", "/", path + "/")) + if merge_slashes and remove_slash: + merged_slashes = re.sub("//+", "/", path) + paths_to_check.append(merged_slashes[:-1]) + + for path in paths_to_check: + path = re.sub("^//+", "/", path) # SECURITY: GHSA-v6wp-4m6f-gcjg + resolves, request = await _check_request_resolves(request, path) + if resolves: + raise redirect_class(request.raw_path + query) + + return await handler(request) + + return impl + + +def _fix_request_current_app(app: "Application") -> Middleware: + @middleware + async def impl(request: Request, handler: Handler) -> StreamResponse: + match_info = request.match_info + prev = match_info.current_app + match_info.current_app = app + try: + return await handler(request) + finally: + match_info.current_app = prev + + return impl diff --git a/lib/python3.12/site-packages/aiohttp/web_protocol.py b/lib/python3.12/site-packages/aiohttp/web_protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..1bd344ae42a65445df53eb6864db1abeb06f56e7 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_protocol.py @@ -0,0 +1,792 @@ +import asyncio +import asyncio.streams +import sys +import traceback +import warnings +from collections import deque +from contextlib import suppress +from html import escape as html_escape +from http import HTTPStatus +from logging import Logger +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Deque, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, +) + +import attr +import yarl +from propcache import under_cached_property + +from .abc import AbstractAccessLogger, AbstractStreamWriter +from .base_protocol import BaseProtocol +from .helpers import ceil_timeout +from .http import ( + HttpProcessingError, + HttpRequestParser, + HttpVersion10, + RawRequestMessage, + StreamWriter, +) +from .http_exceptions import BadHttpMethod +from .log import access_logger, server_logger +from .streams import EMPTY_PAYLOAD, StreamReader +from .tcp_helpers import tcp_keepalive +from .web_exceptions import HTTPException, HTTPInternalServerError +from .web_log import AccessLogger +from .web_request import BaseRequest +from .web_response import Response, StreamResponse + +__all__ = ("RequestHandler", "RequestPayloadError", "PayloadAccessError") + +if TYPE_CHECKING: + import ssl + + from .web_server import Server + + +_RequestFactory = Callable[ + [ + RawRequestMessage, + StreamReader, + "RequestHandler", + AbstractStreamWriter, + "asyncio.Task[None]", + ], + BaseRequest, +] + +_RequestHandler = Callable[[BaseRequest], Awaitable[StreamResponse]] + +ERROR = RawRequestMessage( + "UNKNOWN", + "/", + HttpVersion10, + {}, # type: ignore[arg-type] + {}, # type: ignore[arg-type] + True, + None, + False, + False, + yarl.URL("/"), +) + + +class RequestPayloadError(Exception): + """Payload parsing error.""" + + +class PayloadAccessError(Exception): + """Payload was accessed after response was sent.""" + + +_PAYLOAD_ACCESS_ERROR = PayloadAccessError() + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class _ErrInfo: + status: int + exc: BaseException + message: str + + +_MsgType = Tuple[Union[RawRequestMessage, _ErrInfo], StreamReader] + + +class RequestHandler(BaseProtocol): + """HTTP protocol implementation. + + RequestHandler handles incoming HTTP request. It reads request line, + request headers and request payload and calls handle_request() method. + By default it always returns with 404 response. + + RequestHandler handles errors in incoming request, like bad + status line, bad headers or incomplete payload. If any error occurs, + connection gets closed. + + keepalive_timeout -- number of seconds before closing + keep-alive connection + + tcp_keepalive -- TCP keep-alive is on, default is on + + debug -- enable debug mode + + logger -- custom logger object + + access_log_class -- custom class for access_logger + + access_log -- custom logging object + + access_log_format -- access log format string + + loop -- Optional event loop + + max_line_size -- Optional maximum header line size + + max_field_size -- Optional maximum header field size + + max_headers -- Optional maximum header size + + timeout_ceil_threshold -- Optional value to specify + threshold to ceil() timeout + values + + """ + + __slots__ = ( + "_request_count", + "_keepalive", + "_manager", + "_request_handler", + "_request_factory", + "_tcp_keepalive", + "_next_keepalive_close_time", + "_keepalive_handle", + "_keepalive_timeout", + "_lingering_time", + "_messages", + "_message_tail", + "_handler_waiter", + "_waiter", + "_task_handler", + "_upgrade", + "_payload_parser", + "_request_parser", + "_reading_paused", + "logger", + "debug", + "access_log", + "access_logger", + "_close", + "_force_close", + "_current_request", + "_timeout_ceil_threshold", + "_request_in_progress", + "_logging_enabled", + "_cache", + ) + + def __init__( + self, + manager: "Server", + *, + loop: asyncio.AbstractEventLoop, + # Default should be high enough that it's likely longer than a reverse proxy. + keepalive_timeout: float = 3630, + tcp_keepalive: bool = True, + logger: Logger = server_logger, + access_log_class: Type[AbstractAccessLogger] = AccessLogger, + access_log: Logger = access_logger, + access_log_format: str = AccessLogger.LOG_FORMAT, + debug: bool = False, + max_line_size: int = 8190, + max_headers: int = 32768, + max_field_size: int = 8190, + lingering_time: float = 10.0, + read_bufsize: int = 2**16, + auto_decompress: bool = True, + timeout_ceil_threshold: float = 5, + ): + super().__init__(loop) + + # _request_count is the number of requests processed with the same connection. + self._request_count = 0 + self._keepalive = False + self._current_request: Optional[BaseRequest] = None + self._manager: Optional[Server] = manager + self._request_handler: Optional[_RequestHandler] = manager.request_handler + self._request_factory: Optional[_RequestFactory] = manager.request_factory + + self._tcp_keepalive = tcp_keepalive + # placeholder to be replaced on keepalive timeout setup + self._next_keepalive_close_time = 0.0 + self._keepalive_handle: Optional[asyncio.Handle] = None + self._keepalive_timeout = keepalive_timeout + self._lingering_time = float(lingering_time) + + self._messages: Deque[_MsgType] = deque() + self._message_tail = b"" + + self._waiter: Optional[asyncio.Future[None]] = None + self._handler_waiter: Optional[asyncio.Future[None]] = None + self._task_handler: Optional[asyncio.Task[None]] = None + + self._upgrade = False + self._payload_parser: Any = None + self._request_parser: Optional[HttpRequestParser] = HttpRequestParser( + self, + loop, + read_bufsize, + max_line_size=max_line_size, + max_field_size=max_field_size, + max_headers=max_headers, + payload_exception=RequestPayloadError, + auto_decompress=auto_decompress, + ) + + self._timeout_ceil_threshold: float = 5 + try: + self._timeout_ceil_threshold = float(timeout_ceil_threshold) + except (TypeError, ValueError): + pass + + self.logger = logger + self.debug = debug + self.access_log = access_log + if access_log: + self.access_logger: Optional[AbstractAccessLogger] = access_log_class( + access_log, access_log_format + ) + self._logging_enabled = self.access_logger.enabled + else: + self.access_logger = None + self._logging_enabled = False + + self._close = False + self._force_close = False + self._request_in_progress = False + self._cache: dict[str, Any] = {} + + def __repr__(self) -> str: + return "<{} {}>".format( + self.__class__.__name__, + "connected" if self.transport is not None else "disconnected", + ) + + @under_cached_property + def ssl_context(self) -> Optional["ssl.SSLContext"]: + """Return SSLContext if available.""" + return ( + None + if self.transport is None + else self.transport.get_extra_info("sslcontext") + ) + + @under_cached_property + def peername( + self, + ) -> Optional[Union[str, Tuple[str, int, int, int], Tuple[str, int]]]: + """Return peername if available.""" + return ( + None + if self.transport is None + else self.transport.get_extra_info("peername") + ) + + @property + def keepalive_timeout(self) -> float: + return self._keepalive_timeout + + async def shutdown(self, timeout: Optional[float] = 15.0) -> None: + """Do worker process exit preparations. + + We need to clean up everything and stop accepting requests. + It is especially important for keep-alive connections. + """ + self._force_close = True + + if self._keepalive_handle is not None: + self._keepalive_handle.cancel() + + # Wait for graceful handler completion + if self._request_in_progress: + # The future is only created when we are shutting + # down while the handler is still processing a request + # to avoid creating a future for every request. + self._handler_waiter = self._loop.create_future() + try: + async with ceil_timeout(timeout): + await self._handler_waiter + except (asyncio.CancelledError, asyncio.TimeoutError): + self._handler_waiter = None + if ( + sys.version_info >= (3, 11) + and (task := asyncio.current_task()) + and task.cancelling() + ): + raise + # Then cancel handler and wait + try: + async with ceil_timeout(timeout): + if self._current_request is not None: + self._current_request._cancel(asyncio.CancelledError()) + + if self._task_handler is not None and not self._task_handler.done(): + await asyncio.shield(self._task_handler) + except (asyncio.CancelledError, asyncio.TimeoutError): + if ( + sys.version_info >= (3, 11) + and (task := asyncio.current_task()) + and task.cancelling() + ): + raise + + # force-close non-idle handler + if self._task_handler is not None: + self._task_handler.cancel() + + self.force_close() + + def connection_made(self, transport: asyncio.BaseTransport) -> None: + super().connection_made(transport) + + real_transport = cast(asyncio.Transport, transport) + if self._tcp_keepalive: + tcp_keepalive(real_transport) + + assert self._manager is not None + self._manager.connection_made(self, real_transport) + + loop = self._loop + if sys.version_info >= (3, 12): + task = asyncio.Task(self.start(), loop=loop, eager_start=True) + else: + task = loop.create_task(self.start()) + self._task_handler = task + + def connection_lost(self, exc: Optional[BaseException]) -> None: + if self._manager is None: + return + self._manager.connection_lost(self, exc) + + # Grab value before setting _manager to None. + handler_cancellation = self._manager.handler_cancellation + + self.force_close() + super().connection_lost(exc) + self._manager = None + self._request_factory = None + self._request_handler = None + self._request_parser = None + + if self._keepalive_handle is not None: + self._keepalive_handle.cancel() + + if self._current_request is not None: + if exc is None: + exc = ConnectionResetError("Connection lost") + self._current_request._cancel(exc) + + if handler_cancellation and self._task_handler is not None: + self._task_handler.cancel() + + self._task_handler = None + + if self._payload_parser is not None: + self._payload_parser.feed_eof() + self._payload_parser = None + + def set_parser(self, parser: Any) -> None: + # Actual type is WebReader + assert self._payload_parser is None + + self._payload_parser = parser + + if self._message_tail: + self._payload_parser.feed_data(self._message_tail) + self._message_tail = b"" + + def eof_received(self) -> None: + pass + + def data_received(self, data: bytes) -> None: + if self._force_close or self._close: + return + # parse http messages + messages: Sequence[_MsgType] + if self._payload_parser is None and not self._upgrade: + assert self._request_parser is not None + try: + messages, upgraded, tail = self._request_parser.feed_data(data) + except HttpProcessingError as exc: + messages = [ + (_ErrInfo(status=400, exc=exc, message=exc.message), EMPTY_PAYLOAD) + ] + upgraded = False + tail = b"" + + for msg, payload in messages or (): + self._request_count += 1 + self._messages.append((msg, payload)) + + waiter = self._waiter + if messages and waiter is not None and not waiter.done(): + # don't set result twice + waiter.set_result(None) + + self._upgrade = upgraded + if upgraded and tail: + self._message_tail = tail + + # no parser, just store + elif self._payload_parser is None and self._upgrade and data: + self._message_tail += data + + # feed payload + elif data: + eof, tail = self._payload_parser.feed_data(data) + if eof: + self.close() + + def keep_alive(self, val: bool) -> None: + """Set keep-alive connection mode. + + :param bool val: new state. + """ + self._keepalive = val + if self._keepalive_handle: + self._keepalive_handle.cancel() + self._keepalive_handle = None + + def close(self) -> None: + """Close connection. + + Stop accepting new pipelining messages and close + connection when handlers done processing messages. + """ + self._close = True + if self._waiter: + self._waiter.cancel() + + def force_close(self) -> None: + """Forcefully close connection.""" + self._force_close = True + if self._waiter: + self._waiter.cancel() + if self.transport is not None: + self.transport.close() + self.transport = None + + def log_access( + self, request: BaseRequest, response: StreamResponse, time: Optional[float] + ) -> None: + if self._logging_enabled and self.access_logger is not None: + if TYPE_CHECKING: + assert time is not None + self.access_logger.log(request, response, self._loop.time() - time) + + def log_debug(self, *args: Any, **kw: Any) -> None: + if self.debug: + self.logger.debug(*args, **kw) + + def log_exception(self, *args: Any, **kw: Any) -> None: + self.logger.exception(*args, **kw) + + def _process_keepalive(self) -> None: + self._keepalive_handle = None + if self._force_close or not self._keepalive: + return + + loop = self._loop + now = loop.time() + close_time = self._next_keepalive_close_time + if now < close_time: + # Keep alive close check fired too early, reschedule + self._keepalive_handle = loop.call_at(close_time, self._process_keepalive) + return + + # handler in idle state + if self._waiter and not self._waiter.done(): + self.force_close() + + async def _handle_request( + self, + request: BaseRequest, + start_time: Optional[float], + request_handler: Callable[[BaseRequest], Awaitable[StreamResponse]], + ) -> Tuple[StreamResponse, bool]: + self._request_in_progress = True + try: + try: + self._current_request = request + resp = await request_handler(request) + finally: + self._current_request = None + except HTTPException as exc: + resp = exc + resp, reset = await self.finish_response(request, resp, start_time) + except asyncio.CancelledError: + raise + except asyncio.TimeoutError as exc: + self.log_debug("Request handler timed out.", exc_info=exc) + resp = self.handle_error(request, 504) + resp, reset = await self.finish_response(request, resp, start_time) + except Exception as exc: + resp = self.handle_error(request, 500, exc) + resp, reset = await self.finish_response(request, resp, start_time) + else: + # Deprecation warning (See #2415) + if getattr(resp, "__http_exception__", False): + warnings.warn( + "returning HTTPException object is deprecated " + "(#2415) and will be removed, " + "please raise the exception instead", + DeprecationWarning, + ) + + resp, reset = await self.finish_response(request, resp, start_time) + finally: + self._request_in_progress = False + if self._handler_waiter is not None: + self._handler_waiter.set_result(None) + + return resp, reset + + async def start(self) -> None: + """Process incoming request. + + It reads request line, request headers and request payload, then + calls handle_request() method. Subclass has to override + handle_request(). start() handles various exceptions in request + or response handling. Connection is being closed always unless + keep_alive(True) specified. + """ + loop = self._loop + manager = self._manager + assert manager is not None + keepalive_timeout = self._keepalive_timeout + resp = None + assert self._request_factory is not None + assert self._request_handler is not None + + while not self._force_close: + if not self._messages: + try: + # wait for next request + self._waiter = loop.create_future() + await self._waiter + finally: + self._waiter = None + + message, payload = self._messages.popleft() + + # time is only fetched if logging is enabled as otherwise + # its thrown away and never used. + start = loop.time() if self._logging_enabled else None + + manager.requests_count += 1 + writer = StreamWriter(self, loop) + if isinstance(message, _ErrInfo): + # make request_factory work + request_handler = self._make_error_handler(message) + message = ERROR + else: + request_handler = self._request_handler + + # Important don't hold a reference to the current task + # as on traceback it will prevent the task from being + # collected and will cause a memory leak. + request = self._request_factory( + message, + payload, + self, + writer, + self._task_handler or asyncio.current_task(loop), # type: ignore[arg-type] + ) + try: + # a new task is used for copy context vars (#3406) + coro = self._handle_request(request, start, request_handler) + if sys.version_info >= (3, 12): + task = asyncio.Task(coro, loop=loop, eager_start=True) + else: + task = loop.create_task(coro) + try: + resp, reset = await task + except ConnectionError: + self.log_debug("Ignored premature client disconnection") + break + + # Drop the processed task from asyncio.Task.all_tasks() early + del task + if reset: + self.log_debug("Ignored premature client disconnection 2") + break + + # notify server about keep-alive + self._keepalive = bool(resp.keep_alive) + + # check payload + if not payload.is_eof(): + lingering_time = self._lingering_time + if not self._force_close and lingering_time: + self.log_debug( + "Start lingering close timer for %s sec.", lingering_time + ) + + now = loop.time() + end_t = now + lingering_time + + try: + while not payload.is_eof() and now < end_t: + async with ceil_timeout(end_t - now): + # read and ignore + await payload.readany() + now = loop.time() + except (asyncio.CancelledError, asyncio.TimeoutError): + if ( + sys.version_info >= (3, 11) + and (t := asyncio.current_task()) + and t.cancelling() + ): + raise + + # if payload still uncompleted + if not payload.is_eof() and not self._force_close: + self.log_debug("Uncompleted request.") + self.close() + + payload.set_exception(_PAYLOAD_ACCESS_ERROR) + + except asyncio.CancelledError: + self.log_debug("Ignored premature client disconnection") + self.force_close() + raise + except Exception as exc: + self.log_exception("Unhandled exception", exc_info=exc) + self.force_close() + except BaseException: + self.force_close() + raise + finally: + request._task = None # type: ignore[assignment] # Break reference cycle in case of exception + if self.transport is None and resp is not None: + self.log_debug("Ignored premature client disconnection.") + + if self._keepalive and not self._close and not self._force_close: + # start keep-alive timer + close_time = loop.time() + keepalive_timeout + self._next_keepalive_close_time = close_time + if self._keepalive_handle is None: + self._keepalive_handle = loop.call_at( + close_time, self._process_keepalive + ) + else: + break + + # remove handler, close transport if no handlers left + if not self._force_close: + self._task_handler = None + if self.transport is not None: + self.transport.close() + + async def finish_response( + self, request: BaseRequest, resp: StreamResponse, start_time: Optional[float] + ) -> Tuple[StreamResponse, bool]: + """Prepare the response and write_eof, then log access. + + This has to + be called within the context of any exception so the access logger + can get exception information. Returns True if the client disconnects + prematurely. + """ + request._finish() + if self._request_parser is not None: + self._request_parser.set_upgraded(False) + self._upgrade = False + if self._message_tail: + self._request_parser.feed_data(self._message_tail) + self._message_tail = b"" + try: + prepare_meth = resp.prepare + except AttributeError: + if resp is None: + self.log_exception("Missing return statement on request handler") + else: + self.log_exception( + "Web-handler should return a response instance, " + "got {!r}".format(resp) + ) + exc = HTTPInternalServerError() + resp = Response( + status=exc.status, reason=exc.reason, text=exc.text, headers=exc.headers + ) + prepare_meth = resp.prepare + try: + await prepare_meth(request) + await resp.write_eof() + except ConnectionError: + self.log_access(request, resp, start_time) + return resp, True + + self.log_access(request, resp, start_time) + return resp, False + + def handle_error( + self, + request: BaseRequest, + status: int = 500, + exc: Optional[BaseException] = None, + message: Optional[str] = None, + ) -> StreamResponse: + """Handle errors. + + Returns HTTP response with specific status code. Logs additional + information. It always closes current connection. + """ + if self._request_count == 1 and isinstance(exc, BadHttpMethod): + # BadHttpMethod is common when a client sends non-HTTP + # or encrypted traffic to an HTTP port. This is expected + # to happen when connected to the public internet so we log + # it at the debug level as to not fill logs with noise. + self.logger.debug( + "Error handling request from %s", request.remote, exc_info=exc + ) + else: + self.log_exception( + "Error handling request from %s", request.remote, exc_info=exc + ) + + # some data already got sent, connection is broken + if request.writer.output_size > 0: + raise ConnectionError( + "Response is sent already, cannot send another response " + "with the error message" + ) + + ct = "text/plain" + if status == HTTPStatus.INTERNAL_SERVER_ERROR: + title = "{0.value} {0.phrase}".format(HTTPStatus.INTERNAL_SERVER_ERROR) + msg = HTTPStatus.INTERNAL_SERVER_ERROR.description + tb = None + if self.debug: + with suppress(Exception): + tb = traceback.format_exc() + + if "text/html" in request.headers.get("Accept", ""): + if tb: + tb = html_escape(tb) + msg = f"

Traceback:

\n
{tb}
" + message = ( + "" + "{title}" + "\n

{title}

" + "\n{msg}\n\n" + ).format(title=title, msg=msg) + ct = "text/html" + else: + if tb: + msg = tb + message = title + "\n\n" + msg + + resp = Response(status=status, text=message, content_type=ct) + resp.force_close() + + return resp + + def _make_error_handler( + self, err_info: _ErrInfo + ) -> Callable[[BaseRequest], Awaitable[StreamResponse]]: + async def handler(request: BaseRequest) -> StreamResponse: + return self.handle_error( + request, err_info.status, err_info.exc, err_info.message + ) + + return handler diff --git a/lib/python3.12/site-packages/aiohttp/web_request.py b/lib/python3.12/site-packages/aiohttp/web_request.py new file mode 100644 index 0000000000000000000000000000000000000000..0eafcd6e34cb8418393f1af38a072de947a63502 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_request.py @@ -0,0 +1,914 @@ +import asyncio +import datetime +import io +import re +import socket +import string +import tempfile +import types +import warnings +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Final, + Iterator, + Mapping, + MutableMapping, + Optional, + Pattern, + Tuple, + Union, + cast, +) +from urllib.parse import parse_qsl + +import attr +from multidict import ( + CIMultiDict, + CIMultiDictProxy, + MultiDict, + MultiDictProxy, + MultiMapping, +) +from yarl import URL + +from . import hdrs +from ._cookie_helpers import parse_cookie_header +from .abc import AbstractStreamWriter +from .helpers import ( + _SENTINEL, + DEBUG, + ETAG_ANY, + LIST_QUOTED_ETAG_RE, + ChainMapProxy, + ETag, + HeadersMixin, + parse_http_date, + reify, + sentinel, + set_exception, +) +from .http_parser import RawRequestMessage +from .http_writer import HttpVersion +from .multipart import BodyPartReader, MultipartReader +from .streams import EmptyStreamReader, StreamReader +from .typedefs import ( + DEFAULT_JSON_DECODER, + JSONDecoder, + LooseHeaders, + RawHeaders, + StrOrURL, +) +from .web_exceptions import HTTPRequestEntityTooLarge +from .web_response import StreamResponse + +__all__ = ("BaseRequest", "FileField", "Request") + + +if TYPE_CHECKING: + from .web_app import Application + from .web_protocol import RequestHandler + from .web_urldispatcher import UrlMappingMatchInfo + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class FileField: + name: str + filename: str + file: io.BufferedReader + content_type: str + headers: CIMultiDictProxy[str] + + +_TCHAR: Final[str] = string.digits + string.ascii_letters + r"!#$%&'*+.^_`|~-" +# '-' at the end to prevent interpretation as range in a char class + +_TOKEN: Final[str] = rf"[{_TCHAR}]+" + +_QDTEXT: Final[str] = r"[{}]".format( + r"".join(chr(c) for c in (0x09, 0x20, 0x21) + tuple(range(0x23, 0x7F))) +) +# qdtext includes 0x5C to escape 0x5D ('\]') +# qdtext excludes obs-text (because obsoleted, and encoding not specified) + +_QUOTED_PAIR: Final[str] = r"\\[\t !-~]" + +_QUOTED_STRING: Final[str] = r'"(?:{quoted_pair}|{qdtext})*"'.format( + qdtext=_QDTEXT, quoted_pair=_QUOTED_PAIR +) + +_FORWARDED_PAIR: Final[str] = ( + r"({token})=({token}|{quoted_string})(:\d{{1,4}})?".format( + token=_TOKEN, quoted_string=_QUOTED_STRING + ) +) + +_QUOTED_PAIR_REPLACE_RE: Final[Pattern[str]] = re.compile(r"\\([\t !-~])") +# same pattern as _QUOTED_PAIR but contains a capture group + +_FORWARDED_PAIR_RE: Final[Pattern[str]] = re.compile(_FORWARDED_PAIR) + +############################################################ +# HTTP Request +############################################################ + + +class BaseRequest(MutableMapping[str, Any], HeadersMixin): + + POST_METHODS = { + hdrs.METH_PATCH, + hdrs.METH_POST, + hdrs.METH_PUT, + hdrs.METH_TRACE, + hdrs.METH_DELETE, + } + + ATTRS = HeadersMixin.ATTRS | frozenset( + [ + "_message", + "_protocol", + "_payload_writer", + "_payload", + "_headers", + "_method", + "_version", + "_rel_url", + "_post", + "_read_bytes", + "_state", + "_cache", + "_task", + "_client_max_size", + "_loop", + "_transport_sslcontext", + "_transport_peername", + ] + ) + _post: Optional[MultiDictProxy[Union[str, bytes, FileField]]] = None + _read_bytes: Optional[bytes] = None + + def __init__( + self, + message: RawRequestMessage, + payload: StreamReader, + protocol: "RequestHandler", + payload_writer: AbstractStreamWriter, + task: "asyncio.Task[None]", + loop: asyncio.AbstractEventLoop, + *, + client_max_size: int = 1024**2, + state: Optional[Dict[str, Any]] = None, + scheme: Optional[str] = None, + host: Optional[str] = None, + remote: Optional[str] = None, + ) -> None: + self._message = message + self._protocol = protocol + self._payload_writer = payload_writer + + self._payload = payload + self._headers: CIMultiDictProxy[str] = message.headers + self._method = message.method + self._version = message.version + self._cache: Dict[str, Any] = {} + url = message.url + if url.absolute: + if scheme is not None: + url = url.with_scheme(scheme) + if host is not None: + url = url.with_host(host) + # absolute URL is given, + # override auto-calculating url, host, and scheme + # all other properties should be good + self._cache["url"] = url + self._cache["host"] = url.host + self._cache["scheme"] = url.scheme + self._rel_url = url.relative() + else: + self._rel_url = url + if scheme is not None: + self._cache["scheme"] = scheme + if host is not None: + self._cache["host"] = host + + self._state = {} if state is None else state + self._task = task + self._client_max_size = client_max_size + self._loop = loop + + self._transport_sslcontext = protocol.ssl_context + self._transport_peername = protocol.peername + + if remote is not None: + self._cache["remote"] = remote + + def clone( + self, + *, + method: Union[str, _SENTINEL] = sentinel, + rel_url: Union[StrOrURL, _SENTINEL] = sentinel, + headers: Union[LooseHeaders, _SENTINEL] = sentinel, + scheme: Union[str, _SENTINEL] = sentinel, + host: Union[str, _SENTINEL] = sentinel, + remote: Union[str, _SENTINEL] = sentinel, + client_max_size: Union[int, _SENTINEL] = sentinel, + ) -> "BaseRequest": + """Clone itself with replacement some attributes. + + Creates and returns a new instance of Request object. If no parameters + are given, an exact copy is returned. If a parameter is not passed, it + will reuse the one from the current request object. + """ + if self._read_bytes: + raise RuntimeError("Cannot clone request after reading its content") + + dct: Dict[str, Any] = {} + if method is not sentinel: + dct["method"] = method + if rel_url is not sentinel: + new_url: URL = URL(rel_url) + dct["url"] = new_url + dct["path"] = str(new_url) + if headers is not sentinel: + # a copy semantic + dct["headers"] = CIMultiDictProxy(CIMultiDict(headers)) + dct["raw_headers"] = tuple( + (k.encode("utf-8"), v.encode("utf-8")) + for k, v in dct["headers"].items() + ) + + message = self._message._replace(**dct) + + kwargs = {} + if scheme is not sentinel: + kwargs["scheme"] = scheme + if host is not sentinel: + kwargs["host"] = host + if remote is not sentinel: + kwargs["remote"] = remote + if client_max_size is sentinel: + client_max_size = self._client_max_size + + return self.__class__( + message, + self._payload, + self._protocol, + self._payload_writer, + self._task, + self._loop, + client_max_size=client_max_size, + state=self._state.copy(), + **kwargs, + ) + + @property + def task(self) -> "asyncio.Task[None]": + return self._task + + @property + def protocol(self) -> "RequestHandler": + return self._protocol + + @property + def transport(self) -> Optional[asyncio.Transport]: + if self._protocol is None: + return None + return self._protocol.transport + + @property + def writer(self) -> AbstractStreamWriter: + return self._payload_writer + + @property + def client_max_size(self) -> int: + return self._client_max_size + + @reify + def message(self) -> RawRequestMessage: + warnings.warn("Request.message is deprecated", DeprecationWarning, stacklevel=3) + return self._message + + @reify + def rel_url(self) -> URL: + return self._rel_url + + @reify + def loop(self) -> asyncio.AbstractEventLoop: + warnings.warn( + "request.loop property is deprecated", DeprecationWarning, stacklevel=2 + ) + return self._loop + + # MutableMapping API + + def __getitem__(self, key: str) -> Any: + return self._state[key] + + def __setitem__(self, key: str, value: Any) -> None: + self._state[key] = value + + def __delitem__(self, key: str) -> None: + del self._state[key] + + def __len__(self) -> int: + return len(self._state) + + def __iter__(self) -> Iterator[str]: + return iter(self._state) + + ######## + + @reify + def secure(self) -> bool: + """A bool indicating if the request is handled with SSL.""" + return self.scheme == "https" + + @reify + def forwarded(self) -> Tuple[Mapping[str, str], ...]: + """A tuple containing all parsed Forwarded header(s). + + Makes an effort to parse Forwarded headers as specified by RFC 7239: + + - It adds one (immutable) dictionary per Forwarded 'field-value', ie + per proxy. The element corresponds to the data in the Forwarded + field-value added by the first proxy encountered by the client. Each + subsequent item corresponds to those added by later proxies. + - It checks that every value has valid syntax in general as specified + in section 4: either a 'token' or a 'quoted-string'. + - It un-escapes found escape sequences. + - It does NOT validate 'by' and 'for' contents as specified in section + 6. + - It does NOT validate 'host' contents (Host ABNF). + - It does NOT validate 'proto' contents for valid URI scheme names. + + Returns a tuple containing one or more immutable dicts + """ + elems = [] + for field_value in self._message.headers.getall(hdrs.FORWARDED, ()): + length = len(field_value) + pos = 0 + need_separator = False + elem: Dict[str, str] = {} + elems.append(types.MappingProxyType(elem)) + while 0 <= pos < length: + match = _FORWARDED_PAIR_RE.match(field_value, pos) + if match is not None: # got a valid forwarded-pair + if need_separator: + # bad syntax here, skip to next comma + pos = field_value.find(",", pos) + else: + name, value, port = match.groups() + if value[0] == '"': + # quoted string: remove quotes and unescape + value = _QUOTED_PAIR_REPLACE_RE.sub(r"\1", value[1:-1]) + if port: + value += port + elem[name.lower()] = value + pos += len(match.group(0)) + need_separator = True + elif field_value[pos] == ",": # next forwarded-element + need_separator = False + elem = {} + elems.append(types.MappingProxyType(elem)) + pos += 1 + elif field_value[pos] == ";": # next forwarded-pair + need_separator = False + pos += 1 + elif field_value[pos] in " \t": + # Allow whitespace even between forwarded-pairs, though + # RFC 7239 doesn't. This simplifies code and is in line + # with Postel's law. + pos += 1 + else: + # bad syntax here, skip to next comma + pos = field_value.find(",", pos) + return tuple(elems) + + @reify + def scheme(self) -> str: + """A string representing the scheme of the request. + + Hostname is resolved in this order: + + - overridden value by .clone(scheme=new_scheme) call. + - type of connection to peer: HTTPS if socket is SSL, HTTP otherwise. + + 'http' or 'https'. + """ + if self._transport_sslcontext: + return "https" + else: + return "http" + + @reify + def method(self) -> str: + """Read only property for getting HTTP method. + + The value is upper-cased str like 'GET', 'POST', 'PUT' etc. + """ + return self._method + + @reify + def version(self) -> HttpVersion: + """Read only property for getting HTTP version of request. + + Returns aiohttp.protocol.HttpVersion instance. + """ + return self._version + + @reify + def host(self) -> str: + """Hostname of the request. + + Hostname is resolved in this order: + + - overridden value by .clone(host=new_host) call. + - HOST HTTP header + - socket.getfqdn() value + + For example, 'example.com' or 'localhost:8080'. + + For historical reasons, the port number may be included. + """ + host = self._message.headers.get(hdrs.HOST) + if host is not None: + return host + return socket.getfqdn() + + @reify + def remote(self) -> Optional[str]: + """Remote IP of client initiated HTTP request. + + The IP is resolved in this order: + + - overridden value by .clone(remote=new_remote) call. + - peername of opened socket + """ + if self._transport_peername is None: + return None + if isinstance(self._transport_peername, (list, tuple)): + return str(self._transport_peername[0]) + return str(self._transport_peername) + + @reify + def url(self) -> URL: + """The full URL of the request.""" + # authority is used here because it may include the port number + # and we want yarl to parse it correctly + return URL.build(scheme=self.scheme, authority=self.host).join(self._rel_url) + + @reify + def path(self) -> str: + """The URL including *PATH INFO* without the host or scheme. + + E.g., ``/app/blog`` + """ + return self._rel_url.path + + @reify + def path_qs(self) -> str: + """The URL including PATH_INFO and the query string. + + E.g, /app/blog?id=10 + """ + return str(self._rel_url) + + @reify + def raw_path(self) -> str: + """The URL including raw *PATH INFO* without the host or scheme. + + Warning, the path is unquoted and may contains non valid URL characters + + E.g., ``/my%2Fpath%7Cwith%21some%25strange%24characters`` + """ + return self._message.path + + @reify + def query(self) -> "MultiMapping[str]": + """A multidict with all the variables in the query string.""" + return self._rel_url.query + + @reify + def query_string(self) -> str: + """The query string in the URL. + + E.g., id=10 + """ + return self._rel_url.query_string + + @reify + def headers(self) -> CIMultiDictProxy[str]: + """A case-insensitive multidict proxy with all headers.""" + return self._headers + + @reify + def raw_headers(self) -> RawHeaders: + """A sequence of pairs for all headers.""" + return self._message.raw_headers + + @reify + def if_modified_since(self) -> Optional[datetime.datetime]: + """The value of If-Modified-Since HTTP header, or None. + + This header is represented as a `datetime` object. + """ + return parse_http_date(self.headers.get(hdrs.IF_MODIFIED_SINCE)) + + @reify + def if_unmodified_since(self) -> Optional[datetime.datetime]: + """The value of If-Unmodified-Since HTTP header, or None. + + This header is represented as a `datetime` object. + """ + return parse_http_date(self.headers.get(hdrs.IF_UNMODIFIED_SINCE)) + + @staticmethod + def _etag_values(etag_header: str) -> Iterator[ETag]: + """Extract `ETag` objects from raw header.""" + if etag_header == ETAG_ANY: + yield ETag( + is_weak=False, + value=ETAG_ANY, + ) + else: + for match in LIST_QUOTED_ETAG_RE.finditer(etag_header): + is_weak, value, garbage = match.group(2, 3, 4) + # Any symbol captured by 4th group means + # that the following sequence is invalid. + if garbage: + break + + yield ETag( + is_weak=bool(is_weak), + value=value, + ) + + @classmethod + def _if_match_or_none_impl( + cls, header_value: Optional[str] + ) -> Optional[Tuple[ETag, ...]]: + if not header_value: + return None + + return tuple(cls._etag_values(header_value)) + + @reify + def if_match(self) -> Optional[Tuple[ETag, ...]]: + """The value of If-Match HTTP header, or None. + + This header is represented as a `tuple` of `ETag` objects. + """ + return self._if_match_or_none_impl(self.headers.get(hdrs.IF_MATCH)) + + @reify + def if_none_match(self) -> Optional[Tuple[ETag, ...]]: + """The value of If-None-Match HTTP header, or None. + + This header is represented as a `tuple` of `ETag` objects. + """ + return self._if_match_or_none_impl(self.headers.get(hdrs.IF_NONE_MATCH)) + + @reify + def if_range(self) -> Optional[datetime.datetime]: + """The value of If-Range HTTP header, or None. + + This header is represented as a `datetime` object. + """ + return parse_http_date(self.headers.get(hdrs.IF_RANGE)) + + @reify + def keep_alive(self) -> bool: + """Is keepalive enabled by client?""" + return not self._message.should_close + + @reify + def cookies(self) -> Mapping[str, str]: + """Return request cookies. + + A read-only dictionary-like object. + """ + # Use parse_cookie_header for RFC 6265 compliant Cookie header parsing + # that accepts special characters in cookie names (fixes #2683) + parsed = parse_cookie_header(self.headers.get(hdrs.COOKIE, "")) + # Extract values from Morsel objects + return MappingProxyType({name: morsel.value for name, morsel in parsed}) + + @reify + def http_range(self) -> slice: + """The content of Range HTTP header. + + Return a slice instance. + + """ + rng = self._headers.get(hdrs.RANGE) + start, end = None, None + if rng is not None: + try: + pattern = r"^bytes=(\d*)-(\d*)$" + start, end = re.findall(pattern, rng, re.ASCII)[0] + except IndexError: # pattern was not found in header + raise ValueError("range not in acceptable format") + + end = int(end) if end else None + start = int(start) if start else None + + if start is None and end is not None: + # end with no start is to return tail of content + start = -end + end = None + + if start is not None and end is not None: + # end is inclusive in range header, exclusive for slice + end += 1 + + if start >= end: + raise ValueError("start cannot be after end") + + if start is end is None: # No valid range supplied + raise ValueError("No start or end of range specified") + + return slice(start, end, 1) + + @reify + def content(self) -> StreamReader: + """Return raw payload stream.""" + return self._payload + + @property + def has_body(self) -> bool: + """Return True if request's HTTP BODY can be read, False otherwise.""" + warnings.warn( + "Deprecated, use .can_read_body #2005", DeprecationWarning, stacklevel=2 + ) + return not self._payload.at_eof() + + @property + def can_read_body(self) -> bool: + """Return True if request's HTTP BODY can be read, False otherwise.""" + return not self._payload.at_eof() + + @reify + def body_exists(self) -> bool: + """Return True if request has HTTP BODY, False otherwise.""" + return type(self._payload) is not EmptyStreamReader + + async def release(self) -> None: + """Release request. + + Eat unread part of HTTP BODY if present. + """ + while not self._payload.at_eof(): + await self._payload.readany() + + async def read(self) -> bytes: + """Read request body if present. + + Returns bytes object with full request content. + """ + if self._read_bytes is None: + body = bytearray() + while True: + chunk = await self._payload.readany() + body.extend(chunk) + if self._client_max_size: + body_size = len(body) + if body_size >= self._client_max_size: + raise HTTPRequestEntityTooLarge( + max_size=self._client_max_size, actual_size=body_size + ) + if not chunk: + break + self._read_bytes = bytes(body) + return self._read_bytes + + async def text(self) -> str: + """Return BODY as text using encoding from .charset.""" + bytes_body = await self.read() + encoding = self.charset or "utf-8" + return bytes_body.decode(encoding) + + async def json(self, *, loads: JSONDecoder = DEFAULT_JSON_DECODER) -> Any: + """Return BODY as JSON.""" + body = await self.text() + return loads(body) + + async def multipart(self) -> MultipartReader: + """Return async iterator to process BODY as multipart.""" + return MultipartReader(self._headers, self._payload) + + async def post(self) -> "MultiDictProxy[Union[str, bytes, FileField]]": + """Return POST parameters.""" + if self._post is not None: + return self._post + if self._method not in self.POST_METHODS: + self._post = MultiDictProxy(MultiDict()) + return self._post + + content_type = self.content_type + if content_type not in ( + "", + "application/x-www-form-urlencoded", + "multipart/form-data", + ): + self._post = MultiDictProxy(MultiDict()) + return self._post + + out: MultiDict[Union[str, bytes, FileField]] = MultiDict() + + if content_type == "multipart/form-data": + multipart = await self.multipart() + max_size = self._client_max_size + + size = 0 + while (field := await multipart.next()) is not None: + field_ct = field.headers.get(hdrs.CONTENT_TYPE) + + if isinstance(field, BodyPartReader): + if field.name is None: + raise ValueError("Multipart field missing name.") + + # Note that according to RFC 7578, the Content-Type header + # is optional, even for files, so we can't assume it's + # present. + # https://tools.ietf.org/html/rfc7578#section-4.4 + if field.filename: + # store file in temp file + tmp = await self._loop.run_in_executor( + None, tempfile.TemporaryFile + ) + chunk = await field.read_chunk(size=2**16) + while chunk: + chunk = await field.decode(chunk) + await self._loop.run_in_executor(None, tmp.write, chunk) + size += len(chunk) + if 0 < max_size < size: + await self._loop.run_in_executor(None, tmp.close) + raise HTTPRequestEntityTooLarge( + max_size=max_size, actual_size=size + ) + chunk = await field.read_chunk(size=2**16) + await self._loop.run_in_executor(None, tmp.seek, 0) + + if field_ct is None: + field_ct = "application/octet-stream" + + ff = FileField( + field.name, + field.filename, + cast(io.BufferedReader, tmp), + field_ct, + field.headers, + ) + out.add(field.name, ff) + else: + # deal with ordinary data + value = await field.read(decode=True) + if field_ct is None or field_ct.startswith("text/"): + charset = field.get_charset(default="utf-8") + out.add(field.name, value.decode(charset)) + else: + out.add(field.name, value) + size += len(value) + if 0 < max_size < size: + raise HTTPRequestEntityTooLarge( + max_size=max_size, actual_size=size + ) + else: + raise ValueError( + "To decode nested multipart you need to use custom reader", + ) + else: + data = await self.read() + if data: + charset = self.charset or "utf-8" + out.extend( + parse_qsl( + data.rstrip().decode(charset), + keep_blank_values=True, + encoding=charset, + ) + ) + + self._post = MultiDictProxy(out) + return self._post + + def get_extra_info(self, name: str, default: Any = None) -> Any: + """Extra info from protocol transport""" + protocol = self._protocol + if protocol is None: + return default + + transport = protocol.transport + if transport is None: + return default + + return transport.get_extra_info(name, default) + + def __repr__(self) -> str: + ascii_encodable_path = self.path.encode("ascii", "backslashreplace").decode( + "ascii" + ) + return "<{} {} {} >".format( + self.__class__.__name__, self._method, ascii_encodable_path + ) + + def __eq__(self, other: object) -> bool: + return id(self) == id(other) + + def __bool__(self) -> bool: + return True + + async def _prepare_hook(self, response: StreamResponse) -> None: + return + + def _cancel(self, exc: BaseException) -> None: + set_exception(self._payload, exc) + + def _finish(self) -> None: + if self._post is None or self.content_type != "multipart/form-data": + return + + # NOTE: Release file descriptors for the + # NOTE: `tempfile.Temporaryfile`-created `_io.BufferedRandom` + # NOTE: instances of files sent within multipart request body + # NOTE: via HTTP POST request. + for file_name, file_field_object in self._post.items(): + if isinstance(file_field_object, FileField): + file_field_object.file.close() + + +class Request(BaseRequest): + + ATTRS = BaseRequest.ATTRS | frozenset(["_match_info"]) + + _match_info: Optional["UrlMappingMatchInfo"] = None + + if DEBUG: + + def __setattr__(self, name: str, val: Any) -> None: + if name not in self.ATTRS: + warnings.warn( + "Setting custom {}.{} attribute " + "is discouraged".format(self.__class__.__name__, name), + DeprecationWarning, + stacklevel=2, + ) + super().__setattr__(name, val) + + def clone( + self, + *, + method: Union[str, _SENTINEL] = sentinel, + rel_url: Union[StrOrURL, _SENTINEL] = sentinel, + headers: Union[LooseHeaders, _SENTINEL] = sentinel, + scheme: Union[str, _SENTINEL] = sentinel, + host: Union[str, _SENTINEL] = sentinel, + remote: Union[str, _SENTINEL] = sentinel, + client_max_size: Union[int, _SENTINEL] = sentinel, + ) -> "Request": + ret = super().clone( + method=method, + rel_url=rel_url, + headers=headers, + scheme=scheme, + host=host, + remote=remote, + client_max_size=client_max_size, + ) + new_ret = cast(Request, ret) + new_ret._match_info = self._match_info + return new_ret + + @reify + def match_info(self) -> "UrlMappingMatchInfo": + """Result of route resolving.""" + match_info = self._match_info + assert match_info is not None + return match_info + + @property + def app(self) -> "Application": + """Application instance.""" + match_info = self._match_info + assert match_info is not None + return match_info.current_app + + @property + def config_dict(self) -> ChainMapProxy: + match_info = self._match_info + assert match_info is not None + lst = match_info.apps + app = self.app + idx = lst.index(app) + sublist = list(reversed(lst[: idx + 1])) + return ChainMapProxy(sublist) + + async def _prepare_hook(self, response: StreamResponse) -> None: + match_info = self._match_info + if match_info is None: + return + for app in match_info._apps: + if on_response_prepare := app.on_response_prepare: + await on_response_prepare.send(self, response) diff --git a/lib/python3.12/site-packages/aiohttp/web_response.py b/lib/python3.12/site-packages/aiohttp/web_response.py new file mode 100644 index 0000000000000000000000000000000000000000..e5f8b6cd652c36274386cb0ce5586d8a5fdf044c --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_response.py @@ -0,0 +1,856 @@ +import asyncio +import collections.abc +import datetime +import enum +import json +import math +import time +import warnings +from concurrent.futures import Executor +from http import HTTPStatus +from http.cookies import SimpleCookie +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterator, + MutableMapping, + Optional, + Union, + cast, +) + +from multidict import CIMultiDict, istr + +from . import hdrs, payload +from .abc import AbstractStreamWriter +from .compression_utils import ZLibCompressor +from .helpers import ( + ETAG_ANY, + QUOTED_ETAG_RE, + ETag, + HeadersMixin, + must_be_empty_body, + parse_http_date, + rfc822_formatted_time, + sentinel, + should_remove_content_length, + validate_etag_value, +) +from .http import SERVER_SOFTWARE, HttpVersion10, HttpVersion11 +from .payload import Payload +from .typedefs import JSONEncoder, LooseHeaders + +REASON_PHRASES = {http_status.value: http_status.phrase for http_status in HTTPStatus} +LARGE_BODY_SIZE = 1024**2 + +__all__ = ("ContentCoding", "StreamResponse", "Response", "json_response") + + +if TYPE_CHECKING: + from .web_request import BaseRequest + + BaseClass = MutableMapping[str, Any] +else: + BaseClass = collections.abc.MutableMapping + + +# TODO(py311): Convert to StrEnum for wider use +class ContentCoding(enum.Enum): + # The content codings that we have support for. + # + # Additional registered codings are listed at: + # https://www.iana.org/assignments/http-parameters/http-parameters.xhtml#content-coding + deflate = "deflate" + gzip = "gzip" + identity = "identity" + + +CONTENT_CODINGS = {coding.value: coding for coding in ContentCoding} + +############################################################ +# HTTP Response classes +############################################################ + + +class StreamResponse(BaseClass, HeadersMixin): + + _body: Union[None, bytes, bytearray, Payload] + _length_check = True + _body = None + _keep_alive: Optional[bool] = None + _chunked: bool = False + _compression: bool = False + _compression_strategy: Optional[int] = None + _compression_force: Optional[ContentCoding] = None + _req: Optional["BaseRequest"] = None + _payload_writer: Optional[AbstractStreamWriter] = None + _eof_sent: bool = False + _must_be_empty_body: Optional[bool] = None + _body_length = 0 + _cookies: Optional[SimpleCookie] = None + _send_headers_immediately = True + + def __init__( + self, + *, + status: int = 200, + reason: Optional[str] = None, + headers: Optional[LooseHeaders] = None, + _real_headers: Optional[CIMultiDict[str]] = None, + ) -> None: + """Initialize a new stream response object. + + _real_headers is an internal parameter used to pass a pre-populated + headers object. It is used by the `Response` class to avoid copying + the headers when creating a new response object. It is not intended + to be used by external code. + """ + self._state: Dict[str, Any] = {} + + if _real_headers is not None: + self._headers = _real_headers + elif headers is not None: + self._headers: CIMultiDict[str] = CIMultiDict(headers) + else: + self._headers = CIMultiDict() + + self._set_status(status, reason) + + @property + def prepared(self) -> bool: + return self._eof_sent or self._payload_writer is not None + + @property + def task(self) -> "Optional[asyncio.Task[None]]": + if self._req: + return self._req.task + else: + return None + + @property + def status(self) -> int: + return self._status + + @property + def chunked(self) -> bool: + return self._chunked + + @property + def compression(self) -> bool: + return self._compression + + @property + def reason(self) -> str: + return self._reason + + def set_status( + self, + status: int, + reason: Optional[str] = None, + ) -> None: + assert ( + not self.prepared + ), "Cannot change the response status code after the headers have been sent" + self._set_status(status, reason) + + def _set_status(self, status: int, reason: Optional[str]) -> None: + self._status = int(status) + if reason is None: + reason = REASON_PHRASES.get(self._status, "") + elif "\n" in reason: + raise ValueError("Reason cannot contain \\n") + self._reason = reason + + @property + def keep_alive(self) -> Optional[bool]: + return self._keep_alive + + def force_close(self) -> None: + self._keep_alive = False + + @property + def body_length(self) -> int: + return self._body_length + + @property + def output_length(self) -> int: + warnings.warn("output_length is deprecated", DeprecationWarning) + assert self._payload_writer + return self._payload_writer.buffer_size + + def enable_chunked_encoding(self, chunk_size: Optional[int] = None) -> None: + """Enables automatic chunked transfer encoding.""" + if hdrs.CONTENT_LENGTH in self._headers: + raise RuntimeError( + "You can't enable chunked encoding when a content length is set" + ) + if chunk_size is not None: + warnings.warn("Chunk size is deprecated #1615", DeprecationWarning) + self._chunked = True + + def enable_compression( + self, + force: Optional[Union[bool, ContentCoding]] = None, + strategy: Optional[int] = None, + ) -> None: + """Enables response compression encoding.""" + # Backwards compatibility for when force was a bool <0.17. + if isinstance(force, bool): + force = ContentCoding.deflate if force else ContentCoding.identity + warnings.warn( + "Using boolean for force is deprecated #3318", DeprecationWarning + ) + elif force is not None: + assert isinstance( + force, ContentCoding + ), "force should one of None, bool or ContentEncoding" + + self._compression = True + self._compression_force = force + self._compression_strategy = strategy + + @property + def headers(self) -> "CIMultiDict[str]": + return self._headers + + @property + def cookies(self) -> SimpleCookie: + if self._cookies is None: + self._cookies = SimpleCookie() + return self._cookies + + def set_cookie( + self, + name: str, + value: str, + *, + expires: Optional[str] = None, + domain: Optional[str] = None, + max_age: Optional[Union[int, str]] = None, + path: str = "/", + secure: Optional[bool] = None, + httponly: Optional[bool] = None, + version: Optional[str] = None, + samesite: Optional[str] = None, + partitioned: Optional[bool] = None, + ) -> None: + """Set or update response cookie. + + Sets new cookie or updates existent with new value. + Also updates only those params which are not None. + """ + if self._cookies is None: + self._cookies = SimpleCookie() + + self._cookies[name] = value + c = self._cookies[name] + + if expires is not None: + c["expires"] = expires + elif c.get("expires") == "Thu, 01 Jan 1970 00:00:00 GMT": + del c["expires"] + + if domain is not None: + c["domain"] = domain + + if max_age is not None: + c["max-age"] = str(max_age) + elif "max-age" in c: + del c["max-age"] + + c["path"] = path + + if secure is not None: + c["secure"] = secure + if httponly is not None: + c["httponly"] = httponly + if version is not None: + c["version"] = version + if samesite is not None: + c["samesite"] = samesite + + if partitioned is not None: + c["partitioned"] = partitioned + + def del_cookie( + self, + name: str, + *, + domain: Optional[str] = None, + path: str = "/", + secure: Optional[bool] = None, + httponly: Optional[bool] = None, + samesite: Optional[str] = None, + ) -> None: + """Delete cookie. + + Creates new empty expired cookie. + """ + # TODO: do we need domain/path here? + if self._cookies is not None: + self._cookies.pop(name, None) + self.set_cookie( + name, + "", + max_age=0, + expires="Thu, 01 Jan 1970 00:00:00 GMT", + domain=domain, + path=path, + secure=secure, + httponly=httponly, + samesite=samesite, + ) + + @property + def content_length(self) -> Optional[int]: + # Just a placeholder for adding setter + return super().content_length + + @content_length.setter + def content_length(self, value: Optional[int]) -> None: + if value is not None: + value = int(value) + if self._chunked: + raise RuntimeError( + "You can't set content length when chunked encoding is enable" + ) + self._headers[hdrs.CONTENT_LENGTH] = str(value) + else: + self._headers.pop(hdrs.CONTENT_LENGTH, None) + + @property + def content_type(self) -> str: + # Just a placeholder for adding setter + return super().content_type + + @content_type.setter + def content_type(self, value: str) -> None: + self.content_type # read header values if needed + self._content_type = str(value) + self._generate_content_type_header() + + @property + def charset(self) -> Optional[str]: + # Just a placeholder for adding setter + return super().charset + + @charset.setter + def charset(self, value: Optional[str]) -> None: + ctype = self.content_type # read header values if needed + if ctype == "application/octet-stream": + raise RuntimeError( + "Setting charset for application/octet-stream " + "doesn't make sense, setup content_type first" + ) + assert self._content_dict is not None + if value is None: + self._content_dict.pop("charset", None) + else: + self._content_dict["charset"] = str(value).lower() + self._generate_content_type_header() + + @property + def last_modified(self) -> Optional[datetime.datetime]: + """The value of Last-Modified HTTP header, or None. + + This header is represented as a `datetime` object. + """ + return parse_http_date(self._headers.get(hdrs.LAST_MODIFIED)) + + @last_modified.setter + def last_modified( + self, value: Optional[Union[int, float, datetime.datetime, str]] + ) -> None: + if value is None: + self._headers.pop(hdrs.LAST_MODIFIED, None) + elif isinstance(value, (int, float)): + self._headers[hdrs.LAST_MODIFIED] = time.strftime( + "%a, %d %b %Y %H:%M:%S GMT", time.gmtime(math.ceil(value)) + ) + elif isinstance(value, datetime.datetime): + self._headers[hdrs.LAST_MODIFIED] = time.strftime( + "%a, %d %b %Y %H:%M:%S GMT", value.utctimetuple() + ) + elif isinstance(value, str): + self._headers[hdrs.LAST_MODIFIED] = value + else: + msg = f"Unsupported type for last_modified: {type(value).__name__}" + raise TypeError(msg) + + @property + def etag(self) -> Optional[ETag]: + quoted_value = self._headers.get(hdrs.ETAG) + if not quoted_value: + return None + elif quoted_value == ETAG_ANY: + return ETag(value=ETAG_ANY) + match = QUOTED_ETAG_RE.fullmatch(quoted_value) + if not match: + return None + is_weak, value = match.group(1, 2) + return ETag( + is_weak=bool(is_weak), + value=value, + ) + + @etag.setter + def etag(self, value: Optional[Union[ETag, str]]) -> None: + if value is None: + self._headers.pop(hdrs.ETAG, None) + elif (isinstance(value, str) and value == ETAG_ANY) or ( + isinstance(value, ETag) and value.value == ETAG_ANY + ): + self._headers[hdrs.ETAG] = ETAG_ANY + elif isinstance(value, str): + validate_etag_value(value) + self._headers[hdrs.ETAG] = f'"{value}"' + elif isinstance(value, ETag) and isinstance(value.value, str): + validate_etag_value(value.value) + hdr_value = f'W/"{value.value}"' if value.is_weak else f'"{value.value}"' + self._headers[hdrs.ETAG] = hdr_value + else: + raise ValueError( + f"Unsupported etag type: {type(value)}. " + f"etag must be str, ETag or None" + ) + + def _generate_content_type_header( + self, CONTENT_TYPE: istr = hdrs.CONTENT_TYPE + ) -> None: + assert self._content_dict is not None + assert self._content_type is not None + params = "; ".join(f"{k}={v}" for k, v in self._content_dict.items()) + if params: + ctype = self._content_type + "; " + params + else: + ctype = self._content_type + self._headers[CONTENT_TYPE] = ctype + + async def _do_start_compression(self, coding: ContentCoding) -> None: + if coding is ContentCoding.identity: + return + assert self._payload_writer is not None + self._headers[hdrs.CONTENT_ENCODING] = coding.value + self._payload_writer.enable_compression( + coding.value, self._compression_strategy + ) + # Compressed payload may have different content length, + # remove the header + self._headers.popall(hdrs.CONTENT_LENGTH, None) + + async def _start_compression(self, request: "BaseRequest") -> None: + if self._compression_force: + await self._do_start_compression(self._compression_force) + return + # Encoding comparisons should be case-insensitive + # https://www.rfc-editor.org/rfc/rfc9110#section-8.4.1 + accept_encoding = request.headers.get(hdrs.ACCEPT_ENCODING, "").lower() + for value, coding in CONTENT_CODINGS.items(): + if value in accept_encoding: + await self._do_start_compression(coding) + return + + async def prepare(self, request: "BaseRequest") -> Optional[AbstractStreamWriter]: + if self._eof_sent: + return None + if self._payload_writer is not None: + return self._payload_writer + self._must_be_empty_body = must_be_empty_body(request.method, self.status) + return await self._start(request) + + async def _start(self, request: "BaseRequest") -> AbstractStreamWriter: + self._req = request + writer = self._payload_writer = request._payload_writer + + await self._prepare_headers() + await request._prepare_hook(self) + await self._write_headers() + + return writer + + async def _prepare_headers(self) -> None: + request = self._req + assert request is not None + writer = self._payload_writer + assert writer is not None + keep_alive = self._keep_alive + if keep_alive is None: + keep_alive = request.keep_alive + self._keep_alive = keep_alive + + version = request.version + + headers = self._headers + if self._cookies: + for cookie in self._cookies.values(): + value = cookie.output(header="")[1:] + headers.add(hdrs.SET_COOKIE, value) + + if self._compression: + await self._start_compression(request) + + if self._chunked: + if version != HttpVersion11: + raise RuntimeError( + "Using chunked encoding is forbidden " + "for HTTP/{0.major}.{0.minor}".format(request.version) + ) + if not self._must_be_empty_body: + writer.enable_chunking() + headers[hdrs.TRANSFER_ENCODING] = "chunked" + elif self._length_check: # Disabled for WebSockets + writer.length = self.content_length + if writer.length is None: + if version >= HttpVersion11: + if not self._must_be_empty_body: + writer.enable_chunking() + headers[hdrs.TRANSFER_ENCODING] = "chunked" + elif not self._must_be_empty_body: + keep_alive = False + + # HTTP 1.1: https://tools.ietf.org/html/rfc7230#section-3.3.2 + # HTTP 1.0: https://tools.ietf.org/html/rfc1945#section-10.4 + if self._must_be_empty_body: + if hdrs.CONTENT_LENGTH in headers and should_remove_content_length( + request.method, self.status + ): + del headers[hdrs.CONTENT_LENGTH] + # https://datatracker.ietf.org/doc/html/rfc9112#section-6.1-10 + # https://datatracker.ietf.org/doc/html/rfc9112#section-6.1-13 + if hdrs.TRANSFER_ENCODING in headers: + del headers[hdrs.TRANSFER_ENCODING] + elif (writer.length if self._length_check else self.content_length) != 0: + # https://www.rfc-editor.org/rfc/rfc9110#section-8.3-5 + headers.setdefault(hdrs.CONTENT_TYPE, "application/octet-stream") + headers.setdefault(hdrs.DATE, rfc822_formatted_time()) + headers.setdefault(hdrs.SERVER, SERVER_SOFTWARE) + + # connection header + if hdrs.CONNECTION not in headers: + if keep_alive: + if version == HttpVersion10: + headers[hdrs.CONNECTION] = "keep-alive" + elif version == HttpVersion11: + headers[hdrs.CONNECTION] = "close" + + async def _write_headers(self) -> None: + request = self._req + assert request is not None + writer = self._payload_writer + assert writer is not None + # status line + version = request.version + status_line = f"HTTP/{version[0]}.{version[1]} {self._status} {self._reason}" + await writer.write_headers(status_line, self._headers) + # Send headers immediately if not opted into buffering + if self._send_headers_immediately: + writer.send_headers() + + async def write(self, data: Union[bytes, bytearray, memoryview]) -> None: + assert isinstance( + data, (bytes, bytearray, memoryview) + ), "data argument must be byte-ish (%r)" % type(data) + + if self._eof_sent: + raise RuntimeError("Cannot call write() after write_eof()") + if self._payload_writer is None: + raise RuntimeError("Cannot call write() before prepare()") + + await self._payload_writer.write(data) + + async def drain(self) -> None: + assert not self._eof_sent, "EOF has already been sent" + assert self._payload_writer is not None, "Response has not been started" + warnings.warn( + "drain method is deprecated, use await resp.write()", + DeprecationWarning, + stacklevel=2, + ) + await self._payload_writer.drain() + + async def write_eof(self, data: bytes = b"") -> None: + assert isinstance( + data, (bytes, bytearray, memoryview) + ), "data argument must be byte-ish (%r)" % type(data) + + if self._eof_sent: + return + + assert self._payload_writer is not None, "Response has not been started" + + await self._payload_writer.write_eof(data) + self._eof_sent = True + self._req = None + self._body_length = self._payload_writer.output_size + self._payload_writer = None + + def __repr__(self) -> str: + if self._eof_sent: + info = "eof" + elif self.prepared: + assert self._req is not None + info = f"{self._req.method} {self._req.path} " + else: + info = "not prepared" + return f"<{self.__class__.__name__} {self.reason} {info}>" + + def __getitem__(self, key: str) -> Any: + return self._state[key] + + def __setitem__(self, key: str, value: Any) -> None: + self._state[key] = value + + def __delitem__(self, key: str) -> None: + del self._state[key] + + def __len__(self) -> int: + return len(self._state) + + def __iter__(self) -> Iterator[str]: + return iter(self._state) + + def __hash__(self) -> int: + return hash(id(self)) + + def __eq__(self, other: object) -> bool: + return self is other + + def __bool__(self) -> bool: + return True + + +class Response(StreamResponse): + + _compressed_body: Optional[bytes] = None + _send_headers_immediately = False + + def __init__( + self, + *, + body: Any = None, + status: int = 200, + reason: Optional[str] = None, + text: Optional[str] = None, + headers: Optional[LooseHeaders] = None, + content_type: Optional[str] = None, + charset: Optional[str] = None, + zlib_executor_size: Optional[int] = None, + zlib_executor: Optional[Executor] = None, + ) -> None: + if body is not None and text is not None: + raise ValueError("body and text are not allowed together") + + if headers is None: + real_headers: CIMultiDict[str] = CIMultiDict() + else: + real_headers = CIMultiDict(headers) + + if content_type is not None and "charset" in content_type: + raise ValueError("charset must not be in content_type argument") + + if text is not None: + if hdrs.CONTENT_TYPE in real_headers: + if content_type or charset: + raise ValueError( + "passing both Content-Type header and " + "content_type or charset params " + "is forbidden" + ) + else: + # fast path for filling headers + if not isinstance(text, str): + raise TypeError("text argument must be str (%r)" % type(text)) + if content_type is None: + content_type = "text/plain" + if charset is None: + charset = "utf-8" + real_headers[hdrs.CONTENT_TYPE] = content_type + "; charset=" + charset + body = text.encode(charset) + text = None + elif hdrs.CONTENT_TYPE in real_headers: + if content_type is not None or charset is not None: + raise ValueError( + "passing both Content-Type header and " + "content_type or charset params " + "is forbidden" + ) + elif content_type is not None: + if charset is not None: + content_type += "; charset=" + charset + real_headers[hdrs.CONTENT_TYPE] = content_type + + super().__init__(status=status, reason=reason, _real_headers=real_headers) + + if text is not None: + self.text = text + else: + self.body = body + + self._zlib_executor_size = zlib_executor_size + self._zlib_executor = zlib_executor + + @property + def body(self) -> Optional[Union[bytes, Payload]]: + return self._body + + @body.setter + def body(self, body: Any) -> None: + if body is None: + self._body = None + elif isinstance(body, (bytes, bytearray)): + self._body = body + else: + try: + self._body = body = payload.PAYLOAD_REGISTRY.get(body) + except payload.LookupError: + raise ValueError("Unsupported body type %r" % type(body)) + + headers = self._headers + + # set content-type + if hdrs.CONTENT_TYPE not in headers: + headers[hdrs.CONTENT_TYPE] = body.content_type + + # copy payload headers + if body.headers: + for key, value in body.headers.items(): + if key not in headers: + headers[key] = value + + self._compressed_body = None + + @property + def text(self) -> Optional[str]: + if self._body is None: + return None + # Note: When _body is a Payload (e.g. FilePayload), this may do blocking I/O + # This is generally safe as most common payloads (BytesPayload, StringPayload) + # don't do blocking I/O, but be careful with file-based payloads + return self._body.decode(self.charset or "utf-8") + + @text.setter + def text(self, text: str) -> None: + assert text is None or isinstance( + text, str + ), "text argument must be str (%r)" % type(text) + + if self.content_type == "application/octet-stream": + self.content_type = "text/plain" + if self.charset is None: + self.charset = "utf-8" + + self._body = text.encode(self.charset) + self._compressed_body = None + + @property + def content_length(self) -> Optional[int]: + if self._chunked: + return None + + if hdrs.CONTENT_LENGTH in self._headers: + return int(self._headers[hdrs.CONTENT_LENGTH]) + + if self._compressed_body is not None: + # Return length of the compressed body + return len(self._compressed_body) + elif isinstance(self._body, Payload): + # A payload without content length, or a compressed payload + return None + elif self._body is not None: + return len(self._body) + else: + return 0 + + @content_length.setter + def content_length(self, value: Optional[int]) -> None: + raise RuntimeError("Content length is set automatically") + + async def write_eof(self, data: bytes = b"") -> None: + if self._eof_sent: + return + if self._compressed_body is None: + body: Optional[Union[bytes, Payload]] = self._body + else: + body = self._compressed_body + assert not data, f"data arg is not supported, got {data!r}" + assert self._req is not None + assert self._payload_writer is not None + if body is None or self._must_be_empty_body: + await super().write_eof() + elif isinstance(self._body, Payload): + await self._body.write(self._payload_writer) + await self._body.close() + await super().write_eof() + else: + await super().write_eof(cast(bytes, body)) + + async def _start(self, request: "BaseRequest") -> AbstractStreamWriter: + if hdrs.CONTENT_LENGTH in self._headers: + if should_remove_content_length(request.method, self.status): + del self._headers[hdrs.CONTENT_LENGTH] + elif not self._chunked: + if isinstance(self._body, Payload): + if (size := self._body.size) is not None: + self._headers[hdrs.CONTENT_LENGTH] = str(size) + else: + body_len = len(self._body) if self._body else "0" + # https://www.rfc-editor.org/rfc/rfc9110.html#section-8.6-7 + if body_len != "0" or ( + self.status != 304 and request.method not in hdrs.METH_HEAD_ALL + ): + self._headers[hdrs.CONTENT_LENGTH] = str(body_len) + + return await super()._start(request) + + async def _do_start_compression(self, coding: ContentCoding) -> None: + if self._chunked or isinstance(self._body, Payload): + return await super()._do_start_compression(coding) + if coding is ContentCoding.identity: + return + # Instead of using _payload_writer.enable_compression, + # compress the whole body + compressor = ZLibCompressor( + encoding=coding.value, + max_sync_chunk_size=self._zlib_executor_size, + executor=self._zlib_executor, + ) + assert self._body is not None + if self._zlib_executor_size is None and len(self._body) > LARGE_BODY_SIZE: + warnings.warn( + "Synchronous compression of large response bodies " + f"({len(self._body)} bytes) might block the async event loop. " + "Consider providing a custom value to zlib_executor_size/" + "zlib_executor response properties or disabling compression on it." + ) + self._compressed_body = ( + await compressor.compress(self._body) + compressor.flush() + ) + self._headers[hdrs.CONTENT_ENCODING] = coding.value + self._headers[hdrs.CONTENT_LENGTH] = str(len(self._compressed_body)) + + +def json_response( + data: Any = sentinel, + *, + text: Optional[str] = None, + body: Optional[bytes] = None, + status: int = 200, + reason: Optional[str] = None, + headers: Optional[LooseHeaders] = None, + content_type: str = "application/json", + dumps: JSONEncoder = json.dumps, +) -> Response: + if data is not sentinel: + if text or body: + raise ValueError("only one of data, text, or body should be specified") + else: + text = dumps(data) + return Response( + text=text, + body=body, + status=status, + reason=reason, + headers=headers, + content_type=content_type, + ) diff --git a/lib/python3.12/site-packages/aiohttp/web_routedef.py b/lib/python3.12/site-packages/aiohttp/web_routedef.py new file mode 100644 index 0000000000000000000000000000000000000000..f51b6cd00815a4daeabf7ef269a3225b2b764503 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_routedef.py @@ -0,0 +1,214 @@ +import abc +import os # noqa +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterator, + List, + Optional, + Sequence, + Type, + Union, + overload, +) + +import attr + +from . import hdrs +from .abc import AbstractView +from .typedefs import Handler, PathLike + +if TYPE_CHECKING: + from .web_request import Request + from .web_response import StreamResponse + from .web_urldispatcher import AbstractRoute, UrlDispatcher +else: + Request = StreamResponse = UrlDispatcher = AbstractRoute = None + + +__all__ = ( + "AbstractRouteDef", + "RouteDef", + "StaticDef", + "RouteTableDef", + "head", + "options", + "get", + "post", + "patch", + "put", + "delete", + "route", + "view", + "static", +) + + +class AbstractRouteDef(abc.ABC): + @abc.abstractmethod + def register(self, router: UrlDispatcher) -> List[AbstractRoute]: + pass # pragma: no cover + + +_HandlerType = Union[Type[AbstractView], Handler] + + +@attr.s(auto_attribs=True, frozen=True, repr=False, slots=True) +class RouteDef(AbstractRouteDef): + method: str + path: str + handler: _HandlerType + kwargs: Dict[str, Any] + + def __repr__(self) -> str: + info = [] + for name, value in sorted(self.kwargs.items()): + info.append(f", {name}={value!r}") + return " {handler.__name__!r}{info}>".format( + method=self.method, path=self.path, handler=self.handler, info="".join(info) + ) + + def register(self, router: UrlDispatcher) -> List[AbstractRoute]: + if self.method in hdrs.METH_ALL: + reg = getattr(router, "add_" + self.method.lower()) + return [reg(self.path, self.handler, **self.kwargs)] + else: + return [ + router.add_route(self.method, self.path, self.handler, **self.kwargs) + ] + + +@attr.s(auto_attribs=True, frozen=True, repr=False, slots=True) +class StaticDef(AbstractRouteDef): + prefix: str + path: PathLike + kwargs: Dict[str, Any] + + def __repr__(self) -> str: + info = [] + for name, value in sorted(self.kwargs.items()): + info.append(f", {name}={value!r}") + return " {path}{info}>".format( + prefix=self.prefix, path=self.path, info="".join(info) + ) + + def register(self, router: UrlDispatcher) -> List[AbstractRoute]: + resource = router.add_static(self.prefix, self.path, **self.kwargs) + routes = resource.get_info().get("routes", {}) + return list(routes.values()) + + +def route(method: str, path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return RouteDef(method, path, handler, kwargs) + + +def head(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_HEAD, path, handler, **kwargs) + + +def options(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_OPTIONS, path, handler, **kwargs) + + +def get( + path: str, + handler: _HandlerType, + *, + name: Optional[str] = None, + allow_head: bool = True, + **kwargs: Any, +) -> RouteDef: + return route( + hdrs.METH_GET, path, handler, name=name, allow_head=allow_head, **kwargs + ) + + +def post(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_POST, path, handler, **kwargs) + + +def put(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_PUT, path, handler, **kwargs) + + +def patch(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_PATCH, path, handler, **kwargs) + + +def delete(path: str, handler: _HandlerType, **kwargs: Any) -> RouteDef: + return route(hdrs.METH_DELETE, path, handler, **kwargs) + + +def view(path: str, handler: Type[AbstractView], **kwargs: Any) -> RouteDef: + return route(hdrs.METH_ANY, path, handler, **kwargs) + + +def static(prefix: str, path: PathLike, **kwargs: Any) -> StaticDef: + return StaticDef(prefix, path, kwargs) + + +_Deco = Callable[[_HandlerType], _HandlerType] + + +class RouteTableDef(Sequence[AbstractRouteDef]): + """Route definition table""" + + def __init__(self) -> None: + self._items: List[AbstractRouteDef] = [] + + def __repr__(self) -> str: + return f"" + + @overload + def __getitem__(self, index: int) -> AbstractRouteDef: ... + + @overload + def __getitem__(self, index: slice) -> List[AbstractRouteDef]: ... + + def __getitem__(self, index): # type: ignore[no-untyped-def] + return self._items[index] + + def __iter__(self) -> Iterator[AbstractRouteDef]: + return iter(self._items) + + def __len__(self) -> int: + return len(self._items) + + def __contains__(self, item: object) -> bool: + return item in self._items + + def route(self, method: str, path: str, **kwargs: Any) -> _Deco: + def inner(handler: _HandlerType) -> _HandlerType: + self._items.append(RouteDef(method, path, handler, kwargs)) + return handler + + return inner + + def head(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_HEAD, path, **kwargs) + + def get(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_GET, path, **kwargs) + + def post(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_POST, path, **kwargs) + + def put(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_PUT, path, **kwargs) + + def patch(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_PATCH, path, **kwargs) + + def delete(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_DELETE, path, **kwargs) + + def options(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_OPTIONS, path, **kwargs) + + def view(self, path: str, **kwargs: Any) -> _Deco: + return self.route(hdrs.METH_ANY, path, **kwargs) + + def static(self, prefix: str, path: PathLike, **kwargs: Any) -> None: + self._items.append(StaticDef(prefix, path, kwargs)) diff --git a/lib/python3.12/site-packages/aiohttp/web_runner.py b/lib/python3.12/site-packages/aiohttp/web_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..bcfec727c8419bbc6518085ecedde1f7de8992c9 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_runner.py @@ -0,0 +1,399 @@ +import asyncio +import signal +import socket +import warnings +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, List, Optional, Set + +from yarl import URL + +from .typedefs import PathLike +from .web_app import Application +from .web_server import Server + +if TYPE_CHECKING: + from ssl import SSLContext +else: + try: + from ssl import SSLContext + except ImportError: # pragma: no cover + SSLContext = object # type: ignore[misc,assignment] + +__all__ = ( + "BaseSite", + "TCPSite", + "UnixSite", + "NamedPipeSite", + "SockSite", + "BaseRunner", + "AppRunner", + "ServerRunner", + "GracefulExit", +) + + +class GracefulExit(SystemExit): + code = 1 + + +def _raise_graceful_exit() -> None: + raise GracefulExit() + + +class BaseSite(ABC): + __slots__ = ("_runner", "_ssl_context", "_backlog", "_server") + + def __init__( + self, + runner: "BaseRunner", + *, + shutdown_timeout: float = 60.0, + ssl_context: Optional[SSLContext] = None, + backlog: int = 128, + ) -> None: + if runner.server is None: + raise RuntimeError("Call runner.setup() before making a site") + if shutdown_timeout != 60.0: + msg = "shutdown_timeout should be set on BaseRunner" + warnings.warn(msg, DeprecationWarning, stacklevel=2) + runner._shutdown_timeout = shutdown_timeout + self._runner = runner + self._ssl_context = ssl_context + self._backlog = backlog + self._server: Optional[asyncio.AbstractServer] = None + + @property + @abstractmethod + def name(self) -> str: + pass # pragma: no cover + + @abstractmethod + async def start(self) -> None: + self._runner._reg_site(self) + + async def stop(self) -> None: + self._runner._check_site(self) + if self._server is not None: # Maybe not started yet + self._server.close() + + self._runner._unreg_site(self) + + +class TCPSite(BaseSite): + __slots__ = ("_host", "_port", "_reuse_address", "_reuse_port") + + def __init__( + self, + runner: "BaseRunner", + host: Optional[str] = None, + port: Optional[int] = None, + *, + shutdown_timeout: float = 60.0, + ssl_context: Optional[SSLContext] = None, + backlog: int = 128, + reuse_address: Optional[bool] = None, + reuse_port: Optional[bool] = None, + ) -> None: + super().__init__( + runner, + shutdown_timeout=shutdown_timeout, + ssl_context=ssl_context, + backlog=backlog, + ) + self._host = host + if port is None: + port = 8443 if self._ssl_context else 8080 + self._port = port + self._reuse_address = reuse_address + self._reuse_port = reuse_port + + @property + def name(self) -> str: + scheme = "https" if self._ssl_context else "http" + host = "0.0.0.0" if not self._host else self._host + return str(URL.build(scheme=scheme, host=host, port=self._port)) + + async def start(self) -> None: + await super().start() + loop = asyncio.get_event_loop() + server = self._runner.server + assert server is not None + self._server = await loop.create_server( + server, + self._host, + self._port, + ssl=self._ssl_context, + backlog=self._backlog, + reuse_address=self._reuse_address, + reuse_port=self._reuse_port, + ) + + +class UnixSite(BaseSite): + __slots__ = ("_path",) + + def __init__( + self, + runner: "BaseRunner", + path: PathLike, + *, + shutdown_timeout: float = 60.0, + ssl_context: Optional[SSLContext] = None, + backlog: int = 128, + ) -> None: + super().__init__( + runner, + shutdown_timeout=shutdown_timeout, + ssl_context=ssl_context, + backlog=backlog, + ) + self._path = path + + @property + def name(self) -> str: + scheme = "https" if self._ssl_context else "http" + return f"{scheme}://unix:{self._path}:" + + async def start(self) -> None: + await super().start() + loop = asyncio.get_event_loop() + server = self._runner.server + assert server is not None + self._server = await loop.create_unix_server( + server, + self._path, + ssl=self._ssl_context, + backlog=self._backlog, + ) + + +class NamedPipeSite(BaseSite): + __slots__ = ("_path",) + + def __init__( + self, runner: "BaseRunner", path: str, *, shutdown_timeout: float = 60.0 + ) -> None: + loop = asyncio.get_event_loop() + if not isinstance( + loop, asyncio.ProactorEventLoop # type: ignore[attr-defined] + ): + raise RuntimeError( + "Named Pipes only available in proactor loop under windows" + ) + super().__init__(runner, shutdown_timeout=shutdown_timeout) + self._path = path + + @property + def name(self) -> str: + return self._path + + async def start(self) -> None: + await super().start() + loop = asyncio.get_event_loop() + server = self._runner.server + assert server is not None + _server = await loop.start_serving_pipe( # type: ignore[attr-defined] + server, self._path + ) + self._server = _server[0] + + +class SockSite(BaseSite): + __slots__ = ("_sock", "_name") + + def __init__( + self, + runner: "BaseRunner", + sock: socket.socket, + *, + shutdown_timeout: float = 60.0, + ssl_context: Optional[SSLContext] = None, + backlog: int = 128, + ) -> None: + super().__init__( + runner, + shutdown_timeout=shutdown_timeout, + ssl_context=ssl_context, + backlog=backlog, + ) + self._sock = sock + scheme = "https" if self._ssl_context else "http" + if hasattr(socket, "AF_UNIX") and sock.family == socket.AF_UNIX: + name = f"{scheme}://unix:{sock.getsockname()}:" + else: + host, port = sock.getsockname()[:2] + name = str(URL.build(scheme=scheme, host=host, port=port)) + self._name = name + + @property + def name(self) -> str: + return self._name + + async def start(self) -> None: + await super().start() + loop = asyncio.get_event_loop() + server = self._runner.server + assert server is not None + self._server = await loop.create_server( + server, sock=self._sock, ssl=self._ssl_context, backlog=self._backlog + ) + + +class BaseRunner(ABC): + __slots__ = ("_handle_signals", "_kwargs", "_server", "_sites", "_shutdown_timeout") + + def __init__( + self, + *, + handle_signals: bool = False, + shutdown_timeout: float = 60.0, + **kwargs: Any, + ) -> None: + self._handle_signals = handle_signals + self._kwargs = kwargs + self._server: Optional[Server] = None + self._sites: List[BaseSite] = [] + self._shutdown_timeout = shutdown_timeout + + @property + def server(self) -> Optional[Server]: + return self._server + + @property + def addresses(self) -> List[Any]: + ret: List[Any] = [] + for site in self._sites: + server = site._server + if server is not None: + sockets = server.sockets # type: ignore[attr-defined] + if sockets is not None: + for sock in sockets: + ret.append(sock.getsockname()) + return ret + + @property + def sites(self) -> Set[BaseSite]: + return set(self._sites) + + async def setup(self) -> None: + loop = asyncio.get_event_loop() + + if self._handle_signals: + try: + loop.add_signal_handler(signal.SIGINT, _raise_graceful_exit) + loop.add_signal_handler(signal.SIGTERM, _raise_graceful_exit) + except NotImplementedError: # pragma: no cover + # add_signal_handler is not implemented on Windows + pass + + self._server = await self._make_server() + + @abstractmethod + async def shutdown(self) -> None: + """Call any shutdown hooks to help server close gracefully.""" + + async def cleanup(self) -> None: + # The loop over sites is intentional, an exception on gather() + # leaves self._sites in unpredictable state. + # The loop guaranties that a site is either deleted on success or + # still present on failure + for site in list(self._sites): + await site.stop() + + if self._server: # If setup succeeded + # Yield to event loop to ensure incoming requests prior to stopping the sites + # have all started to be handled before we proceed to close idle connections. + await asyncio.sleep(0) + self._server.pre_shutdown() + await self.shutdown() + await self._server.shutdown(self._shutdown_timeout) + await self._cleanup_server() + + self._server = None + if self._handle_signals: + loop = asyncio.get_running_loop() + try: + loop.remove_signal_handler(signal.SIGINT) + loop.remove_signal_handler(signal.SIGTERM) + except NotImplementedError: # pragma: no cover + # remove_signal_handler is not implemented on Windows + pass + + @abstractmethod + async def _make_server(self) -> Server: + pass # pragma: no cover + + @abstractmethod + async def _cleanup_server(self) -> None: + pass # pragma: no cover + + def _reg_site(self, site: BaseSite) -> None: + if site in self._sites: + raise RuntimeError(f"Site {site} is already registered in runner {self}") + self._sites.append(site) + + def _check_site(self, site: BaseSite) -> None: + if site not in self._sites: + raise RuntimeError(f"Site {site} is not registered in runner {self}") + + def _unreg_site(self, site: BaseSite) -> None: + if site not in self._sites: + raise RuntimeError(f"Site {site} is not registered in runner {self}") + self._sites.remove(site) + + +class ServerRunner(BaseRunner): + """Low-level web server runner""" + + __slots__ = ("_web_server",) + + def __init__( + self, web_server: Server, *, handle_signals: bool = False, **kwargs: Any + ) -> None: + super().__init__(handle_signals=handle_signals, **kwargs) + self._web_server = web_server + + async def shutdown(self) -> None: + pass + + async def _make_server(self) -> Server: + return self._web_server + + async def _cleanup_server(self) -> None: + pass + + +class AppRunner(BaseRunner): + """Web Application runner""" + + __slots__ = ("_app",) + + def __init__( + self, app: Application, *, handle_signals: bool = False, **kwargs: Any + ) -> None: + super().__init__(handle_signals=handle_signals, **kwargs) + if not isinstance(app, Application): + raise TypeError( + "The first argument should be web.Application " + "instance, got {!r}".format(app) + ) + self._app = app + + @property + def app(self) -> Application: + return self._app + + async def shutdown(self) -> None: + await self._app.shutdown() + + async def _make_server(self) -> Server: + loop = asyncio.get_event_loop() + self._app._set_loop(loop) + self._app.on_startup.freeze() + await self._app.startup() + self._app.freeze() + + return self._app._make_handler(loop=loop, **self._kwargs) + + async def _cleanup_server(self) -> None: + await self._app.cleanup() diff --git a/lib/python3.12/site-packages/aiohttp/web_server.py b/lib/python3.12/site-packages/aiohttp/web_server.py new file mode 100644 index 0000000000000000000000000000000000000000..328aca1e405ef87e4df8a992c32eac092b4af8f0 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_server.py @@ -0,0 +1,84 @@ +"""Low level HTTP server.""" + +import asyncio +from typing import Any, Awaitable, Callable, Dict, List, Optional # noqa + +from .abc import AbstractStreamWriter +from .http_parser import RawRequestMessage +from .streams import StreamReader +from .web_protocol import RequestHandler, _RequestFactory, _RequestHandler +from .web_request import BaseRequest + +__all__ = ("Server",) + + +class Server: + def __init__( + self, + handler: _RequestHandler, + *, + request_factory: Optional[_RequestFactory] = None, + handler_cancellation: bool = False, + loop: Optional[asyncio.AbstractEventLoop] = None, + **kwargs: Any, + ) -> None: + self._loop = loop or asyncio.get_running_loop() + self._connections: Dict[RequestHandler, asyncio.Transport] = {} + self._kwargs = kwargs + # requests_count is the number of requests being processed by the server + # for the lifetime of the server. + self.requests_count = 0 + self.request_handler = handler + self.request_factory = request_factory or self._make_request + self.handler_cancellation = handler_cancellation + + @property + def connections(self) -> List[RequestHandler]: + return list(self._connections.keys()) + + def connection_made( + self, handler: RequestHandler, transport: asyncio.Transport + ) -> None: + self._connections[handler] = transport + + def connection_lost( + self, handler: RequestHandler, exc: Optional[BaseException] = None + ) -> None: + if handler in self._connections: + if handler._task_handler: + handler._task_handler.add_done_callback( + lambda f: self._connections.pop(handler, None) + ) + else: + del self._connections[handler] + + def _make_request( + self, + message: RawRequestMessage, + payload: StreamReader, + protocol: RequestHandler, + writer: AbstractStreamWriter, + task: "asyncio.Task[None]", + ) -> BaseRequest: + return BaseRequest(message, payload, protocol, writer, task, self._loop) + + def pre_shutdown(self) -> None: + for conn in self._connections: + conn.close() + + async def shutdown(self, timeout: Optional[float] = None) -> None: + coros = (conn.shutdown(timeout) for conn in self._connections) + await asyncio.gather(*coros) + self._connections.clear() + + def __call__(self) -> RequestHandler: + try: + return RequestHandler(self, loop=self._loop, **self._kwargs) + except TypeError: + # Failsafe creation: remove all custom handler_args + kwargs = { + k: v + for k, v in self._kwargs.items() + if k in ["debug", "access_log_class"] + } + return RequestHandler(self, loop=self._loop, **kwargs) diff --git a/lib/python3.12/site-packages/aiohttp/web_urldispatcher.py b/lib/python3.12/site-packages/aiohttp/web_urldispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..cfa57a310046c78636d1872f6e4c2e27b6a18a76 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_urldispatcher.py @@ -0,0 +1,1305 @@ +import abc +import asyncio +import base64 +import functools +import hashlib +import html +import inspect +import keyword +import os +import platform +import re +import sys +import warnings +from functools import wraps +from pathlib import Path +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Container, + Dict, + Final, + Generator, + Iterable, + Iterator, + List, + Mapping, + NoReturn, + Optional, + Pattern, + Set, + Sized, + Tuple, + Type, + TypedDict, + Union, + cast, +) + +from yarl import URL, __version__ as yarl_version + +from . import hdrs +from .abc import AbstractMatchInfo, AbstractRouter, AbstractView +from .helpers import DEBUG +from .http import HttpVersion11 +from .typedefs import Handler, PathLike +from .web_exceptions import ( + HTTPException, + HTTPExpectationFailed, + HTTPForbidden, + HTTPMethodNotAllowed, + HTTPNotFound, +) +from .web_fileresponse import FileResponse +from .web_request import Request +from .web_response import Response, StreamResponse +from .web_routedef import AbstractRouteDef + +__all__ = ( + "UrlDispatcher", + "UrlMappingMatchInfo", + "AbstractResource", + "Resource", + "PlainResource", + "DynamicResource", + "AbstractRoute", + "ResourceRoute", + "StaticResource", + "View", +) + + +if TYPE_CHECKING: + from .web_app import Application + + BaseDict = Dict[str, str] +else: + BaseDict = dict + +CIRCULAR_SYMLINK_ERROR = ( + (OSError,) + if sys.version_info < (3, 10) and sys.platform.startswith("win32") + else (RuntimeError,) if sys.version_info < (3, 13) else () +) + +YARL_VERSION: Final[Tuple[int, ...]] = tuple(map(int, yarl_version.split(".")[:2])) + +HTTP_METHOD_RE: Final[Pattern[str]] = re.compile( + r"^[0-9A-Za-z!#\$%&'\*\+\-\.\^_`\|~]+$" +) +ROUTE_RE: Final[Pattern[str]] = re.compile( + r"(\{[_a-zA-Z][^{}]*(?:\{[^{}]*\}[^{}]*)*\})" +) +PATH_SEP: Final[str] = re.escape("/") + +IS_WINDOWS: Final[bool] = platform.system() == "Windows" + +_ExpectHandler = Callable[[Request], Awaitable[Optional[StreamResponse]]] +_Resolve = Tuple[Optional["UrlMappingMatchInfo"], Set[str]] + +html_escape = functools.partial(html.escape, quote=True) + + +class _InfoDict(TypedDict, total=False): + path: str + + formatter: str + pattern: Pattern[str] + + directory: Path + prefix: str + routes: Mapping[str, "AbstractRoute"] + + app: "Application" + + domain: str + + rule: "AbstractRuleMatching" + + http_exception: HTTPException + + +class AbstractResource(Sized, Iterable["AbstractRoute"]): + def __init__(self, *, name: Optional[str] = None) -> None: + self._name = name + + @property + def name(self) -> Optional[str]: + return self._name + + @property + @abc.abstractmethod + def canonical(self) -> str: + """Exposes the resource's canonical path. + + For example '/foo/bar/{name}' + + """ + + @abc.abstractmethod # pragma: no branch + def url_for(self, **kwargs: str) -> URL: + """Construct url for resource with additional params.""" + + @abc.abstractmethod # pragma: no branch + async def resolve(self, request: Request) -> _Resolve: + """Resolve resource. + + Return (UrlMappingMatchInfo, allowed_methods) pair. + """ + + @abc.abstractmethod + def add_prefix(self, prefix: str) -> None: + """Add a prefix to processed URLs. + + Required for subapplications support. + """ + + @abc.abstractmethod + def get_info(self) -> _InfoDict: + """Return a dict with additional info useful for introspection""" + + def freeze(self) -> None: + pass + + @abc.abstractmethod + def raw_match(self, path: str) -> bool: + """Perform a raw match against path""" + + +class AbstractRoute(abc.ABC): + def __init__( + self, + method: str, + handler: Union[Handler, Type[AbstractView]], + *, + expect_handler: Optional[_ExpectHandler] = None, + resource: Optional[AbstractResource] = None, + ) -> None: + + if expect_handler is None: + expect_handler = _default_expect_handler + + assert inspect.iscoroutinefunction(expect_handler) or ( + sys.version_info < (3, 14) and asyncio.iscoroutinefunction(expect_handler) + ), f"Coroutine is expected, got {expect_handler!r}" + + method = method.upper() + if not HTTP_METHOD_RE.match(method): + raise ValueError(f"{method} is not allowed HTTP method") + + assert callable(handler), handler + if inspect.iscoroutinefunction(handler) or ( + sys.version_info < (3, 14) and asyncio.iscoroutinefunction(handler) + ): + pass + elif inspect.isgeneratorfunction(handler): + if TYPE_CHECKING: + assert False + warnings.warn( + "Bare generators are deprecated, use @coroutine wrapper", + DeprecationWarning, + ) + elif isinstance(handler, type) and issubclass(handler, AbstractView): + pass + else: + warnings.warn( + "Bare functions are deprecated, use async ones", DeprecationWarning + ) + + @wraps(handler) + async def handler_wrapper(request: Request) -> StreamResponse: + result = old_handler(request) # type: ignore[call-arg] + if asyncio.iscoroutine(result): + result = await result + assert isinstance(result, StreamResponse) + return result + + old_handler = handler + handler = handler_wrapper + + self._method = method + self._handler = handler + self._expect_handler = expect_handler + self._resource = resource + + @property + def method(self) -> str: + return self._method + + @property + def handler(self) -> Handler: + return self._handler + + @property + @abc.abstractmethod + def name(self) -> Optional[str]: + """Optional route's name, always equals to resource's name.""" + + @property + def resource(self) -> Optional[AbstractResource]: + return self._resource + + @abc.abstractmethod + def get_info(self) -> _InfoDict: + """Return a dict with additional info useful for introspection""" + + @abc.abstractmethod # pragma: no branch + def url_for(self, *args: str, **kwargs: str) -> URL: + """Construct url for route with additional params.""" + + async def handle_expect_header(self, request: Request) -> Optional[StreamResponse]: + return await self._expect_handler(request) + + +class UrlMappingMatchInfo(BaseDict, AbstractMatchInfo): + + __slots__ = ("_route", "_apps", "_current_app", "_frozen") + + def __init__(self, match_dict: Dict[str, str], route: AbstractRoute) -> None: + super().__init__(match_dict) + self._route = route + self._apps: List[Application] = [] + self._current_app: Optional[Application] = None + self._frozen = False + + @property + def handler(self) -> Handler: + return self._route.handler + + @property + def route(self) -> AbstractRoute: + return self._route + + @property + def expect_handler(self) -> _ExpectHandler: + return self._route.handle_expect_header + + @property + def http_exception(self) -> Optional[HTTPException]: + return None + + def get_info(self) -> _InfoDict: # type: ignore[override] + return self._route.get_info() + + @property + def apps(self) -> Tuple["Application", ...]: + return tuple(self._apps) + + def add_app(self, app: "Application") -> None: + if self._frozen: + raise RuntimeError("Cannot change apps stack after .freeze() call") + if self._current_app is None: + self._current_app = app + self._apps.insert(0, app) + + @property + def current_app(self) -> "Application": + app = self._current_app + assert app is not None + return app + + @current_app.setter + def current_app(self, app: "Application") -> None: + if DEBUG: # pragma: no cover + if app not in self._apps: + raise RuntimeError( + "Expected one of the following apps {!r}, got {!r}".format( + self._apps, app + ) + ) + self._current_app = app + + def freeze(self) -> None: + self._frozen = True + + def __repr__(self) -> str: + return f"" + + +class MatchInfoError(UrlMappingMatchInfo): + + __slots__ = ("_exception",) + + def __init__(self, http_exception: HTTPException) -> None: + self._exception = http_exception + super().__init__({}, SystemRoute(self._exception)) + + @property + def http_exception(self) -> HTTPException: + return self._exception + + def __repr__(self) -> str: + return "".format( + self._exception.status, self._exception.reason + ) + + +async def _default_expect_handler(request: Request) -> None: + """Default handler for Expect header. + + Just send "100 Continue" to client. + raise HTTPExpectationFailed if value of header is not "100-continue" + """ + expect = request.headers.get(hdrs.EXPECT, "") + if request.version == HttpVersion11: + if expect.lower() == "100-continue": + await request.writer.write(b"HTTP/1.1 100 Continue\r\n\r\n") + # Reset output_size as we haven't started the main body yet. + request.writer.output_size = 0 + else: + raise HTTPExpectationFailed(text="Unknown Expect: %s" % expect) + + +class Resource(AbstractResource): + def __init__(self, *, name: Optional[str] = None) -> None: + super().__init__(name=name) + self._routes: Dict[str, ResourceRoute] = {} + self._any_route: Optional[ResourceRoute] = None + self._allowed_methods: Set[str] = set() + + def add_route( + self, + method: str, + handler: Union[Type[AbstractView], Handler], + *, + expect_handler: Optional[_ExpectHandler] = None, + ) -> "ResourceRoute": + if route := self._routes.get(method, self._any_route): + raise RuntimeError( + "Added route will never be executed, " + f"method {route.method} is already " + "registered" + ) + + route_obj = ResourceRoute(method, handler, self, expect_handler=expect_handler) + self.register_route(route_obj) + return route_obj + + def register_route(self, route: "ResourceRoute") -> None: + assert isinstance( + route, ResourceRoute + ), f"Instance of Route class is required, got {route!r}" + if route.method == hdrs.METH_ANY: + self._any_route = route + self._allowed_methods.add(route.method) + self._routes[route.method] = route + + async def resolve(self, request: Request) -> _Resolve: + if (match_dict := self._match(request.rel_url.path_safe)) is None: + return None, set() + if route := self._routes.get(request.method, self._any_route): + return UrlMappingMatchInfo(match_dict, route), self._allowed_methods + return None, self._allowed_methods + + @abc.abstractmethod + def _match(self, path: str) -> Optional[Dict[str, str]]: + pass # pragma: no cover + + def __len__(self) -> int: + return len(self._routes) + + def __iter__(self) -> Iterator["ResourceRoute"]: + return iter(self._routes.values()) + + # TODO: implement all abstract methods + + +class PlainResource(Resource): + def __init__(self, path: str, *, name: Optional[str] = None) -> None: + super().__init__(name=name) + assert not path or path.startswith("/") + self._path = path + + @property + def canonical(self) -> str: + return self._path + + def freeze(self) -> None: + if not self._path: + self._path = "/" + + def add_prefix(self, prefix: str) -> None: + assert prefix.startswith("/") + assert not prefix.endswith("/") + assert len(prefix) > 1 + self._path = prefix + self._path + + def _match(self, path: str) -> Optional[Dict[str, str]]: + # string comparison is about 10 times faster than regexp matching + if self._path == path: + return {} + return None + + def raw_match(self, path: str) -> bool: + return self._path == path + + def get_info(self) -> _InfoDict: + return {"path": self._path} + + def url_for(self) -> URL: # type: ignore[override] + return URL.build(path=self._path, encoded=True) + + def __repr__(self) -> str: + name = "'" + self.name + "' " if self.name is not None else "" + return f"" + + +class DynamicResource(Resource): + + DYN = re.compile(r"\{(?P[_a-zA-Z][_a-zA-Z0-9]*)\}") + DYN_WITH_RE = re.compile(r"\{(?P[_a-zA-Z][_a-zA-Z0-9]*):(?P.+)\}") + GOOD = r"[^{}/]+" + + def __init__(self, path: str, *, name: Optional[str] = None) -> None: + super().__init__(name=name) + self._orig_path = path + pattern = "" + formatter = "" + for part in ROUTE_RE.split(path): + match = self.DYN.fullmatch(part) + if match: + pattern += "(?P<{}>{})".format(match.group("var"), self.GOOD) + formatter += "{" + match.group("var") + "}" + continue + + match = self.DYN_WITH_RE.fullmatch(part) + if match: + pattern += "(?P<{var}>{re})".format(**match.groupdict()) + formatter += "{" + match.group("var") + "}" + continue + + if "{" in part or "}" in part: + raise ValueError(f"Invalid path '{path}'['{part}']") + + part = _requote_path(part) + formatter += part + pattern += re.escape(part) + + try: + compiled = re.compile(pattern) + except re.error as exc: + raise ValueError(f"Bad pattern '{pattern}': {exc}") from None + assert compiled.pattern.startswith(PATH_SEP) + assert formatter.startswith("/") + self._pattern = compiled + self._formatter = formatter + + @property + def canonical(self) -> str: + return self._formatter + + def add_prefix(self, prefix: str) -> None: + assert prefix.startswith("/") + assert not prefix.endswith("/") + assert len(prefix) > 1 + self._pattern = re.compile(re.escape(prefix) + self._pattern.pattern) + self._formatter = prefix + self._formatter + + def _match(self, path: str) -> Optional[Dict[str, str]]: + match = self._pattern.fullmatch(path) + if match is None: + return None + return { + key: _unquote_path_safe(value) for key, value in match.groupdict().items() + } + + def raw_match(self, path: str) -> bool: + return self._orig_path == path + + def get_info(self) -> _InfoDict: + return {"formatter": self._formatter, "pattern": self._pattern} + + def url_for(self, **parts: str) -> URL: + url = self._formatter.format_map({k: _quote_path(v) for k, v in parts.items()}) + return URL.build(path=url, encoded=True) + + def __repr__(self) -> str: + name = "'" + self.name + "' " if self.name is not None else "" + return "".format( + name=name, formatter=self._formatter + ) + + +class PrefixResource(AbstractResource): + def __init__(self, prefix: str, *, name: Optional[str] = None) -> None: + assert not prefix or prefix.startswith("/"), prefix + assert prefix in ("", "/") or not prefix.endswith("/"), prefix + super().__init__(name=name) + self._prefix = _requote_path(prefix) + self._prefix2 = self._prefix + "/" + + @property + def canonical(self) -> str: + return self._prefix + + def add_prefix(self, prefix: str) -> None: + assert prefix.startswith("/") + assert not prefix.endswith("/") + assert len(prefix) > 1 + self._prefix = prefix + self._prefix + self._prefix2 = self._prefix + "/" + + def raw_match(self, prefix: str) -> bool: + return False + + # TODO: impl missing abstract methods + + +class StaticResource(PrefixResource): + VERSION_KEY = "v" + + def __init__( + self, + prefix: str, + directory: PathLike, + *, + name: Optional[str] = None, + expect_handler: Optional[_ExpectHandler] = None, + chunk_size: int = 256 * 1024, + show_index: bool = False, + follow_symlinks: bool = False, + append_version: bool = False, + ) -> None: + super().__init__(prefix, name=name) + try: + directory = Path(directory).expanduser().resolve(strict=True) + except FileNotFoundError as error: + raise ValueError(f"'{directory}' does not exist") from error + if not directory.is_dir(): + raise ValueError(f"'{directory}' is not a directory") + self._directory = directory + self._show_index = show_index + self._chunk_size = chunk_size + self._follow_symlinks = follow_symlinks + self._expect_handler = expect_handler + self._append_version = append_version + + self._routes = { + "GET": ResourceRoute( + "GET", self._handle, self, expect_handler=expect_handler + ), + "HEAD": ResourceRoute( + "HEAD", self._handle, self, expect_handler=expect_handler + ), + } + self._allowed_methods = set(self._routes) + + def url_for( # type: ignore[override] + self, + *, + filename: PathLike, + append_version: Optional[bool] = None, + ) -> URL: + if append_version is None: + append_version = self._append_version + filename = str(filename).lstrip("/") + + url = URL.build(path=self._prefix, encoded=True) + # filename is not encoded + if YARL_VERSION < (1, 6): + url = url / filename.replace("%", "%25") + else: + url = url / filename + + if append_version: + unresolved_path = self._directory.joinpath(filename) + try: + if self._follow_symlinks: + normalized_path = Path(os.path.normpath(unresolved_path)) + normalized_path.relative_to(self._directory) + filepath = normalized_path.resolve() + else: + filepath = unresolved_path.resolve() + filepath.relative_to(self._directory) + except (ValueError, FileNotFoundError): + # ValueError for case when path point to symlink + # with follow_symlinks is False + return url # relatively safe + if filepath.is_file(): + # TODO cache file content + # with file watcher for cache invalidation + with filepath.open("rb") as f: + file_bytes = f.read() + h = self._get_file_hash(file_bytes) + url = url.with_query({self.VERSION_KEY: h}) + return url + return url + + @staticmethod + def _get_file_hash(byte_array: bytes) -> str: + m = hashlib.sha256() # todo sha256 can be configurable param + m.update(byte_array) + b64 = base64.urlsafe_b64encode(m.digest()) + return b64.decode("ascii") + + def get_info(self) -> _InfoDict: + return { + "directory": self._directory, + "prefix": self._prefix, + "routes": self._routes, + } + + def set_options_route(self, handler: Handler) -> None: + if "OPTIONS" in self._routes: + raise RuntimeError("OPTIONS route was set already") + self._routes["OPTIONS"] = ResourceRoute( + "OPTIONS", handler, self, expect_handler=self._expect_handler + ) + self._allowed_methods.add("OPTIONS") + + async def resolve(self, request: Request) -> _Resolve: + path = request.rel_url.path_safe + method = request.method + # We normalise here to avoid matches that traverse below the static root. + # e.g. /static/../../../../home/user/webapp/static/ + norm_path = os.path.normpath(path) + if IS_WINDOWS: + norm_path = norm_path.replace("\\", "/") + if not norm_path.startswith(self._prefix2) and norm_path != self._prefix: + return None, set() + + allowed_methods = self._allowed_methods + if method not in allowed_methods: + return None, allowed_methods + + match_dict = {"filename": _unquote_path_safe(path[len(self._prefix) + 1 :])} + return (UrlMappingMatchInfo(match_dict, self._routes[method]), allowed_methods) + + def __len__(self) -> int: + return len(self._routes) + + def __iter__(self) -> Iterator[AbstractRoute]: + return iter(self._routes.values()) + + async def _handle(self, request: Request) -> StreamResponse: + filename = request.match_info["filename"] + unresolved_path = self._directory.joinpath(filename) + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, self._resolve_path_to_response, unresolved_path + ) + + def _resolve_path_to_response(self, unresolved_path: Path) -> StreamResponse: + """Take the unresolved path and query the file system to form a response.""" + # Check for access outside the root directory. For follow symlinks, URI + # cannot traverse out, but symlinks can. Otherwise, no access outside + # root is permitted. + try: + if self._follow_symlinks: + normalized_path = Path(os.path.normpath(unresolved_path)) + normalized_path.relative_to(self._directory) + file_path = normalized_path.resolve() + else: + file_path = unresolved_path.resolve() + file_path.relative_to(self._directory) + except (ValueError, *CIRCULAR_SYMLINK_ERROR) as error: + # ValueError is raised for the relative check. Circular symlinks + # raise here on resolving for python < 3.13. + raise HTTPNotFound() from error + + # if path is a directory, return the contents if permitted. Note the + # directory check will raise if a segment is not readable. + try: + if file_path.is_dir(): + if self._show_index: + return Response( + text=self._directory_as_html(file_path), + content_type="text/html", + ) + else: + raise HTTPForbidden() + except PermissionError as error: + raise HTTPForbidden() from error + + # Return the file response, which handles all other checks. + return FileResponse(file_path, chunk_size=self._chunk_size) + + def _directory_as_html(self, dir_path: Path) -> str: + """returns directory's index as html.""" + assert dir_path.is_dir() + + relative_path_to_dir = dir_path.relative_to(self._directory).as_posix() + index_of = f"Index of /{html_escape(relative_path_to_dir)}" + h1 = f"

{index_of}

" + + index_list = [] + dir_index = dir_path.iterdir() + for _file in sorted(dir_index): + # show file url as relative to static path + rel_path = _file.relative_to(self._directory).as_posix() + quoted_file_url = _quote_path(f"{self._prefix}/{rel_path}") + + # if file is a directory, add '/' to the end of the name + if _file.is_dir(): + file_name = f"{_file.name}/" + else: + file_name = _file.name + + index_list.append( + f'
  • {html_escape(file_name)}
    • ' + ) + ul = "
      \n{}\n
    ".format("\n".join(index_list)) + body = f"\n{h1}\n{ul}\n" + + head_str = f"\n{index_of}\n" + html = f"\n{head_str}\n{body}\n" + + return html + + def __repr__(self) -> str: + name = "'" + self.name + "'" if self.name is not None else "" + return " {directory!r}>".format( + name=name, path=self._prefix, directory=self._directory + ) + + +class PrefixedSubAppResource(PrefixResource): + def __init__(self, prefix: str, app: "Application") -> None: + super().__init__(prefix) + self._app = app + self._add_prefix_to_resources(prefix) + + def add_prefix(self, prefix: str) -> None: + super().add_prefix(prefix) + self._add_prefix_to_resources(prefix) + + def _add_prefix_to_resources(self, prefix: str) -> None: + router = self._app.router + for resource in router.resources(): + # Since the canonical path of a resource is about + # to change, we need to unindex it and then reindex + router.unindex_resource(resource) + resource.add_prefix(prefix) + router.index_resource(resource) + + def url_for(self, *args: str, **kwargs: str) -> URL: + raise RuntimeError(".url_for() is not supported by sub-application root") + + def get_info(self) -> _InfoDict: + return {"app": self._app, "prefix": self._prefix} + + async def resolve(self, request: Request) -> _Resolve: + match_info = await self._app.router.resolve(request) + match_info.add_app(self._app) + if isinstance(match_info.http_exception, HTTPMethodNotAllowed): + methods = match_info.http_exception.allowed_methods + else: + methods = set() + return match_info, methods + + def __len__(self) -> int: + return len(self._app.router.routes()) + + def __iter__(self) -> Iterator[AbstractRoute]: + return iter(self._app.router.routes()) + + def __repr__(self) -> str: + return " {app!r}>".format( + prefix=self._prefix, app=self._app + ) + + +class AbstractRuleMatching(abc.ABC): + @abc.abstractmethod # pragma: no branch + async def match(self, request: Request) -> bool: + """Return bool if the request satisfies the criteria""" + + @abc.abstractmethod # pragma: no branch + def get_info(self) -> _InfoDict: + """Return a dict with additional info useful for introspection""" + + @property + @abc.abstractmethod # pragma: no branch + def canonical(self) -> str: + """Return a str""" + + +class Domain(AbstractRuleMatching): + re_part = re.compile(r"(?!-)[a-z\d-]{1,63}(? None: + super().__init__() + self._domain = self.validation(domain) + + @property + def canonical(self) -> str: + return self._domain + + def validation(self, domain: str) -> str: + if not isinstance(domain, str): + raise TypeError("Domain must be str") + domain = domain.rstrip(".").lower() + if not domain: + raise ValueError("Domain cannot be empty") + elif "://" in domain: + raise ValueError("Scheme not supported") + url = URL("http://" + domain) + assert url.raw_host is not None + if not all(self.re_part.fullmatch(x) for x in url.raw_host.split(".")): + raise ValueError("Domain not valid") + if url.port == 80: + return url.raw_host + return f"{url.raw_host}:{url.port}" + + async def match(self, request: Request) -> bool: + host = request.headers.get(hdrs.HOST) + if not host: + return False + return self.match_domain(host) + + def match_domain(self, host: str) -> bool: + return host.lower() == self._domain + + def get_info(self) -> _InfoDict: + return {"domain": self._domain} + + +class MaskDomain(Domain): + re_part = re.compile(r"(?!-)[a-z\d\*-]{1,63}(? None: + super().__init__(domain) + mask = self._domain.replace(".", r"\.").replace("*", ".*") + self._mask = re.compile(mask) + + @property + def canonical(self) -> str: + return self._mask.pattern + + def match_domain(self, host: str) -> bool: + return self._mask.fullmatch(host) is not None + + +class MatchedSubAppResource(PrefixedSubAppResource): + def __init__(self, rule: AbstractRuleMatching, app: "Application") -> None: + AbstractResource.__init__(self) + self._prefix = "" + self._app = app + self._rule = rule + + @property + def canonical(self) -> str: + return self._rule.canonical + + def get_info(self) -> _InfoDict: + return {"app": self._app, "rule": self._rule} + + async def resolve(self, request: Request) -> _Resolve: + if not await self._rule.match(request): + return None, set() + match_info = await self._app.router.resolve(request) + match_info.add_app(self._app) + if isinstance(match_info.http_exception, HTTPMethodNotAllowed): + methods = match_info.http_exception.allowed_methods + else: + methods = set() + return match_info, methods + + def __repr__(self) -> str: + return f" {self._app!r}>" + + +class ResourceRoute(AbstractRoute): + """A route with resource""" + + def __init__( + self, + method: str, + handler: Union[Handler, Type[AbstractView]], + resource: AbstractResource, + *, + expect_handler: Optional[_ExpectHandler] = None, + ) -> None: + super().__init__( + method, handler, expect_handler=expect_handler, resource=resource + ) + + def __repr__(self) -> str: + return " {handler!r}".format( + method=self.method, resource=self._resource, handler=self.handler + ) + + @property + def name(self) -> Optional[str]: + if self._resource is None: + return None + return self._resource.name + + def url_for(self, *args: str, **kwargs: str) -> URL: + """Construct url for route with additional params.""" + assert self._resource is not None + return self._resource.url_for(*args, **kwargs) + + def get_info(self) -> _InfoDict: + assert self._resource is not None + return self._resource.get_info() + + +class SystemRoute(AbstractRoute): + def __init__(self, http_exception: HTTPException) -> None: + super().__init__(hdrs.METH_ANY, self._handle) + self._http_exception = http_exception + + def url_for(self, *args: str, **kwargs: str) -> URL: + raise RuntimeError(".url_for() is not allowed for SystemRoute") + + @property + def name(self) -> Optional[str]: + return None + + def get_info(self) -> _InfoDict: + return {"http_exception": self._http_exception} + + async def _handle(self, request: Request) -> StreamResponse: + raise self._http_exception + + @property + def status(self) -> int: + return self._http_exception.status + + @property + def reason(self) -> str: + return self._http_exception.reason + + def __repr__(self) -> str: + return "".format(self=self) + + +class View(AbstractView): + async def _iter(self) -> StreamResponse: + if self.request.method not in hdrs.METH_ALL: + self._raise_allowed_methods() + method: Optional[Callable[[], Awaitable[StreamResponse]]] + method = getattr(self, self.request.method.lower(), None) + if method is None: + self._raise_allowed_methods() + ret = await method() + assert isinstance(ret, StreamResponse) + return ret + + def __await__(self) -> Generator[None, None, StreamResponse]: + return self._iter().__await__() + + def _raise_allowed_methods(self) -> NoReturn: + allowed_methods = {m for m in hdrs.METH_ALL if hasattr(self, m.lower())} + raise HTTPMethodNotAllowed(self.request.method, allowed_methods) + + +class ResourcesView(Sized, Iterable[AbstractResource], Container[AbstractResource]): + def __init__(self, resources: List[AbstractResource]) -> None: + self._resources = resources + + def __len__(self) -> int: + return len(self._resources) + + def __iter__(self) -> Iterator[AbstractResource]: + yield from self._resources + + def __contains__(self, resource: object) -> bool: + return resource in self._resources + + +class RoutesView(Sized, Iterable[AbstractRoute], Container[AbstractRoute]): + def __init__(self, resources: List[AbstractResource]): + self._routes: List[AbstractRoute] = [] + for resource in resources: + for route in resource: + self._routes.append(route) + + def __len__(self) -> int: + return len(self._routes) + + def __iter__(self) -> Iterator[AbstractRoute]: + yield from self._routes + + def __contains__(self, route: object) -> bool: + return route in self._routes + + +class UrlDispatcher(AbstractRouter, Mapping[str, AbstractResource]): + + NAME_SPLIT_RE = re.compile(r"[.:-]") + + def __init__(self) -> None: + super().__init__() + self._resources: List[AbstractResource] = [] + self._named_resources: Dict[str, AbstractResource] = {} + self._resource_index: dict[str, list[AbstractResource]] = {} + self._matched_sub_app_resources: List[MatchedSubAppResource] = [] + + async def resolve(self, request: Request) -> UrlMappingMatchInfo: + resource_index = self._resource_index + allowed_methods: Set[str] = set() + + # MatchedSubAppResource is primarily used to match on domain names + # (though custom rules could match on other things). This means that + # the traversal algorithm below can't be applied, and that we likely + # need to check these first so a sub app that defines the same path + # as a parent app will get priority if there's a domain match. + # + # For most cases we do not expect there to be many of these since + # currently they are only added by `.add_domain()`. + for resource in self._matched_sub_app_resources: + match_dict, allowed = await resource.resolve(request) + if match_dict is not None: + return match_dict + else: + allowed_methods |= allowed + + # Walk the url parts looking for candidates. We walk the url backwards + # to ensure the most explicit match is found first. If there are multiple + # candidates for a given url part because there are multiple resources + # registered for the same canonical path, we resolve them in a linear + # fashion to ensure registration order is respected. + url_part = request.rel_url.path_safe + while url_part: + for candidate in resource_index.get(url_part, ()): + match_dict, allowed = await candidate.resolve(request) + if match_dict is not None: + return match_dict + else: + allowed_methods |= allowed + if url_part == "/": + break + url_part = url_part.rpartition("/")[0] or "/" + + if allowed_methods: + return MatchInfoError(HTTPMethodNotAllowed(request.method, allowed_methods)) + + return MatchInfoError(HTTPNotFound()) + + def __iter__(self) -> Iterator[str]: + return iter(self._named_resources) + + def __len__(self) -> int: + return len(self._named_resources) + + def __contains__(self, resource: object) -> bool: + return resource in self._named_resources + + def __getitem__(self, name: str) -> AbstractResource: + return self._named_resources[name] + + def resources(self) -> ResourcesView: + return ResourcesView(self._resources) + + def routes(self) -> RoutesView: + return RoutesView(self._resources) + + def named_resources(self) -> Mapping[str, AbstractResource]: + return MappingProxyType(self._named_resources) + + def register_resource(self, resource: AbstractResource) -> None: + assert isinstance( + resource, AbstractResource + ), f"Instance of AbstractResource class is required, got {resource!r}" + if self.frozen: + raise RuntimeError("Cannot register a resource into frozen router.") + + name = resource.name + + if name is not None: + parts = self.NAME_SPLIT_RE.split(name) + for part in parts: + if keyword.iskeyword(part): + raise ValueError( + f"Incorrect route name {name!r}, " + "python keywords cannot be used " + "for route name" + ) + if not part.isidentifier(): + raise ValueError( + "Incorrect route name {!r}, " + "the name should be a sequence of " + "python identifiers separated " + "by dash, dot or column".format(name) + ) + if name in self._named_resources: + raise ValueError( + "Duplicate {!r}, " + "already handled by {!r}".format(name, self._named_resources[name]) + ) + self._named_resources[name] = resource + self._resources.append(resource) + + if isinstance(resource, MatchedSubAppResource): + # We cannot index match sub-app resources because they have match rules + self._matched_sub_app_resources.append(resource) + else: + self.index_resource(resource) + + def _get_resource_index_key(self, resource: AbstractResource) -> str: + """Return a key to index the resource in the resource index.""" + if "{" in (index_key := resource.canonical): + # strip at the first { to allow for variables, and than + # rpartition at / to allow for variable parts in the path + # For example if the canonical path is `/core/locations{tail:.*}` + # the index key will be `/core` since index is based on the + # url parts split by `/` + index_key = index_key.partition("{")[0].rpartition("/")[0] + return index_key.rstrip("/") or "/" + + def index_resource(self, resource: AbstractResource) -> None: + """Add a resource to the resource index.""" + resource_key = self._get_resource_index_key(resource) + # There may be multiple resources for a canonical path + # so we keep them in a list to ensure that registration + # order is respected. + self._resource_index.setdefault(resource_key, []).append(resource) + + def unindex_resource(self, resource: AbstractResource) -> None: + """Remove a resource from the resource index.""" + resource_key = self._get_resource_index_key(resource) + self._resource_index[resource_key].remove(resource) + + def add_resource(self, path: str, *, name: Optional[str] = None) -> Resource: + if path and not path.startswith("/"): + raise ValueError("path should be started with / or be empty") + # Reuse last added resource if path and name are the same + if self._resources: + resource = self._resources[-1] + if resource.name == name and resource.raw_match(path): + return cast(Resource, resource) + if not ("{" in path or "}" in path or ROUTE_RE.search(path)): + resource = PlainResource(path, name=name) + self.register_resource(resource) + return resource + resource = DynamicResource(path, name=name) + self.register_resource(resource) + return resource + + def add_route( + self, + method: str, + path: str, + handler: Union[Handler, Type[AbstractView]], + *, + name: Optional[str] = None, + expect_handler: Optional[_ExpectHandler] = None, + ) -> AbstractRoute: + resource = self.add_resource(path, name=name) + return resource.add_route(method, handler, expect_handler=expect_handler) + + def add_static( + self, + prefix: str, + path: PathLike, + *, + name: Optional[str] = None, + expect_handler: Optional[_ExpectHandler] = None, + chunk_size: int = 256 * 1024, + show_index: bool = False, + follow_symlinks: bool = False, + append_version: bool = False, + ) -> AbstractResource: + """Add static files view. + + prefix - url prefix + path - folder with files + + """ + assert prefix.startswith("/") + if prefix.endswith("/"): + prefix = prefix[:-1] + resource = StaticResource( + prefix, + path, + name=name, + expect_handler=expect_handler, + chunk_size=chunk_size, + show_index=show_index, + follow_symlinks=follow_symlinks, + append_version=append_version, + ) + self.register_resource(resource) + return resource + + def add_head(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method HEAD.""" + return self.add_route(hdrs.METH_HEAD, path, handler, **kwargs) + + def add_options(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method OPTIONS.""" + return self.add_route(hdrs.METH_OPTIONS, path, handler, **kwargs) + + def add_get( + self, + path: str, + handler: Handler, + *, + name: Optional[str] = None, + allow_head: bool = True, + **kwargs: Any, + ) -> AbstractRoute: + """Shortcut for add_route with method GET. + + If allow_head is true, another + route is added allowing head requests to the same endpoint. + """ + resource = self.add_resource(path, name=name) + if allow_head: + resource.add_route(hdrs.METH_HEAD, handler, **kwargs) + return resource.add_route(hdrs.METH_GET, handler, **kwargs) + + def add_post(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method POST.""" + return self.add_route(hdrs.METH_POST, path, handler, **kwargs) + + def add_put(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method PUT.""" + return self.add_route(hdrs.METH_PUT, path, handler, **kwargs) + + def add_patch(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method PATCH.""" + return self.add_route(hdrs.METH_PATCH, path, handler, **kwargs) + + def add_delete(self, path: str, handler: Handler, **kwargs: Any) -> AbstractRoute: + """Shortcut for add_route with method DELETE.""" + return self.add_route(hdrs.METH_DELETE, path, handler, **kwargs) + + def add_view( + self, path: str, handler: Type[AbstractView], **kwargs: Any + ) -> AbstractRoute: + """Shortcut for add_route with ANY methods for a class-based view.""" + return self.add_route(hdrs.METH_ANY, path, handler, **kwargs) + + def freeze(self) -> None: + super().freeze() + for resource in self._resources: + resource.freeze() + + def add_routes(self, routes: Iterable[AbstractRouteDef]) -> List[AbstractRoute]: + """Append routes to route table. + + Parameter should be a sequence of RouteDef objects. + + Returns a list of registered AbstractRoute instances. + """ + registered_routes = [] + for route_def in routes: + registered_routes.extend(route_def.register(self)) + return registered_routes + + +def _quote_path(value: str) -> str: + if YARL_VERSION < (1, 6): + value = value.replace("%", "%25") + return URL.build(path=value, encoded=False).raw_path + + +def _unquote_path_safe(value: str) -> str: + if "%" not in value: + return value + return value.replace("%2F", "/").replace("%25", "%") + + +def _requote_path(value: str) -> str: + # Quote non-ascii characters and other characters which must be quoted, + # but preserve existing %-sequences. + result = _quote_path(value) + if "%" in value: + result = result.replace("%25", "%") + return result diff --git a/lib/python3.12/site-packages/aiohttp/web_ws.py b/lib/python3.12/site-packages/aiohttp/web_ws.py new file mode 100644 index 0000000000000000000000000000000000000000..575f9a3dc8507d1e6b766333c9daec389313febd --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/web_ws.py @@ -0,0 +1,631 @@ +import asyncio +import base64 +import binascii +import hashlib +import json +import sys +from typing import Any, Final, Iterable, Optional, Tuple, Union, cast + +import attr +from multidict import CIMultiDict + +from . import hdrs +from ._websocket.reader import WebSocketDataQueue +from ._websocket.writer import DEFAULT_LIMIT +from .abc import AbstractStreamWriter +from .client_exceptions import WSMessageTypeError +from .helpers import calculate_timeout_when, set_exception, set_result +from .http import ( + WS_CLOSED_MESSAGE, + WS_CLOSING_MESSAGE, + WS_KEY, + WebSocketError, + WebSocketReader, + WebSocketWriter, + WSCloseCode, + WSMessage, + WSMsgType as WSMsgType, + ws_ext_gen, + ws_ext_parse, +) +from .http_websocket import _INTERNAL_RECEIVE_TYPES +from .log import ws_logger +from .streams import EofStream +from .typedefs import JSONDecoder, JSONEncoder +from .web_exceptions import HTTPBadRequest, HTTPException +from .web_request import BaseRequest +from .web_response import StreamResponse + +if sys.version_info >= (3, 11): + import asyncio as async_timeout +else: + import async_timeout + +__all__ = ( + "WebSocketResponse", + "WebSocketReady", + "WSMsgType", +) + +THRESHOLD_CONNLOST_ACCESS: Final[int] = 5 + + +@attr.s(auto_attribs=True, frozen=True, slots=True) +class WebSocketReady: + ok: bool + protocol: Optional[str] + + def __bool__(self) -> bool: + return self.ok + + +class WebSocketResponse(StreamResponse): + + _length_check: bool = False + _ws_protocol: Optional[str] = None + _writer: Optional[WebSocketWriter] = None + _reader: Optional[WebSocketDataQueue] = None + _closed: bool = False + _closing: bool = False + _conn_lost: int = 0 + _close_code: Optional[int] = None + _loop: Optional[asyncio.AbstractEventLoop] = None + _waiting: bool = False + _close_wait: Optional[asyncio.Future[None]] = None + _exception: Optional[BaseException] = None + _heartbeat_when: float = 0.0 + _heartbeat_cb: Optional[asyncio.TimerHandle] = None + _pong_response_cb: Optional[asyncio.TimerHandle] = None + _ping_task: Optional[asyncio.Task[None]] = None + + def __init__( + self, + *, + timeout: float = 10.0, + receive_timeout: Optional[float] = None, + autoclose: bool = True, + autoping: bool = True, + heartbeat: Optional[float] = None, + protocols: Iterable[str] = (), + compress: bool = True, + max_msg_size: int = 4 * 1024 * 1024, + writer_limit: int = DEFAULT_LIMIT, + ) -> None: + super().__init__(status=101) + self._protocols = protocols + self._timeout = timeout + self._receive_timeout = receive_timeout + self._autoclose = autoclose + self._autoping = autoping + self._heartbeat = heartbeat + if heartbeat is not None: + self._pong_heartbeat = heartbeat / 2.0 + self._compress: Union[bool, int] = compress + self._max_msg_size = max_msg_size + self._writer_limit = writer_limit + + def _cancel_heartbeat(self) -> None: + self._cancel_pong_response_cb() + if self._heartbeat_cb is not None: + self._heartbeat_cb.cancel() + self._heartbeat_cb = None + if self._ping_task is not None: + self._ping_task.cancel() + self._ping_task = None + + def _cancel_pong_response_cb(self) -> None: + if self._pong_response_cb is not None: + self._pong_response_cb.cancel() + self._pong_response_cb = None + + def _reset_heartbeat(self) -> None: + if self._heartbeat is None: + return + self._cancel_pong_response_cb() + req = self._req + timeout_ceil_threshold = ( + req._protocol._timeout_ceil_threshold if req is not None else 5 + ) + loop = self._loop + assert loop is not None + now = loop.time() + when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold) + self._heartbeat_when = when + if self._heartbeat_cb is None: + # We do not cancel the previous heartbeat_cb here because + # it generates a significant amount of TimerHandle churn + # which causes asyncio to rebuild the heap frequently. + # Instead _send_heartbeat() will reschedule the next + # heartbeat if it fires too early. + self._heartbeat_cb = loop.call_at(when, self._send_heartbeat) + + def _send_heartbeat(self) -> None: + self._heartbeat_cb = None + loop = self._loop + assert loop is not None and self._writer is not None + now = loop.time() + if now < self._heartbeat_when: + # Heartbeat fired too early, reschedule + self._heartbeat_cb = loop.call_at( + self._heartbeat_when, self._send_heartbeat + ) + return + + req = self._req + timeout_ceil_threshold = ( + req._protocol._timeout_ceil_threshold if req is not None else 5 + ) + when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold) + self._cancel_pong_response_cb() + self._pong_response_cb = loop.call_at(when, self._pong_not_received) + + coro = self._writer.send_frame(b"", WSMsgType.PING) + if sys.version_info >= (3, 12): + # Optimization for Python 3.12, try to send the ping + # immediately to avoid having to schedule + # the task on the event loop. + ping_task = asyncio.Task(coro, loop=loop, eager_start=True) + else: + ping_task = loop.create_task(coro) + + if not ping_task.done(): + self._ping_task = ping_task + ping_task.add_done_callback(self._ping_task_done) + else: + self._ping_task_done(ping_task) + + def _ping_task_done(self, task: "asyncio.Task[None]") -> None: + """Callback for when the ping task completes.""" + if not task.cancelled() and (exc := task.exception()): + self._handle_ping_pong_exception(exc) + self._ping_task = None + + def _pong_not_received(self) -> None: + if self._req is not None and self._req.transport is not None: + self._handle_ping_pong_exception( + asyncio.TimeoutError( + f"No PONG received after {self._pong_heartbeat} seconds" + ) + ) + + def _handle_ping_pong_exception(self, exc: BaseException) -> None: + """Handle exceptions raised during ping/pong processing.""" + if self._closed: + return + self._set_closed() + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + self._exception = exc + if self._waiting and not self._closing and self._reader is not None: + self._reader.feed_data(WSMessage(WSMsgType.ERROR, exc, None), 0) + + def _set_closed(self) -> None: + """Set the connection to closed. + + Cancel any heartbeat timers and set the closed flag. + """ + self._closed = True + self._cancel_heartbeat() + + async def prepare(self, request: BaseRequest) -> AbstractStreamWriter: + # make pre-check to don't hide it by do_handshake() exceptions + if self._payload_writer is not None: + return self._payload_writer + + protocol, writer = self._pre_start(request) + payload_writer = await super().prepare(request) + assert payload_writer is not None + self._post_start(request, protocol, writer) + await payload_writer.drain() + return payload_writer + + def _handshake( + self, request: BaseRequest + ) -> Tuple["CIMultiDict[str]", Optional[str], int, bool]: + headers = request.headers + if "websocket" != headers.get(hdrs.UPGRADE, "").lower().strip(): + raise HTTPBadRequest( + text=( + "No WebSocket UPGRADE hdr: {}\n Can " + '"Upgrade" only to "WebSocket".' + ).format(headers.get(hdrs.UPGRADE)) + ) + + if "upgrade" not in headers.get(hdrs.CONNECTION, "").lower(): + raise HTTPBadRequest( + text="No CONNECTION upgrade hdr: {}".format( + headers.get(hdrs.CONNECTION) + ) + ) + + # find common sub-protocol between client and server + protocol: Optional[str] = None + if hdrs.SEC_WEBSOCKET_PROTOCOL in headers: + req_protocols = [ + str(proto.strip()) + for proto in headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(",") + ] + + for proto in req_protocols: + if proto in self._protocols: + protocol = proto + break + else: + # No overlap found: Return no protocol as per spec + ws_logger.warning( + "%s: Client protocols %r don’t overlap server-known ones %r", + request.remote, + req_protocols, + self._protocols, + ) + + # check supported version + version = headers.get(hdrs.SEC_WEBSOCKET_VERSION, "") + if version not in ("13", "8", "7"): + raise HTTPBadRequest(text=f"Unsupported version: {version}") + + # check client handshake for validity + key = headers.get(hdrs.SEC_WEBSOCKET_KEY) + try: + if not key or len(base64.b64decode(key)) != 16: + raise HTTPBadRequest(text=f"Handshake error: {key!r}") + except binascii.Error: + raise HTTPBadRequest(text=f"Handshake error: {key!r}") from None + + accept_val = base64.b64encode( + hashlib.sha1(key.encode() + WS_KEY).digest() + ).decode() + response_headers = CIMultiDict( + { + hdrs.UPGRADE: "websocket", + hdrs.CONNECTION: "upgrade", + hdrs.SEC_WEBSOCKET_ACCEPT: accept_val, + } + ) + + notakeover = False + compress = 0 + if self._compress: + extensions = headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS) + # Server side always get return with no exception. + # If something happened, just drop compress extension + compress, notakeover = ws_ext_parse(extensions, isserver=True) + if compress: + enabledext = ws_ext_gen( + compress=compress, isserver=True, server_notakeover=notakeover + ) + response_headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = enabledext + + if protocol: + response_headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = protocol + return ( + response_headers, + protocol, + compress, + notakeover, + ) + + def _pre_start(self, request: BaseRequest) -> Tuple[Optional[str], WebSocketWriter]: + self._loop = request._loop + + headers, protocol, compress, notakeover = self._handshake(request) + + self.set_status(101) + self.headers.update(headers) + self.force_close() + self._compress = compress + transport = request._protocol.transport + assert transport is not None + writer = WebSocketWriter( + request._protocol, + transport, + compress=compress, + notakeover=notakeover, + limit=self._writer_limit, + ) + + return protocol, writer + + def _post_start( + self, request: BaseRequest, protocol: Optional[str], writer: WebSocketWriter + ) -> None: + self._ws_protocol = protocol + self._writer = writer + + self._reset_heartbeat() + + loop = self._loop + assert loop is not None + self._reader = WebSocketDataQueue(request._protocol, 2**16, loop=loop) + request.protocol.set_parser( + WebSocketReader( + self._reader, self._max_msg_size, compress=bool(self._compress) + ) + ) + # disable HTTP keepalive for WebSocket + request.protocol.keep_alive(False) + + def can_prepare(self, request: BaseRequest) -> WebSocketReady: + if self._writer is not None: + raise RuntimeError("Already started") + try: + _, protocol, _, _ = self._handshake(request) + except HTTPException: + return WebSocketReady(False, None) + else: + return WebSocketReady(True, protocol) + + @property + def prepared(self) -> bool: + return self._writer is not None + + @property + def closed(self) -> bool: + return self._closed + + @property + def close_code(self) -> Optional[int]: + return self._close_code + + @property + def ws_protocol(self) -> Optional[str]: + return self._ws_protocol + + @property + def compress(self) -> Union[int, bool]: + return self._compress + + def get_extra_info(self, name: str, default: Any = None) -> Any: + """Get optional transport information. + + If no value associated with ``name`` is found, ``default`` is returned. + """ + writer = self._writer + if writer is None: + return default + transport = writer.transport + if transport is None: + return default + return transport.get_extra_info(name, default) + + def exception(self) -> Optional[BaseException]: + return self._exception + + async def ping(self, message: bytes = b"") -> None: + if self._writer is None: + raise RuntimeError("Call .prepare() first") + await self._writer.send_frame(message, WSMsgType.PING) + + async def pong(self, message: bytes = b"") -> None: + # unsolicited pong + if self._writer is None: + raise RuntimeError("Call .prepare() first") + await self._writer.send_frame(message, WSMsgType.PONG) + + async def send_frame( + self, message: bytes, opcode: WSMsgType, compress: Optional[int] = None + ) -> None: + """Send a frame over the websocket.""" + if self._writer is None: + raise RuntimeError("Call .prepare() first") + await self._writer.send_frame(message, opcode, compress) + + async def send_str(self, data: str, compress: Optional[int] = None) -> None: + if self._writer is None: + raise RuntimeError("Call .prepare() first") + if not isinstance(data, str): + raise TypeError("data argument must be str (%r)" % type(data)) + await self._writer.send_frame( + data.encode("utf-8"), WSMsgType.TEXT, compress=compress + ) + + async def send_bytes(self, data: bytes, compress: Optional[int] = None) -> None: + if self._writer is None: + raise RuntimeError("Call .prepare() first") + if not isinstance(data, (bytes, bytearray, memoryview)): + raise TypeError("data argument must be byte-ish (%r)" % type(data)) + await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress) + + async def send_json( + self, + data: Any, + compress: Optional[int] = None, + *, + dumps: JSONEncoder = json.dumps, + ) -> None: + await self.send_str(dumps(data), compress=compress) + + async def write_eof(self) -> None: # type: ignore[override] + if self._eof_sent: + return + if self._payload_writer is None: + raise RuntimeError("Response has not been started") + + await self.close() + self._eof_sent = True + + async def close( + self, *, code: int = WSCloseCode.OK, message: bytes = b"", drain: bool = True + ) -> bool: + """Close websocket connection.""" + if self._writer is None: + raise RuntimeError("Call .prepare() first") + + if self._closed: + return False + self._set_closed() + + try: + await self._writer.close(code, message) + writer = self._payload_writer + assert writer is not None + if drain: + await writer.drain() + except (asyncio.CancelledError, asyncio.TimeoutError): + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + raise + except Exception as exc: + self._exception = exc + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + return True + + reader = self._reader + assert reader is not None + # we need to break `receive()` cycle before we can call + # `reader.read()` as `close()` may be called from different task + if self._waiting: + assert self._loop is not None + assert self._close_wait is None + self._close_wait = self._loop.create_future() + reader.feed_data(WS_CLOSING_MESSAGE, 0) + await self._close_wait + + if self._closing: + self._close_transport() + return True + + try: + async with async_timeout.timeout(self._timeout): + while True: + msg = await reader.read() + if msg.type is WSMsgType.CLOSE: + self._set_code_close_transport(msg.data) + return True + except asyncio.CancelledError: + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + raise + except Exception as exc: + self._exception = exc + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + return True + + def _set_closing(self, code: WSCloseCode) -> None: + """Set the close code and mark the connection as closing.""" + self._closing = True + self._close_code = code + self._cancel_heartbeat() + + def _set_code_close_transport(self, code: WSCloseCode) -> None: + """Set the close code and close the transport.""" + self._close_code = code + self._close_transport() + + def _close_transport(self) -> None: + """Close the transport.""" + if self._req is not None and self._req.transport is not None: + self._req.transport.close() + + async def receive(self, timeout: Optional[float] = None) -> WSMessage: + if self._reader is None: + raise RuntimeError("Call .prepare() first") + + receive_timeout = timeout or self._receive_timeout + while True: + if self._waiting: + raise RuntimeError("Concurrent call to receive() is not allowed") + + if self._closed: + self._conn_lost += 1 + if self._conn_lost >= THRESHOLD_CONNLOST_ACCESS: + raise RuntimeError("WebSocket connection is closed.") + return WS_CLOSED_MESSAGE + elif self._closing: + return WS_CLOSING_MESSAGE + + try: + self._waiting = True + try: + if receive_timeout: + # Entering the context manager and creating + # Timeout() object can take almost 50% of the + # run time in this loop so we avoid it if + # there is no read timeout. + async with async_timeout.timeout(receive_timeout): + msg = await self._reader.read() + else: + msg = await self._reader.read() + self._reset_heartbeat() + finally: + self._waiting = False + if self._close_wait: + set_result(self._close_wait, None) + except asyncio.TimeoutError: + raise + except EofStream: + self._close_code = WSCloseCode.OK + await self.close() + return WSMessage(WSMsgType.CLOSED, None, None) + except WebSocketError as exc: + self._close_code = exc.code + await self.close(code=exc.code) + return WSMessage(WSMsgType.ERROR, exc, None) + except Exception as exc: + self._exception = exc + self._set_closing(WSCloseCode.ABNORMAL_CLOSURE) + await self.close() + return WSMessage(WSMsgType.ERROR, exc, None) + + if msg.type not in _INTERNAL_RECEIVE_TYPES: + # If its not a close/closing/ping/pong message + # we can return it immediately + return msg + + if msg.type is WSMsgType.CLOSE: + self._set_closing(msg.data) + # Could be closed while awaiting reader. + if not self._closed and self._autoclose: + # The client is likely going to close the + # connection out from under us so we do not + # want to drain any pending writes as it will + # likely result writing to a broken pipe. + await self.close(drain=False) + elif msg.type is WSMsgType.CLOSING: + self._set_closing(WSCloseCode.OK) + elif msg.type is WSMsgType.PING and self._autoping: + await self.pong(msg.data) + continue + elif msg.type is WSMsgType.PONG and self._autoping: + continue + + return msg + + async def receive_str(self, *, timeout: Optional[float] = None) -> str: + msg = await self.receive(timeout) + if msg.type is not WSMsgType.TEXT: + raise WSMessageTypeError( + f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT" + ) + return cast(str, msg.data) + + async def receive_bytes(self, *, timeout: Optional[float] = None) -> bytes: + msg = await self.receive(timeout) + if msg.type is not WSMsgType.BINARY: + raise WSMessageTypeError( + f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY" + ) + return cast(bytes, msg.data) + + async def receive_json( + self, *, loads: JSONDecoder = json.loads, timeout: Optional[float] = None + ) -> Any: + data = await self.receive_str(timeout=timeout) + return loads(data) + + async def write(self, data: bytes) -> None: + raise RuntimeError("Cannot call .write() for websocket") + + def __aiter__(self) -> "WebSocketResponse": + return self + + async def __anext__(self) -> WSMessage: + msg = await self.receive() + if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): + raise StopAsyncIteration + return msg + + def _cancel(self, exc: BaseException) -> None: + # web_protocol calls this from connection_lost + # or when the server is shutting down. + self._closing = True + self._cancel_heartbeat() + if self._reader is not None: + set_exception(self._reader, exc) diff --git a/lib/python3.12/site-packages/aiohttp/worker.py b/lib/python3.12/site-packages/aiohttp/worker.py new file mode 100644 index 0000000000000000000000000000000000000000..f7281bfde7541412c3174aa5fdcb859fa1b7a996 --- /dev/null +++ b/lib/python3.12/site-packages/aiohttp/worker.py @@ -0,0 +1,255 @@ +"""Async gunicorn worker for aiohttp.web""" + +import asyncio +import inspect +import os +import re +import signal +import sys +from types import FrameType +from typing import TYPE_CHECKING, Any, Optional + +from gunicorn.config import AccessLogFormat as GunicornAccessLogFormat +from gunicorn.workers import base + +from aiohttp import web + +from .helpers import set_result +from .web_app import Application +from .web_log import AccessLogger + +if TYPE_CHECKING: + import ssl + + SSLContext = ssl.SSLContext +else: + try: + import ssl + + SSLContext = ssl.SSLContext + except ImportError: # pragma: no cover + ssl = None # type: ignore[assignment] + SSLContext = object # type: ignore[misc,assignment] + + +__all__ = ("GunicornWebWorker", "GunicornUVLoopWebWorker") + + +class GunicornWebWorker(base.Worker): # type: ignore[misc,no-any-unimported] + + DEFAULT_AIOHTTP_LOG_FORMAT = AccessLogger.LOG_FORMAT + DEFAULT_GUNICORN_LOG_FORMAT = GunicornAccessLogFormat.default + + def __init__(self, *args: Any, **kw: Any) -> None: # pragma: no cover + super().__init__(*args, **kw) + + self._task: Optional[asyncio.Task[None]] = None + self.exit_code = 0 + self._notify_waiter: Optional[asyncio.Future[bool]] = None + + def init_process(self) -> None: + # create new event_loop after fork + asyncio.get_event_loop().close() + + self.loop = asyncio.new_event_loop() + asyncio.set_event_loop(self.loop) + + super().init_process() + + def run(self) -> None: + self._task = self.loop.create_task(self._run()) + + try: # ignore all finalization problems + self.loop.run_until_complete(self._task) + except Exception: + self.log.exception("Exception in gunicorn worker") + self.loop.run_until_complete(self.loop.shutdown_asyncgens()) + self.loop.close() + + sys.exit(self.exit_code) + + async def _run(self) -> None: + runner = None + if isinstance(self.wsgi, Application): + app = self.wsgi + elif inspect.iscoroutinefunction(self.wsgi) or ( + sys.version_info < (3, 14) and asyncio.iscoroutinefunction(self.wsgi) + ): + wsgi = await self.wsgi() + if isinstance(wsgi, web.AppRunner): + runner = wsgi + app = runner.app + else: + app = wsgi + else: + raise RuntimeError( + "wsgi app should be either Application or " + "async function returning Application, got {}".format(self.wsgi) + ) + + if runner is None: + access_log = self.log.access_log if self.cfg.accesslog else None + runner = web.AppRunner( + app, + logger=self.log, + keepalive_timeout=self.cfg.keepalive, + access_log=access_log, + access_log_format=self._get_valid_log_format( + self.cfg.access_log_format + ), + shutdown_timeout=self.cfg.graceful_timeout / 100 * 95, + ) + await runner.setup() + + ctx = self._create_ssl_context(self.cfg) if self.cfg.is_ssl else None + + runner = runner + assert runner is not None + server = runner.server + assert server is not None + for sock in self.sockets: + site = web.SockSite( + runner, + sock, + ssl_context=ctx, + ) + await site.start() + + # If our parent changed then we shut down. + pid = os.getpid() + try: + while self.alive: # type: ignore[has-type] + self.notify() + + cnt = server.requests_count + if self.max_requests and cnt > self.max_requests: + self.alive = False + self.log.info("Max requests, shutting down: %s", self) + + elif pid == os.getpid() and self.ppid != os.getppid(): + self.alive = False + self.log.info("Parent changed, shutting down: %s", self) + else: + await self._wait_next_notify() + except BaseException: + pass + + await runner.cleanup() + + def _wait_next_notify(self) -> "asyncio.Future[bool]": + self._notify_waiter_done() + + loop = self.loop + assert loop is not None + self._notify_waiter = waiter = loop.create_future() + self.loop.call_later(1.0, self._notify_waiter_done, waiter) + + return waiter + + def _notify_waiter_done( + self, waiter: Optional["asyncio.Future[bool]"] = None + ) -> None: + if waiter is None: + waiter = self._notify_waiter + if waiter is not None: + set_result(waiter, True) + + if waiter is self._notify_waiter: + self._notify_waiter = None + + def init_signals(self) -> None: + # Set up signals through the event loop API. + + self.loop.add_signal_handler( + signal.SIGQUIT, self.handle_quit, signal.SIGQUIT, None + ) + + self.loop.add_signal_handler( + signal.SIGTERM, self.handle_exit, signal.SIGTERM, None + ) + + self.loop.add_signal_handler( + signal.SIGINT, self.handle_quit, signal.SIGINT, None + ) + + self.loop.add_signal_handler( + signal.SIGWINCH, self.handle_winch, signal.SIGWINCH, None + ) + + self.loop.add_signal_handler( + signal.SIGUSR1, self.handle_usr1, signal.SIGUSR1, None + ) + + self.loop.add_signal_handler( + signal.SIGABRT, self.handle_abort, signal.SIGABRT, None + ) + + # Don't let SIGTERM and SIGUSR1 disturb active requests + # by interrupting system calls + signal.siginterrupt(signal.SIGTERM, False) + signal.siginterrupt(signal.SIGUSR1, False) + # Reset signals so Gunicorn doesn't swallow subprocess return codes + # See: https://github.com/aio-libs/aiohttp/issues/6130 + + def handle_quit(self, sig: int, frame: Optional[FrameType]) -> None: + self.alive = False + + # worker_int callback + self.cfg.worker_int(self) + + # wakeup closing process + self._notify_waiter_done() + + def handle_abort(self, sig: int, frame: Optional[FrameType]) -> None: + self.alive = False + self.exit_code = 1 + self.cfg.worker_abort(self) + sys.exit(1) + + @staticmethod + def _create_ssl_context(cfg: Any) -> "SSLContext": + """Creates SSLContext instance for usage in asyncio.create_server. + + See ssl.SSLSocket.__init__ for more details. + """ + if ssl is None: # pragma: no cover + raise RuntimeError("SSL is not supported.") + + ctx = ssl.SSLContext(cfg.ssl_version) + ctx.load_cert_chain(cfg.certfile, cfg.keyfile) + ctx.verify_mode = cfg.cert_reqs + if cfg.ca_certs: + ctx.load_verify_locations(cfg.ca_certs) + if cfg.ciphers: + ctx.set_ciphers(cfg.ciphers) + return ctx + + def _get_valid_log_format(self, source_format: str) -> str: + if source_format == self.DEFAULT_GUNICORN_LOG_FORMAT: + return self.DEFAULT_AIOHTTP_LOG_FORMAT + elif re.search(r"%\([^\)]+\)", source_format): + raise ValueError( + "Gunicorn's style options in form of `%(name)s` are not " + "supported for the log formatting. Please use aiohttp's " + "format specification to configure access log formatting: " + "http://docs.aiohttp.org/en/stable/logging.html" + "#format-specification" + ) + else: + return source_format + + +class GunicornUVLoopWebWorker(GunicornWebWorker): + def init_process(self) -> None: + import uvloop + + # Close any existing event loop before setting a + # new policy. + asyncio.get_event_loop().close() + + # Setup uvloop policy, so that every + # asyncio.get_event_loop() will create an instance + # of uvloop event loop. + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + super().init_process() diff --git a/lib/python3.12/site-packages/flash_attn/__init__.py b/lib/python3.12/site-packages/flash_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db131242dd4925eac3c41fb3c7623d1c4a765b03 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/__init__.py @@ -0,0 +1,11 @@ +__version__ = "2.7.4.post1" + +from flash_attn.flash_attn_interface import ( + flash_attn_func, + flash_attn_kvpacked_func, + flash_attn_qkvpacked_func, + flash_attn_varlen_func, + flash_attn_varlen_kvpacked_func, + flash_attn_varlen_qkvpacked_func, + flash_attn_with_kvcache, +) diff --git a/lib/python3.12/site-packages/flash_attn/bert_padding.py b/lib/python3.12/site-packages/flash_attn/bert_padding.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2d35159a014a9d03aabead9e52e009168696ea --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/bert_padding.py @@ -0,0 +1,218 @@ +# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + + +class IndexFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, input, indices): + ctx.save_for_backward(indices) + assert input.ndim >= 2 + ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] + second_dim = other_shape.numel() + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + # return input[indices] + return torch.gather( + rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) + ).reshape(-1, *other_shape) + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + assert grad_output.ndim >= 2 + other_shape = grad_output.shape[1:] + grad_output = rearrange(grad_output, "b ... -> b (...)") + grad_input = torch.zeros( + [ctx.first_axis_dim, grad_output.shape[1]], + device=grad_output.device, + dtype=grad_output.dtype, + ) + # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + # grad_input[indices] = grad_output + grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) + return grad_input.reshape(ctx.first_axis_dim, *other_shape), None + + +index_first_axis = IndexFirstAxis.apply + + +class IndexPutFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, values, indices, first_axis_dim): + ctx.save_for_backward(indices) + assert indices.ndim == 1 + assert values.ndim >= 2 + output = torch.zeros( + first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype + ) + # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + output[indices] = values + # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) + return output + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + grad_values = grad_output[indices] + # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) + return grad_values, None, None + + +index_put_first_axis = IndexPutFirstAxis.apply + + +class IndexFirstAxisResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, input, indices): + ctx.save_for_backward(indices) + assert input.ndim >= 2 + ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] + second_dim = other_shape.numel() + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + output = input[indices] + # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last + # memory format to channel_first. In other words, input might not be contiguous. + # If we don't detach, Pytorch complains about output being a view and is being modified inplace + return output, input.detach() + + @staticmethod + def backward(ctx, grad_output, grad_residual): + (indices,) = ctx.saved_tensors + assert grad_output.ndim >= 2 + other_shape = grad_output.shape[1:] + assert grad_residual.shape[1:] == other_shape + grad_input = grad_residual + # grad_input[indices] += grad_output + indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1))) + indices = indices.expand_as(grad_output) + grad_input.scatter_add_(0, indices, grad_output) + return grad_input.reshape(ctx.first_axis_dim, *other_shape), None + + +index_first_axis_residual = IndexFirstAxisResidual.apply + + +def unpad_input(hidden_states, attention_mask, unused_mask=None): + """ + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. + unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. + indices: (total_nnz), the indices of masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. + """ + all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask + seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) + used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the + # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim + # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to + # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, + # so we write custom forward and backward to make it a bit faster. + return ( + index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + used_seqlens_in_batch, + ) + + +def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length): + """ + Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model). + The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286). + + For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is: + ``` + [ + [2, 3, 0, 0, 0, 0], + [3, 2, 0, 0, 0, 0], + [6, 0, 0, 0, 0, 0] + ] + ``` + , which refers to the 3D-attention mask: + ``` + [ + [ + [1, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [0, 0, 1, 0, 0, 0], + [0, 0, 1, 1, 0, 0], + [0, 0, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 1] + ], + [ + [1, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [0, 0, 0, 1, 0, 0], + [0, 0, 0, 1, 1, 0], + [0, 0, 0, 0, 0, 1] + ], + [ + [1, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 0, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 1] + ] + ] + ```. + + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + """ + length = attention_mask_in_length.sum(dim=-1) + seqlen = attention_mask_in_length.size(-1) + attention_mask_2d = torch.arange(seqlen, device=length.device, dtype=length.dtype).expand(len(length), seqlen) < length.unsqueeze(1) + real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten() + seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx] + indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the + # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim + # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to + # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, + # so we write custom forward and backward to make it a bit faster. + return ( + index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def pad_input(hidden_states, indices, batch, seqlen): + """ + Arguments: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. + batch: int, batch size for the padded sequence. + seqlen: int, maximum sequence length for the padded sequence. + Return: + hidden_states: (batch, seqlen, ...) + """ + dim = hidden_states.shape[-1] + # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) + # output[indices] = hidden_states + output = index_put_first_axis(hidden_states, indices, batch * seqlen) + return rearrange(output, "(b s) ... -> b s ...", b=batch) diff --git a/lib/python3.12/site-packages/flash_attn/flash_attn_interface.py b/lib/python3.12/site-packages/flash_attn/flash_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..30134990d688fd0df0e4c374ed0eb2d93f9a718f --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/flash_attn_interface.py @@ -0,0 +1,1611 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +import os + +# isort: off +# We need to import the CUDA kernels after importing torch +USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE" +if USE_TRITON_ROCM: + from .flash_attn_triton_amd import interface_fa as flash_attn_gpu +else: + import flash_attn_2_cuda as flash_attn_gpu + +# isort: on + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _get_block_size_n(device, head_dim, is_dropout, is_causal): + # This should match the block sizes in the CUDA kernel + assert head_dim <= 256 + major, minor = torch.cuda.get_device_capability(device) + is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100) + is_sm80 = major == 8 and minor == 0 + is_sm90 = major == 9 and minor == 0 + if head_dim <= 32: + return 128 + if head_dim <= 64: + return 128 if not is_dropout else 64 + elif head_dim <= 96: + return 64 + elif head_dim <= 128: + if is_sm8x: + return 64 if (not is_dropout and is_causal) else 32 + else: + return 64 if not is_dropout else 32 + elif head_dim <= 160: + if is_sm8x: + return 64 + else: + return 32 + elif head_dim <= 192: + return 64 + elif head_dim <= 224: + return 64 + elif head_dim <= 256: + return 64 + + +def round_multiple(x, m): + return (x + m - 1) // m * m + + +# torch.compile() support is only enabled for pytorch >= 2.4 +# The reason for this is that we are using the new custom_op and register_fake +# APIs, which support inplace modification of inputs in the function itself +if torch.__version__ >= "2.4.0": + _torch_custom_op_wrapper = torch.library.custom_op + _torch_register_fake_wrapper = torch.library.register_fake +else: + def noop_custom_op_wrapper(name, fn=None, /, *, mutates_args, device_types=None, schema=None): + def wrap(func): + return func + if fn is None: + return wrap + return fn + def noop_register_fake_wrapper(op, fn=None, /, *, lib=None, _stacklevel=1): + def wrap(func): + return func + if fn is None: + return wrap + return fn + _torch_custom_op_wrapper = noop_custom_op_wrapper + _torch_register_fake_wrapper = noop_register_fake_wrapper + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_forward", mutates_args=(), device_types="cuda") +def _flash_attn_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + return_softmax: bool +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.fwd( + q, + k, + v, + None, + alibi_slopes, + dropout_p, + softmax_scale, + causal, + window_size_left, + window_size_right, + softcap, + return_softmax, + None, + ) + return out, softmax_lse, S_dmask, rng_state + + +@_torch_register_fake_wrapper("flash_attn::_flash_attn_forward") +def _flash_attn_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + return_softmax: bool +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + batch_size, seqlen_q, num_heads, head_size = q.shape + seqlen_k = k.shape[1] + out = torch.empty_like(q) + softmax_lse = torch.empty((batch_size, num_heads, seqlen_q), dtype=torch.float32, device=q.device, layout=q.layout) + p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) + if return_softmax: + p = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128), round_multiple(seqlen_k, 128)), dtype=q.dtype, device=q.device, layout=q.layout) + rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) + + return out, softmax_lse, p, rng_state + + +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward +else: + _wrapped_flash_attn_forward = _flash_attn_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_forward", mutates_args=(), device_types="cuda") +def _flash_attn_varlen_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + alibi_slopes: Optional[torch.Tensor] = None, + return_softmax: bool = False, + block_table: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + zero_tensors: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.varlen_fwd( + q, + k, + v, + None, + cu_seqlens_q, + cu_seqlens_k, + seqused_k, + leftpad_k, + block_table, + alibi_slopes, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + zero_tensors, + causal, + window_size_left, + window_size_right, + softcap, + return_softmax, + None, + ) + # if out.isnan().any() or softmax_lse.isnan().any(): + # breakpoint() + return out, softmax_lse, S_dmask, rng_state + + +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_forward") +def _flash_attn_varlen_forward_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int = -1, + window_size_right: int = -1, + softcap: float = 0.0, + alibi_slopes: Optional[torch.Tensor] = None, + return_softmax: bool = False, + block_table: Optional[torch.Tensor] = None, + leftpad_k: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + zero_tensors: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + paged_kv = block_table is not None + batch_size = cu_seqlens_q.numel() - 1 + total_q, num_heads, _ = q.shape + + out = torch.empty_like(q) + softmax_lse = torch.empty((num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout) + p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout) + seqlen_q_rounded = round_multiple(max_seqlen_q, 128) + seqlen_k_rounded = round_multiple(max_seqlen_k, 128) + if return_softmax: + p = torch.empty((batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded), dtype=q.dtype, device=q.device, layout=q.layout) + rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) + return out, softmax_lse, p, rng_state + + +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_forward = torch.ops.flash_attn._flash_attn_varlen_forward +else: + _wrapped_flash_attn_varlen_forward = _flash_attn_varlen_forward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_backward", mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + dq: Optional[torch.Tensor], + dk: Optional[torch.Tensor], + dv: Optional[torch.Tensor], + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + deterministic: bool, + rng_state: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + ( + dq, + dk, + dv, + softmax_d, + ) = flash_attn_gpu.bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + alibi_slopes, + dropout_p, + softmax_scale, + causal, + window_size_left, + window_size_right, + softcap, + deterministic, + None, + rng_state, + ) + return softmax_d + + +@_torch_register_fake_wrapper("flash_attn::_flash_attn_backward") +def _flash_attn_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + dq: Optional[torch.Tensor], + dk: Optional[torch.Tensor], + dv: Optional[torch.Tensor], + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + deterministic: bool, + rng_state: Optional[torch.Tensor] = None, +) -> torch.Tensor: + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + if dq is None: + dq = torch.empty_like(q) + if dk is None: + dk = torch.empty_like(k) + if dv is None: + dv = torch.empty_like(v) + batch_size, seqlen_q, num_heads, _ = q.shape + softmax_d = torch.empty((batch_size, num_heads, round_multiple(seqlen_q, 128)), device=q.device, dtype=torch.float32) + + return softmax_d + + +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward +else: + _wrapped_flash_attn_backward = _flash_attn_backward + + +@_torch_custom_op_wrapper("flash_attn::_flash_attn_varlen_backward", mutates_args=("dq", "dk", "dv"), device_types="cuda") +def _flash_attn_varlen_backward( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + dq: Optional[torch.Tensor], + dk: Optional[torch.Tensor], + dv: Optional[torch.Tensor], + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + deterministic: bool, + rng_state: Optional[torch.Tensor] = None, + zero_tensors: bool = False, +) -> torch.Tensor: + # dq, dk, dv are allocated by us so they should already be contiguous + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + ( + dq, + dk, + dv, + softmax_d, + ) = flash_attn_gpu.varlen_bwd( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + alibi_slopes, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + zero_tensors, + causal, + window_size_left, + window_size_right, + softcap, + deterministic, + None, + rng_state, + ) + # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): + # breakpoint() + return softmax_d + + +@_torch_register_fake_wrapper("flash_attn::_flash_attn_varlen_backward") +def _flash_attn_varlen_backward_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + softmax_lse: torch.Tensor, + dq: Optional[torch.Tensor], + dk: Optional[torch.Tensor], + dv: Optional[torch.Tensor], + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float, + softmax_scale: float, + causal: bool, + window_size_left: int, + window_size_right: int, + softcap: float, + alibi_slopes: Optional[torch.Tensor], + deterministic: bool, + rng_state: Optional[torch.Tensor] = None, + zero_tensors: bool = False, +) -> torch.Tensor: + dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)] + batch_size = cu_seqlens_q.numel() - 1 + total_q, num_heads, _ = q.shape + + if dq is None: + dq = torch.empty_like(q) + if dk is None: + dk = torch.empty_like(k) + if dv is None: + dv = torch.empty_like(v) + softmax_d = torch.empty((num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32) + + return softmax_d + + +if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_varlen_backward = torch.ops.flash_attn._flash_attn_varlen_backward +else: + _wrapped_flash_attn_varlen_backward = _flash_attn_varlen_backward + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + is_grad_enabled, + ): + is_grad = is_grad_enabled and qkv.requires_grad + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach() + head_size_og = q.size(3) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( + q, + k, + v, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + ) + if is_grad: + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) + ctx.dropout_p = dropout_p + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + head_size_og = dout.size(3) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dqkv[:, :, 0], + dqkv[:, :, 1], + dqkv[:, :, 2], + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + qkv, + cu_seqlens, + max_seqlen, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + is_grad_enabled, + ): + is_grad = is_grad_enabled and qkv.requires_grad + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach() + head_size_og = q.size(2) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( + q, + k, + v, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + block_table=None, + ) + if is_grad: + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state) + ctx.dropout_p = dropout_p + ctx.max_seqlen = max_seqlen + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors + qkv_shape = q.shape[:-2] + (3, *q.shape[-2:]) + dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device) + head_size_og = dout.size(2) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_varlen_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dqkv[:, 0], + dqkv[:, 1], + dqkv[:, 2], + cu_seqlens, + cu_seqlens, + ctx.max_seqlen, + ctx.max_seqlen, + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dqkv = dqkv[..., : dout.shape[-1]] # We could have padded the head dimension + return dqkv, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + kv, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + is_grad_enabled, + ): + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach() + head_size_og = q.size(3) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( + q, + k, + v, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + ) + if is_grad: + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) + ctx.dropout_p = dropout_p + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors + dq = torch.empty_like(q) + kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) + dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) + head_size_og = dout.size(3) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dq, + dkv[:, :, 0], + dkv[:, :, 1], + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension + dkv = dkv[..., : dout.shape[-1]] + return dq, dkv, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + kv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + is_grad_enabled, + ): + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, kv] + ) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + k, v = kv[:, 0].detach(), kv[:, 1].detach() + head_size_og = q.size(2) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + block_table=None, + ) + if is_grad: + ctx.save_for_backward( + q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state + ) + ctx.dropout_p = dropout_p + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors + dq = torch.empty_like(q) + kv_shape = k.shape[:-2] + (2, *k.shape[-2:]) + dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device) + head_size_og = dout.size(2) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_varlen_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dq, + dkv[:, 0], + dkv[:, 1], + cu_seqlens_q, + cu_seqlens_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension + dkv = dkv[..., : dout.shape[-1]] + return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class FlashAttnFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + is_grad_enabled, + ): + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + head_size_og = q.size(3) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward( + q, + k, + v, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + ) + if is_grad: + ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state) + ctx.dropout_p = dropout_p + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + head_size_og = dout.size(3) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension + dk = dk[..., : dout.shape[-1]] + dv = dv[..., : dout.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None + + +class FlashAttnVarlenFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_softmax, + block_table, + is_grad_enabled, + ): + is_grad = is_grad_enabled and any( + x.requires_grad for x in [q, k, v] + ) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + head_size_og = q.size(2) + if head_size_og % 8 != 0: + q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8]) + k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8]) + v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8]) + out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=return_softmax and dropout_p > 0, + block_table=block_table, + ) + if is_grad: + ctx.save_for_backward( + q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state + ) + ctx.dropout_p = dropout_p + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + + out = out_padded[..., :head_size_og] + return out if not return_softmax else (out, softmax_lse, S_dmask) + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors + dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + head_size_og = dout.size(2) + dout_padded = dout + if head_size_og % 8 != 0: + dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8]) + _wrapped_flash_attn_varlen_backward( + dout_padded, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + ctx.max_seqlen_q, + ctx.max_seqlen_k, + ctx.dropout_p, + ctx.softmax_scale, + ctx.causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state=rng_state, + ) + dq = dq[..., : dout.shape[-1]] # We could have padded the head dimension + dk = dk[..., : dout.shape[-1]] + dv = dv[..., : dout.shape[-1]] + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +def flash_attn_qkvpacked_func( + qkv, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # <=0.0 means deactivate + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_kvpacked_func and flash_attn_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (batch_size, seqlen, 3, nheads, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to + the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnQKVPackedFunc.apply( + qkv, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + torch.is_grad_enabled(), + ) + + +def flash_attn_kvpacked_func( + q, + kv, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of K, V. + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + kv: (batch_size, seqlen, 2, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnKVPackedFunc.apply( + q, + kv, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + torch.is_grad_enabled(), + ) + + +def flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k: (batch_size, seqlen, nheads_k, headdim) + v: (batch_size, seqlen, nheads_k, headdim) + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnFunc.apply( + q, + k, + v, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + torch.is_grad_enabled(), + ) + + +def flash_attn_varlen_qkvpacked_func( + qkv, + cu_seqlens, + max_seqlen, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If Q, K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of Q, K, V. + For multi-query and grouped-query attention (MQA/GQA), please see + flash_attn_varlen_kvpacked_func and flash_attn_varlen_func. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive. + + Arguments: + qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch. + cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into qkv. + max_seqlen: int. Maximum sequence length in the batch. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (total, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnVarlenQKVPackedFunc.apply( + qkv, + cu_seqlens, + max_seqlen, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + torch.is_grad_enabled(), + ) + + +def flash_attn_varlen_kvpacked_func( + q, + kv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, +): + """dropout_p should be set to 0.0 during evaluation + If K, V are already stacked into 1 tensor, this function will be faster than + calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation + of the gradients of K, V. + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. + kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch. + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into q. + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into kv. + max_seqlen_q: int. Maximum query sequence length in the batch. + max_seqlen_k: int. Maximum key sequence length in the batch. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (total, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnVarlenKVPackedFunc.apply( + q, + kv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + torch.is_grad_enabled(), + ) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, + block_table=None, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. + k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. + v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into q. + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into kv. + max_seqlen_q: int. Maximum query sequence length in the batch. + max_seqlen_k: int. Maximum key sequence length in the batch. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (total, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + softcap, + alibi_slopes, + deterministic, + return_attn_probs, + block_table, + torch.is_grad_enabled(), + ) + + +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + block_table: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + alibi_slopes=None, + num_splits=0, + return_softmax_lse=False, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table, + or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k. + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0. + block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" + assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) + if cache_seqlens is not None and isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + cache_seqlens = maybe_contiguous(cache_seqlens) + cache_batch_idx = maybe_contiguous(cache_batch_idx) + block_table = maybe_contiguous(block_table) + out, softmax_lse = flash_attn_gpu.fwd_kvcache( + q, + k_cache, + v_cache, + k, + v, + cache_seqlens, + rotary_cos, + rotary_sin, + cache_batch_idx, + cache_leftpad, + block_table, + alibi_slopes, + None, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + rotary_interleaved, + num_splits, + ) + return (out, softmax_lse) if return_softmax_lse else out diff --git a/lib/python3.12/site-packages/flash_attn/flash_attn_triton.py b/lib/python3.12/site-packages/flash_attn/flash_attn_triton.py new file mode 100644 index 0000000000000000000000000000000000000000..30420c057adf1916e16403d3f0d02d0e26c8b7a3 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/flash_attn_triton.py @@ -0,0 +1,1160 @@ +""" +*Experimental* implementation of FlashAttention in Triton. +Tested with triton==2.0.0.dev20221202. +Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions +other than 64: +https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207 +We'll update this implementation with the new Triton backend once this is fixed. + +We use the FlashAttention implementation from Phil Tillet a starting point. +https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py + +Changes: +- Implement both causal and non-causal attention. +- Implement both self-attention and cross-attention. +- Support arbitrary seqlens (not just multiples of 128), for both forward and backward. +- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward. +- Support attention bias. +- Speed up the forward pass a bit, and only store the LSE instead of m and l. +- Make the backward for d=128 much faster by reducing register spilling. +- Optionally parallelize the backward pass across seqlen_k, to deal with the case of +small batch size * nheads. + +Caution: +- This is an *experimental* implementation. The forward pass should be quite robust but +I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler). +- This implementation has only been tested on A100. +- If you plan to use headdim other than 64 and 128, you should test for race conditions +(due to the Triton compiler), as done in tests/test_flash_attn.py +"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions +for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident +that there are none left for other head dimensions. + +Differences between this Triton version and the CUDA version: +- Triton version doesn't support dropout. +- Triton forward is generally faster than CUDA forward, while Triton backward is +generally slower than CUDA backward. Overall Triton forward + backward is slightly slower +than CUDA forward + backward. +- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor). +- Triton version supports attention bias, while CUDA version doesn't. +""" + +import math + +import torch +import triton +import triton.language as tl + + +# Disabling autotune for now, set num_warps=4 if headdim=64 and num_warps=8 if headdim=128 +# @triton.autotune( +# configs=[ +# triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1), +# # This config has a race condition when EVEN_M == False, disabling it for now. +# # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1), +# ], +# key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'] +# ) +@triton.heuristics( + { + "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0, + "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0, + "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"], + } +) +@triton.jit +def _fwd_kernel( + Q, + K, + V, + Bias, + Out, + Lse, + TMP, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug + softmax_scale, + stride_qb, + stride_qh, + stride_qm, + stride_kb, + stride_kh, + stride_kn, + stride_vb, + stride_vh, + stride_vn, + stride_bb, + stride_bh, + stride_bm, + stride_ob, + stride_oh, + stride_om, + nheads, + seqlen_q, + seqlen_k, + seqlen_q_rounded, + headdim, + CACHE_KEY_SEQLEN_Q, + CACHE_KEY_SEQLEN_K, + BIAS_TYPE: tl.constexpr, + IS_CAUSAL: tl.constexpr, + BLOCK_HEADDIM: tl.constexpr, + EVEN_M: tl.constexpr, + EVEN_N: tl.constexpr, + EVEN_HEADDIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + start_m = tl.program_id(0) + off_hb = tl.program_id(1) + off_b = off_hb // nheads + off_h = off_hb % nheads + # off_b = tl.program_id(1) + # off_h = tl.program_id(2) + # off_hb = off_b * nheads + off_h + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_HEADDIM) + # Initialize pointers to Q, K, V + # Adding parenthesis around indexing might use int32 math instead of int64 math? + # https://github.com/openai/triton/issues/741 + # I'm seeing a tiny bit of difference (5-7us) + q_ptrs = ( + Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :]) + ) + k_ptrs = ( + K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :]) + ) + v_ptrs = ( + V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :]) + ) + if BIAS_TYPE == "vector": + b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n + elif BIAS_TYPE == "matrix": + b_ptrs = ( + Bias + + off_b * stride_bb + + off_h * stride_bh + + (offs_m[:, None] * stride_bm + offs_n[None, :]) + ) + # initialize pointer to m and l + t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m + lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32) + # load q: it will stay in SRAM throughout + # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call + # tl.load(q_ptrs), we get the wrong output! + if EVEN_M & EVEN_N: + if EVEN_HEADDIM: + q = tl.load(q_ptrs) + else: + q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0) + else: + if EVEN_HEADDIM: + q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0) + else: + q = tl.load( + q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0 + ) + # loop over k, v and update accumulator + end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k) + for start_n in range(0, end_n, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + if EVEN_N & EVEN_M: # If we just do "if EVEN_N", there seems to be some race condition + if EVEN_HEADDIM: + k = tl.load(k_ptrs + start_n * stride_kn) + else: + k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0) + else: + if EVEN_HEADDIM: + k = tl.load( + k_ptrs + start_n * stride_kn, + mask=(start_n + offs_n)[:, None] < seqlen_k, + other=0.0, + ) + else: + k = tl.load( + k_ptrs + start_n * stride_kn, + mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), + other=0.0, + ) + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k, trans_b=True) + # Trying to combine the two masks seem to make the result wrong + if not EVEN_N: # Need to mask out otherwise the softmax is wrong + qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float("-inf")) + if IS_CAUSAL: + qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float("-inf")) + if BIAS_TYPE != "none": + if BIAS_TYPE == "vector": + if EVEN_N: + bias = tl.load(b_ptrs + start_n).to(tl.float32) + else: + bias = tl.load( + b_ptrs + start_n, mask=(start_n + offs_n) < seqlen_k, other=0.0 + ).to(tl.float32) + bias = bias[None, :] + elif BIAS_TYPE == "matrix": + if EVEN_M & EVEN_N: + bias = tl.load(b_ptrs + start_n).to(tl.float32) + else: + bias = tl.load( + b_ptrs + start_n, + mask=(offs_m[:, None] < seqlen_q) + & ((start_n + offs_n)[None, :] < seqlen_k), + other=0.0, + ).to(tl.float32) + # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler + # can then fuse the mult and add into an fma instruction. But if we have bias we need to + # to multiply with softmax_scale here. + qk = qk * softmax_scale + bias + m_ij = tl.maximum(tl.max(qk, 1), lse_i) + p = tl.exp(qk - m_ij[:, None]) + else: + m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i) + p = tl.exp(qk * softmax_scale - m_ij[:, None]) + l_ij = tl.sum(p, 1) + + # scale acc_o + acc_o_scale = tl.exp(m_i - m_ij) + + # # -- update output accumulator -- + # BUG: have to store and immediately load + tl.store(t_ptrs, acc_o_scale) + acc_o_scale = tl.load(t_ptrs) + acc_o = acc_o * acc_o_scale[:, None] + # update acc_o + if EVEN_N & EVEN_M: # If we just do "if EVEN_N", there seems to be some race condition + if EVEN_HEADDIM: + v = tl.load(v_ptrs + start_n * stride_vn) + else: + v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0) + else: + if EVEN_HEADDIM: + v = tl.load( + v_ptrs + start_n * stride_vn, + mask=(start_n + offs_n)[:, None] < seqlen_k, + other=0.0, + ) + else: + v = tl.load( + v_ptrs + start_n * stride_vn, + mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), + other=0.0, + ) + p = p.to(v.dtype) + acc_o += tl.dot(p, v) + + # -- update statistics + m_i = m_ij + l_i_new = tl.exp(lse_i - m_ij) + l_ij + lse_i = m_ij + tl.log(l_i_new) + + o_scale = tl.exp(m_i - lse_i) + # BUG: have to store and immediately load + tl.store(t_ptrs, o_scale) + o_scale = tl.load(t_ptrs) + acc_o = acc_o * o_scale[:, None] + # rematerialize offsets to save registers + start_m = tl.program_id(0) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + # write back l and m + lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m + tl.store(lse_ptrs, lse_i) + # initialize pointers to output + offs_d = tl.arange(0, BLOCK_HEADDIM) + out_ptrs = ( + Out + + off_b * stride_ob + + off_h * stride_oh + + (offs_m[:, None] * stride_om + offs_d[None, :]) + ) + if EVEN_M: + if EVEN_HEADDIM: + tl.store(out_ptrs, acc_o) + else: + tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim) + else: + if EVEN_HEADDIM: + tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q) + else: + tl.store( + out_ptrs, acc_o, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim) + ) + + +@triton.jit +def _bwd_preprocess_do_o_dot( + Out, + DO, + Delta, + stride_ob, + stride_oh, + stride_om, + stride_dob, + stride_doh, + stride_dom, + nheads, + seqlen_q, + seqlen_q_rounded, + headdim, + BLOCK_M: tl.constexpr, + BLOCK_HEADDIM: tl.constexpr, +): + start_m = tl.program_id(0) + off_hb = tl.program_id(1) + off_b = off_hb // nheads + off_h = off_hb % nheads + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_HEADDIM) + # load + o = tl.load( + Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :], + mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + other=0.0, + ).to(tl.float32) + do = tl.load( + DO + + off_b * stride_dob + + off_h * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :], + mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + other=0.0, + ).to(tl.float32) + delta = tl.sum(o * do, axis=1) + # write-back + tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta) + + +@triton.jit +def _bwd_store_dk_dv( + dk_ptrs, + dv_ptrs, + dk, + dv, + offs_n, + offs_d, + seqlen_k, + headdim, + EVEN_M: tl.constexpr, + EVEN_N: tl.constexpr, + EVEN_HEADDIM: tl.constexpr, +): + # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False, + # if we just call tl.store(dv_ptrs), there's a race condition + if EVEN_N & EVEN_M: + if EVEN_HEADDIM: + tl.store(dv_ptrs, dv) + tl.store(dk_ptrs, dk) + else: + tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim) + tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim) + else: + if EVEN_HEADDIM: + tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k) + tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k) + else: + tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)) + tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)) + + +@triton.jit +def _bwd_kernel_one_col_block( + start_n, + Q, + K, + V, + Bias, + DO, + DQ, + DK, + DV, + LSE, + D, + softmax_scale, + stride_qm, + stride_kn, + stride_vn, + stride_bm, + stride_dom, + stride_dqm, + stride_dkn, + stride_dvn, + seqlen_q, + seqlen_k, + headdim, + ATOMIC_ADD: tl.constexpr, + BIAS_TYPE: tl.constexpr, + IS_CAUSAL: tl.constexpr, + BLOCK_HEADDIM: tl.constexpr, + EVEN_M: tl.constexpr, + EVEN_N: tl.constexpr, + EVEN_HEADDIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N) + begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M + # initialize row/col offsets + offs_qm = begin_m + tl.arange(0, BLOCK_M) + offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_HEADDIM) + # initialize pointers to value-like data + q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :]) + k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :]) + v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :]) + do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :]) + dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :]) + if BIAS_TYPE == "vector": + b_ptrs = Bias + offs_n + elif BIAS_TYPE == "matrix": + b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :]) + # initialize dv and dk + dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32) + dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32) + # There seems to be some problem with Triton pipelining that makes results wrong for + # headdim=64, seqlen=(113, 255), bias_type='matrix'. In this case the for loop + # may have zero step, and pipelining with the bias matrix could screw it up. + # So we just exit early. + if begin_m >= seqlen_q: + dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :]) + dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :]) + _bwd_store_dk_dv( + dk_ptrs, + dv_ptrs, + dk, + dv, + offs_n, + offs_d, + seqlen_k, + headdim, + EVEN_M=EVEN_M, + EVEN_N=EVEN_N, + EVEN_HEADDIM=EVEN_HEADDIM, + ) + return + # k and v stay in SRAM throughout + # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False, + # if we just call tl.load(k_ptrs), we get the wrong output! + if EVEN_N & EVEN_M: + if EVEN_HEADDIM: + k = tl.load(k_ptrs) + v = tl.load(v_ptrs) + else: + k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0) + v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0) + else: + if EVEN_HEADDIM: + k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0) + v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0) + else: + k = tl.load( + k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0 + ) + v = tl.load( + v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0 + ) + # loop over rows + num_block_m = tl.cdiv(seqlen_q, BLOCK_M) + for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M): + start_m = tl.multiple_of(start_m, BLOCK_M) + offs_m_curr = start_m + offs_m + # load q, k, v, do on-chip + # Same bug as below. Otherwise gives wrong result for headdim=40, seqlen=(128, 117) + if EVEN_M & EVEN_HEADDIM: + q = tl.load(q_ptrs) + else: + if EVEN_HEADDIM: + q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0) + else: + q = tl.load( + q_ptrs, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + other=0.0, + ) + # recompute p = softmax(qk, dim=-1).T + qk = tl.dot(q, k, trans_b=True) + # Trying to combine the two masks seem to make the result wrong + if not EVEN_N: # Need to mask out otherwise the softmax is wrong + qk = tl.where(offs_n[None, :] < seqlen_k, qk, float("-inf")) + if IS_CAUSAL: + qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk, float("-inf")) + if BIAS_TYPE != "none": + tl.debug_barrier() # Race condition otherwise + if BIAS_TYPE == "vector": + if EVEN_N: + bias = tl.load(b_ptrs).to(tl.float32) + else: + bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32) + bias = bias[None, :] + elif BIAS_TYPE == "matrix": + if EVEN_M & EVEN_N: + bias = tl.load(b_ptrs).to(tl.float32) + else: + bias = tl.load( + b_ptrs, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k), + other=0.0, + ).to(tl.float32) + qk = qk * softmax_scale + bias + # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong. + # Also wrong for headdim=64. + if not (EVEN_M & EVEN_HEADDIM): + tl.debug_barrier() + lse_i = tl.load(LSE + offs_m_curr) + if BIAS_TYPE == "none": + p = tl.exp(qk * softmax_scale - lse_i[:, None]) + else: + p = tl.exp(qk - lse_i[:, None]) + # compute dv + # [2022-10-30] TD: A Triton bug: if EVEN_M=True and EVEN_HEADDIM=False, if we call + # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0), we get wrong outputs + # in the case of headdim=48/96, seqlen_q & seqlen_k >= 512. If headdim=40 or seqlen < 512, + # the output is correct. + if EVEN_M & EVEN_HEADDIM: + do = tl.load(do_ptrs) + else: + # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask. + do = tl.load( + do_ptrs, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + other=0.0, + ) + # if EVEN_M: + # if EVEN_HEADDIM: + # do = tl.load(do_ptrs) + # else: + # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0) + # else: + # if EVEN_HEADDIM: + # do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0) + # else: + # do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) + # & (offs_d[None, :] < headdim), other=0.0) + dv += tl.dot(p.to(do.dtype), do, trans_a=True) + # compute dp = dot(v, do) + # There seems to be a race condition when headdim=48/96, and dq, dk are wrong. + # Also wrong for headdim=128, seqlen=(108, 256), and ATOMIC_ADD=True + # Also wrong for headdim=64, seqlen=(1023, 1024), and ATOMIC_ADD=False + if not (EVEN_M & EVEN_HEADDIM): + tl.debug_barrier() + dp = tl.dot(do, v, trans_b=True) + # There's a race condition for headdim=48 + if not EVEN_HEADDIM: + tl.debug_barrier() + # compute ds = p * (dp - delta[:, None]) + # Putting the subtraction after the dp matmul (instead of before) is slightly faster + Di = tl.load(D + offs_m_curr) + # Converting ds to q.dtype here reduces register pressure and makes it much faster + # for BLOCK_HEADDIM=128 + ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype) + # compute dk = dot(ds.T, q) + dk += tl.dot(ds, q, trans_a=True) + # compute dq + if not ( + EVEN_M & EVEN_HEADDIM + ): # Otherewise there's a race condition when BIAS_TYPE='matrix' + tl.debug_barrier() + if not ATOMIC_ADD: + if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M + dq = tl.load(dq_ptrs, eviction_policy="evict_last") + dq += tl.dot(ds, k) + tl.store(dq_ptrs, dq, eviction_policy="evict_last") + else: + if EVEN_HEADDIM: + dq = tl.load( + dq_ptrs, + mask=offs_m_curr[:, None] < seqlen_q, + other=0.0, + eviction_policy="evict_last", + ) + dq += tl.dot(ds, k) + tl.store( + dq_ptrs, + dq, + mask=offs_m_curr[:, None] < seqlen_q, + eviction_policy="evict_last", + ) + else: + dq = tl.load( + dq_ptrs, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + other=0.0, + eviction_policy="evict_last", + ) + dq += tl.dot(ds, k) + tl.store( + dq_ptrs, + dq, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + eviction_policy="evict_last", + ) + else: # If we're parallelizing across the seqlen_k dimension + dq = tl.dot(ds, k) + if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M + tl.atomic_add(dq_ptrs, dq) + else: + if EVEN_HEADDIM: + tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q) + else: + tl.atomic_add( + dq_ptrs, + dq, + mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), + ) + # increment pointers + dq_ptrs += BLOCK_M * stride_dqm + q_ptrs += BLOCK_M * stride_qm + do_ptrs += BLOCK_M * stride_dom + if BIAS_TYPE == "matrix": + b_ptrs += BLOCK_M * stride_bm + # write-back + dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :]) + dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :]) + _bwd_store_dk_dv( + dk_ptrs, + dv_ptrs, + dk, + dv, + offs_n, + offs_d, + seqlen_k, + headdim, + EVEN_M=EVEN_M, + EVEN_N=EVEN_N, + EVEN_HEADDIM=EVEN_HEADDIM, + ) + + +def init_to_zero(name): + return lambda nargs: nargs[name].zero_() + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "SEQUENCE_PARALLEL": False}, + num_warps=8, + num_stages=1, + pre_hook=init_to_zero("DQ"), + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "SEQUENCE_PARALLEL": True}, + num_warps=8, + num_stages=1, + pre_hook=init_to_zero("DQ"), + ), + # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now + # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4* + # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), + # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), + # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')), + # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')), + ], + key=["CACHE_KEY_SEQLEN_Q", "CACHE_KEY_SEQLEN_K", "BIAS_TYPE", "IS_CAUSAL", "BLOCK_HEADDIM"], +) +@triton.heuristics( + { + "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0, + "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0, + "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"], + } +) +@triton.jit +def _bwd_kernel( + Q, + K, + V, + Bias, + DO, + DQ, + DK, + DV, + LSE, + D, + softmax_scale, + stride_qb, + stride_qh, + stride_qm, + stride_kb, + stride_kh, + stride_kn, + stride_vb, + stride_vh, + stride_vn, + stride_bb, + stride_bh, + stride_bm, + stride_dob, + stride_doh, + stride_dom, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dvb, + stride_dvh, + stride_dvn, + nheads, + seqlen_q, + seqlen_k, + seqlen_q_rounded, + headdim, + CACHE_KEY_SEQLEN_Q, + CACHE_KEY_SEQLEN_K, + BIAS_TYPE: tl.constexpr, + IS_CAUSAL: tl.constexpr, + BLOCK_HEADDIM: tl.constexpr, + SEQUENCE_PARALLEL: tl.constexpr, + EVEN_M: tl.constexpr, + EVEN_N: tl.constexpr, + EVEN_HEADDIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + off_hb = tl.program_id(1) + off_b = off_hb // nheads + off_h = off_hb % nheads + # offset pointers for batch/head + Q += off_b * stride_qb + off_h * stride_qh + K += off_b * stride_kb + off_h * stride_kh + V += off_b * stride_vb + off_h * stride_vh + DO += off_b * stride_dob + off_h * stride_doh + DQ += off_b * stride_dqb + off_h * stride_dqh + DK += off_b * stride_dkb + off_h * stride_dkh + DV += off_b * stride_dvb + off_h * stride_dvh + if BIAS_TYPE != "none": + Bias += off_b * stride_bb + off_h * stride_bh + # pointer to row-wise quantities in value-like data + D += off_hb * seqlen_q_rounded + LSE += off_hb * seqlen_q_rounded + if not SEQUENCE_PARALLEL: + num_block_n = tl.cdiv(seqlen_k, BLOCK_N) + for start_n in range(0, num_block_n): + _bwd_kernel_one_col_block( + start_n, + Q, + K, + V, + Bias, + DO, + DQ, + DK, + DV, + LSE, + D, + softmax_scale, + stride_qm, + stride_kn, + stride_vn, + stride_bm, + stride_dom, + stride_dqm, + stride_dkn, + stride_dvn, + seqlen_q, + seqlen_k, + headdim, + ATOMIC_ADD=False, + BIAS_TYPE=BIAS_TYPE, + IS_CAUSAL=IS_CAUSAL, + BLOCK_HEADDIM=BLOCK_HEADDIM, + EVEN_M=EVEN_M, + EVEN_N=EVEN_N, + EVEN_HEADDIM=EVEN_HEADDIM, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + ) + else: + start_n = tl.program_id(0) + _bwd_kernel_one_col_block( + start_n, + Q, + K, + V, + Bias, + DO, + DQ, + DK, + DV, + LSE, + D, + softmax_scale, + stride_qm, + stride_kn, + stride_vn, + stride_bm, + stride_dom, + stride_dqm, + stride_dkn, + stride_dvn, + seqlen_q, + seqlen_k, + headdim, + ATOMIC_ADD=True, + BIAS_TYPE=BIAS_TYPE, + IS_CAUSAL=IS_CAUSAL, + BLOCK_HEADDIM=BLOCK_HEADDIM, + EVEN_M=EVEN_M, + EVEN_N=EVEN_N, + EVEN_HEADDIM=EVEN_HEADDIM, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + ) + + +def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None): + # shape constraints + batch, seqlen_q, nheads, d = q.shape + _, seqlen_k, _, _ = k.shape + assert k.shape == (batch, seqlen_k, nheads, d) + assert v.shape == (batch, seqlen_k, nheads, d) + assert d <= 128, "FlashAttention only support head dimensions up to 128" + assert q.dtype == k.dtype == v.dtype, "All tensors must have the same type" + assert q.dtype in [torch.float16, torch.bfloat16], "Only support fp16 and bf16" + assert q.is_cuda and k.is_cuda and v.is_cuda + softmax_scale = softmax_scale or 1.0 / math.sqrt(d) + + has_bias = bias is not None + bias_type = "none" + if has_bias: + assert bias.dtype in [q.dtype, torch.float] + assert bias.is_cuda + assert bias.dim() == 4 + if bias.stride(-1) != 1: + bias = bias.contiguous() + if bias.shape[2:] == (1, seqlen_k): + bias_type = "vector" + elif bias.shape[2:] == (seqlen_q, seqlen_k): + bias_type = "matrix" + else: + raise RuntimeError( + "Last 2 dimensions of bias must be (1, seqlen_k)" " or (seqlen_q, seqlen_k)" + ) + bias = bias.expand(batch, nheads, seqlen_q, seqlen_k) + bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0) + + seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128 + lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32) + tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32) + o = torch.empty_like(q) + + BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16) + BLOCK = 128 + num_warps = 4 if d <= 64 else 8 + grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads) + _fwd_kernel[grid]( + q, + k, + v, + bias, + o, + lse, + tmp, + softmax_scale, + q.stride(0), + q.stride(2), + q.stride(1), + k.stride(0), + k.stride(2), + k.stride(1), + v.stride(0), + v.stride(2), + v.stride(1), + *bias_strides, + o.stride(0), + o.stride(2), + o.stride(1), + nheads, + seqlen_q, + seqlen_k, + seqlen_q_rounded, + d, + seqlen_q // 32, + seqlen_k // 32, # key for triton cache (limit number of compilations) + # Can't use kwargs here because triton autotune expects key to be args, not kwargs + # IS_CAUSAL=causal, BLOCK_HEADDIM=d, + bias_type, + causal, + BLOCK_HEADDIM, + BLOCK_M=BLOCK, + BLOCK_N=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + return o, lse, softmax_scale # softmax_scale could have been updated + + +def _flash_attn_backward( + do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None +): + # Make sure that the last dimension is contiguous + if do.stride(-1) != 1: + do = do.contiguous() + batch, seqlen_q, nheads, d = q.shape + _, seqlen_k, _, _ = k.shape + # assert d in {16, 32, 64, 128} + assert d <= 128 + seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128 + assert lse.shape == (batch, nheads, seqlen_q_rounded) + assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1 + assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1 + softmax_scale = softmax_scale or 1.0 / math.sqrt(d) + # dq_accum = torch.zeros_like(q, dtype=torch.float32) + dq_accum = torch.empty_like(q, dtype=torch.float32) + delta = torch.empty_like(lse) + # delta = torch.zeros_like(lse) + + BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16) + grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads) + _bwd_preprocess_do_o_dot[grid]( + o, + do, + delta, + o.stride(0), + o.stride(2), + o.stride(1), + do.stride(0), + do.stride(2), + do.stride(1), + nheads, + seqlen_q, + seqlen_q_rounded, + d, + BLOCK_M=128, + BLOCK_HEADDIM=BLOCK_HEADDIM, + ) + + has_bias = bias is not None + bias_type = "none" + if has_bias: + assert bias.dtype in [q.dtype, torch.float] + assert bias.is_cuda + assert bias.dim() == 4 + assert bias.stride(-1) == 1 + if bias.shape[2:] == (1, seqlen_k): + bias_type = "vector" + elif bias.shape[2:] == (seqlen_q, seqlen_k): + bias_type = "matrix" + else: + raise RuntimeError( + "Last 2 dimensions of bias must be (1, seqlen_k)" " or (seqlen_q, seqlen_k)" + ) + bias = bias.expand(batch, nheads, seqlen_q, seqlen_k) + bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0) + + # BLOCK_M = 128 + # BLOCK_N = 64 + # num_warps = 4 + grid = lambda META: ( + triton.cdiv(seqlen_k, META["BLOCK_N"]) if META["SEQUENCE_PARALLEL"] else 1, + batch * nheads, + ) + _bwd_kernel[grid]( + q, + k, + v, + bias, + do, + dq_accum, + dk, + dv, + lse, + delta, + softmax_scale, + q.stride(0), + q.stride(2), + q.stride(1), + k.stride(0), + k.stride(2), + k.stride(1), + v.stride(0), + v.stride(2), + v.stride(1), + *bias_strides, + do.stride(0), + do.stride(2), + do.stride(1), + dq_accum.stride(0), + dq_accum.stride(2), + dq_accum.stride(1), + dk.stride(0), + dk.stride(2), + dk.stride(1), + dv.stride(0), + dv.stride(2), + dv.stride(1), + nheads, + seqlen_q, + seqlen_k, + seqlen_q_rounded, + d, + seqlen_q // 32, + seqlen_k // 32, # key for triton cache (limit number of compilations) + # Can't use kwargs here because triton autotune expects key to be args, not kwargs + # IS_CAUSAL=causal, BLOCK_HEADDIM=d, + bias_type, + causal, + BLOCK_HEADDIM, + # SEQUENCE_PARALLEL=False, + # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, + # num_warps=num_warps, + # num_stages=1, + ) + dq.copy_(dq_accum) + + +class FlashAttnQKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None): + """ + qkv: (batch, seqlen, 3, nheads, headdim) + bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen). + For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen). + ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen) + """ + # Make sure that the last dimension is contiguous + if qkv.stride(-1) != 1: + qkv = qkv.contiguous() + o, lse, ctx.softmax_scale = _flash_attn_forward( + qkv[:, :, 0], + qkv[:, :, 1], + qkv[:, :, 2], + bias=bias, + causal=causal, + softmax_scale=softmax_scale, + ) + ctx.save_for_backward(qkv, o, lse, bias) + ctx.causal = causal + return o + + @staticmethod + def backward(ctx, do): + qkv, o, lse, bias = ctx.saved_tensors + assert not ctx.needs_input_grad[1], "FlashAttention does not support bias gradient yet" + # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd + # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version. + with torch.inference_mode(): + dqkv = torch.empty_like(qkv) + _flash_attn_backward( + do, + qkv[:, :, 0], + qkv[:, :, 1], + qkv[:, :, 2], + o, + lse, + dqkv[:, :, 0], + dqkv[:, :, 1], + dqkv[:, :, 2], + bias=bias, + causal=ctx.causal, + softmax_scale=ctx.softmax_scale, + ) + return dqkv, None, None, None + + +flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply + + +class FlashAttnKVPackedFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None): + """ + q: (batch, seqlen_q, nheads, headdim) + kv: (batch, seqlen_k, 2, nheads, headdim) + bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k). + For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k). + ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k) + """ + # Make sure that the last dimension is contiguous + q, kv = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]] + o, lse, ctx.softmax_scale = _flash_attn_forward( + q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale + ) + ctx.save_for_backward(q, kv, o, lse, bias) + ctx.causal = causal + return o + + @staticmethod + def backward(ctx, do): + q, kv, o, lse, bias = ctx.saved_tensors + if len(ctx.needs_input_grad) >= 3: + assert not ctx.needs_input_grad[2], "FlashAttention does not support bias gradient yet" + # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd + # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version. + with torch.inference_mode(): + dq = torch.empty_like(q) + dkv = torch.empty_like(kv) + _flash_attn_backward( + do, + q, + kv[:, :, 0], + kv[:, :, 1], + o, + lse, + dq, + dkv[:, :, 0], + dkv[:, :, 1], + bias=bias, + causal=ctx.causal, + softmax_scale=ctx.softmax_scale, + ) + return dq, dkv, None, None, None + + +flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply + + +class FlashAttnFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None): + """ + q: (batch_size, seqlen_q, nheads, headdim) + k, v: (batch_size, seqlen_k, nheads, headdim) + bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k). + For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k). + ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k) + """ + # Make sure that the last dimension is contiguous + q, k, v = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]] + o, lse, ctx.softmax_scale = _flash_attn_forward( + q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale + ) + ctx.save_for_backward(q, k, v, o, lse, bias) + ctx.causal = causal + return o + + @staticmethod + def backward(ctx, do): + q, k, v, o, lse, bias = ctx.saved_tensors + assert not ctx.needs_input_grad[3], "FlashAttention does not support bias gradient yet" + # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd + # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version. + with torch.inference_mode(): + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + _flash_attn_backward( + do, + q, + k, + v, + o, + lse, + dq, + dk, + dv, + bias=bias, + causal=ctx.causal, + softmax_scale=ctx.softmax_scale, + ) + return dq, dk, dv, None, None, None + + +flash_attn_func = FlashAttnFunc.apply diff --git a/lib/python3.12/site-packages/flash_attn/flash_attn_triton_og.py b/lib/python3.12/site-packages/flash_attn/flash_attn_triton_og.py new file mode 100644 index 0000000000000000000000000000000000000000..f2ddb99487b4f162745e2f6dd3d1744946dc3fb2 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/flash_attn_triton_og.py @@ -0,0 +1,365 @@ +# [2022-10-23] Downloaded from https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py +# for benchmarking. +# We fixed a few dtype cast to make it work for bf16 + +""" +Fused Attention +=============== +This is a Triton implementation of the Flash Attention algorithm +(see: Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf; Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf) +""" + +import pytest +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fwd_kernel( + Q, + K, + V, + sm_scale, + TMP, + L, + M, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug + Out, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vk, + stride_vn, + stride_oz, + stride_oh, + stride_om, + stride_on, + Z, + H, + N_CTX, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + start_m = tl.program_id(0) + off_hz = tl.program_id(1) + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + off_q = off_hz * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk + off_k = off_hz * stride_qh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk + off_v = off_hz * stride_qh + offs_n[:, None] * stride_qm + offs_d[None, :] * stride_qk + # Initialize pointers to Q, K, V + q_ptrs = Q + off_q + k_ptrs = K + off_k + v_ptrs = V + off_v + # initialize pointer to m and l + t_ptrs = TMP + off_hz * N_CTX + offs_m + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + # load q: it will stay in SRAM throughout + q = tl.load(q_ptrs) + # loop over k, v and update accumulator + for start_n in range(0, (start_m + 1) * BLOCK_M, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load(k_ptrs + start_n * stride_kn) + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k, trans_b=True) + qk *= sm_scale + qk += tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), 0, float("-inf")) + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(m_ij - m_i_new) + l_i_new = alpha * l_i + beta * l_ij + # -- update output accumulator -- + # scale p + p_scale = beta / l_i_new + p = p * p_scale[:, None] + # scale acc + acc_scale = l_i / l_i_new * alpha + tl.store(t_ptrs, acc_scale) + acc_scale = tl.load(t_ptrs) # BUG: have to store and immediately load + acc = acc * acc_scale[:, None] + # update acc + v = tl.load(v_ptrs + start_n * stride_vk) + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + # rematerialize offsets to save registers + start_m = tl.program_id(0) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + # write back l and m + l_ptrs = L + off_hz * N_CTX + offs_m + m_ptrs = M + off_hz * N_CTX + offs_m + tl.store(l_ptrs, l_i) + tl.store(m_ptrs, m_i) + # initialize pointers to output + offs_n = tl.arange(0, BLOCK_DMODEL) + off_o = off_hz * stride_oh + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + out_ptrs = Out + off_o + tl.store(out_ptrs, acc) + + +@triton.jit +def _bwd_preprocess( + Out, + DO, + L, + NewDO, + Delta, + BLOCK_M: tl.constexpr, + D_HEAD: tl.constexpr, +): + off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + off_n = tl.arange(0, D_HEAD) + # load + o = tl.load(Out + off_m[:, None] * D_HEAD + off_n[None, :]).to(tl.float32) + do = tl.load(DO + off_m[:, None] * D_HEAD + off_n[None, :]).to(tl.float32) + denom = tl.load(L + off_m).to(tl.float32) + # compute + do = do / denom[:, None] + delta = tl.sum(o * do, axis=1) + # write-back + tl.store(NewDO + off_m[:, None] * D_HEAD + off_n[None, :], do) + tl.store(Delta + off_m, delta) + + +@triton.jit +def _bwd_kernel( + Q, + K, + V, + sm_scale, + Out, + DO, + DQ, + DK, + DV, + L, + M, + D, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vk, + stride_vn, + Z, + H, + N_CTX, + num_block, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + off_hz = tl.program_id(0) + off_z = off_hz // H + off_h = off_hz % H + # offset pointers for batch/head + Q += off_z * stride_qz + off_h * stride_qh + K += off_z * stride_qz + off_h * stride_qh + V += off_z * stride_qz + off_h * stride_qh + DO += off_z * stride_qz + off_h * stride_qh + DQ += off_z * stride_qz + off_h * stride_qh + DK += off_z * stride_qz + off_h * stride_qh + DV += off_z * stride_qz + off_h * stride_qh + for start_n in range(0, num_block): + lo = start_n * BLOCK_M + # initialize row/col offsets + offs_qm = lo + tl.arange(0, BLOCK_M) + offs_n = start_n * BLOCK_M + tl.arange(0, BLOCK_M) + offs_m = tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_DMODEL) + # initialize pointers to value-like data + q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk) + k_ptrs = K + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk) + v_ptrs = V + (offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk) + do_ptrs = DO + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk) + dq_ptrs = DQ + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk) + # pointer to row-wise quantities in value-like data + D_ptrs = D + off_hz * N_CTX + m_ptrs = M + off_hz * N_CTX + # initialize dv amd dk + dv = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + dk = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + # k and v stay in SRAM throughout + k = tl.load(k_ptrs) + v = tl.load(v_ptrs) + # loop over rows + for start_m in range(lo, num_block * BLOCK_M, BLOCK_M): + offs_m_curr = start_m + offs_m + # load q, k, v, do on-chip + q = tl.load(q_ptrs) + # recompute p = softmax(qk, dim=-1).T + # NOTE: `do` is pre-divided by `l`; no normalization here + qk = tl.dot(q, k, trans_b=True) + qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk, float("-inf")) + m = tl.load(m_ptrs + offs_m_curr) + p = tl.exp(qk * sm_scale - m[:, None]) + # compute dv + do = tl.load(do_ptrs) + dv += tl.dot(p.to(do.dtype), do, trans_a=True) + # compute dp = dot(v, do) + Di = tl.load(D_ptrs + offs_m_curr) + dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - Di[:, None] + dp += tl.dot(do, v, trans_b=True) + # compute ds = p * (dp - delta[:, None]) + ds = p * dp * sm_scale + # compute dk = dot(ds.T, q) + dk += tl.dot(ds.to(q.dtype), q, trans_a=True) + # # compute dq + dq = tl.load(dq_ptrs, eviction_policy="evict_last") + dq += tl.dot(ds.to(k.dtype), k) + tl.store(dq_ptrs, dq, eviction_policy="evict_last") + # # increment pointers + dq_ptrs += BLOCK_M * stride_qm + q_ptrs += BLOCK_M * stride_qm + do_ptrs += BLOCK_M * stride_qm + # write-back + dv_ptrs = DV + (offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk) + dk_ptrs = DK + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk) + tl.store(dv_ptrs, dv) + tl.store(dk_ptrs, dk) + + +class _attention(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, sm_scale): + BLOCK = 128 + # shape constraints + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + assert Lq == Lk and Lk == Lv + assert Lk in {16, 32, 64, 128} + o = torch.empty_like(q) + grid = (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1]) + tmp = torch.empty( + (q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32 + ) + L = torch.empty((q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + m = torch.empty((q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + num_warps = 4 if Lk <= 64 else 8 + + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + tmp, + L, + m, + o, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + q.shape[0], + q.shape[1], + q.shape[2], + BLOCK_M=BLOCK, + BLOCK_N=BLOCK, + BLOCK_DMODEL=Lk, + num_warps=num_warps, + num_stages=1, + ) + ctx.save_for_backward(q, k, v, o, L, m) + ctx.BLOCK = BLOCK + ctx.grid = grid + ctx.sm_scale = sm_scale + ctx.BLOCK_DMODEL = Lk + return o + + @staticmethod + def backward(ctx, do): + q, k, v, o, l, m = ctx.saved_tensors + do = do.contiguous() + dq = torch.zeros_like(q, dtype=torch.float32) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + do_scaled = torch.empty_like(do) + delta = torch.empty_like(l) + _bwd_preprocess[(ctx.grid[0] * ctx.grid[1],)]( + o, + do, + l, + do_scaled, + delta, + BLOCK_M=ctx.BLOCK, + D_HEAD=ctx.BLOCK_DMODEL, + ) + + # NOTE: kernel currently buggy for other values of `num_warps` + num_warps = 8 + _bwd_kernel[(ctx.grid[1],)]( + q, + k, + v, + ctx.sm_scale, + o, + do_scaled, + dq, + dk, + dv, + l, + m, + delta, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + q.shape[0], + q.shape[1], + q.shape[2], + ctx.grid[0], + BLOCK_M=ctx.BLOCK, + BLOCK_N=ctx.BLOCK, + BLOCK_DMODEL=ctx.BLOCK_DMODEL, + num_warps=num_warps, + num_stages=1, + ) + return dq.to(q.dtype), dk, dv, None + + +attention = _attention.apply diff --git a/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attention.py b/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..4c930291063db1350ce05aecf55c6050bad28fd4 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attention.py @@ -0,0 +1,197 @@ +import math + +import hydra +import torch +import torch.nn as nn +from einops import rearrange + +from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input +from flash_attn.flash_blocksparse_attn_interface import ( + convert_blockmask, + flash_blocksparse_attn_func, +) + + +class FlashBlocksparseAttention(nn.Module): + """Implement the scaled dot product attention with softmax. + Arguments + --------- + softmax_temp: The temperature to use for the softmax attention. + (default: 1/sqrt(d_keys) where d_keys is computed at + runtime) + attention_dropout: The dropout rate to apply to the attention + (default: 0.1) + """ + + def __init__( + self, + sparsity_config, + softmax_temp=None, + attention_dropout=0.0, + max_seq_length=2048, + device=None, + dtype=None, + ): + super().__init__() + self.sparsity_config = hydra.utils.instantiate(sparsity_config) + self.softmax_temp = softmax_temp + self.dropout_p = attention_dropout + + # initialize sparse layout and register as buffer + max_seq_length = ((max_seq_length + 256 - 1) // 256) * 256 + layout = self.sparsity_config.make_layout(max_seq_length) + self.register_buffer("layout", layout) + blockmask_converted = convert_blockmask(self.layout, causal=False) + self.register_buffer("blockmask_converted", blockmask_converted) + # logger.info(f'Attention class {self.__class__}: saving={self.layout.float().mean()}') + + def forward( + self, + qkv, + attn_mask=None, + key_padding_mask=None, + causal=False, + cu_seqlens=None, + max_s=None, + need_weights=False, + convert_mask=True, + ): + """Implements the multihead softmax attention. + Arguments + --------- + qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None + attn_mask: An implementation of BaseMask that encodes where each + query can attend to + key_padding_mask: An implementation of BaseMask that encodes how + many query each sequence in the batch consists of + """ + assert not need_weights + assert attn_mask is None + assert qkv.dtype == torch.float16 + assert qkv.is_cuda + + if cu_seqlens is None: + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + # Convert mask to take a subset + seqlen_rounded = ((seqlen + 256 - 1) // 256) * 256 + assert seqlen_rounded // 16 <= self.layout.shape[0], ( + seqlen_rounded // 256 <= self.layout.shape[1] + ) + blockmask = self.layout[: seqlen_rounded // 16, : seqlen_rounded // 256] + if key_padding_mask is None: + qkv = rearrange(qkv, "b s ... -> (b s) ...") + max_s = seqlen + cu_seqlens = torch.arange( + 0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32, device=qkv.device + ) + output = flash_blocksparse_attn_func( + qkv, + cu_seqlens, + blockmask, + self.dropout_p if self.training else 0.0, + max_s, + softmax_scale=self.softmax_temp, + causal=causal, + ) + output = rearrange(output, "(b s) ... -> b s ...", b=batch_size) + else: + key_padding_mask_bool = key_padding_mask.bool_matrix + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask_bool) + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_blocksparse_attn_func( + x_unpad, + cu_seqlens, + blockmask, + self.dropout_p if self.training else 0.0, + max_s, + softmax_scale=self.softmax_temp, + causal=causal, + ) + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen + ), + "b s (h d) -> b s h d", + h=nheads, + ) + else: + assert max_s is not None + seqlen = max_s + # Convert mask to take a subset + seqlen_rounded = ((seqlen + 256 - 1) // 256) * 256 + assert seqlen_rounded // 16 <= self.layout.shape[0], ( + seqlen_rounded // 256 <= self.layout.shape[1] + ) + blockmask = self.layout[: seqlen_rounded // 16, : seqlen_rounded // 256] + if convert_mask: + output = flash_blocksparse_attn_func( + qkv, + cu_seqlens, + blockmask, + self.dropout_p if self.training else 0.0, + max_s, + softmax_scale=self.softmax_temp, + causal=causal, + ) + else: + output = flash_blocksparse_attn_func( + qkv, + cu_seqlens, + self.blockmask_converted, + self.dropout_p if self.training else 0.0, + max_s, + softmax_scale=self.softmax_temp, + causal=causal, + convert_mask=False, + ) + + return output, None + + +class FlashBlocksparseMHA(nn.Module): + def __init__( + self, + embed_dim, + num_heads, + sparsity_config, + bias=True, + batch_first=True, + attention_dropout=0.0, + causal=False, + max_seq_length=2048, + device=None, + dtype=None, + **kwargs, + ) -> None: + assert batch_first + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.embed_dim = embed_dim + self.causal = causal + + self.num_heads = num_heads + assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads" + self.head_dim = self.embed_dim // num_heads + assert self.head_dim in [16, 32, 64], "Only support head_dim == 16, 32, or 64" + + self.Wqkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs) + self.inner_attn = FlashBlocksparseAttention( + sparsity_config, + attention_dropout=attention_dropout, + max_seq_length=max_seq_length, + **factory_kwargs, + ) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs) + + def forward( + self, x, x_ignored_, x_ignored_1_, attn_mask=None, key_padding_mask=None, need_weights=False + ): + qkv = self.Wqkv(x) + qkv = rearrange(qkv, "b s (three h d) -> b s three h d", three=3, h=self.num_heads) + context, attn_weights = self.inner_attn( + qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=self.causal + ) + return self.out_proj(rearrange(context, "b s h d -> b s (h d)")), attn_weights diff --git a/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attn_interface.py b/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attn_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..9ce3fe8c1344dd33165c43e4cc1ef0f70feb5d04 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/flash_blocksparse_attn_interface.py @@ -0,0 +1,200 @@ +# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/fmha.py +import flash_attn_cuda +import torch +import torch.nn as nn + + +def convert_blockmask(blockmask, causal): + """Convert from the 0-1 format to the format used by the CUDA code. + 0 means the block is skipped. + nonzero means the block is not skipped. + Argument: + blockmask: (row, col): a 0-1 tensor + Return: + blockmask_converted: (col, row), dtype torch.int32: for each column, it contains the row + indices of the nonzero blocks, padded with -1 to reach length @row. + The indices are multiplied by 4, with the smallest bit used to encode whether + it is the first nonzero in its row, and the 2nd smallest bit to encode whether it is + the last nonzero in its row.. + """ + assert not causal + # TD [2022-05-13]: The indexing and sorting is very tricky + nrow, ncol = blockmask.shape + # Sort does not support bool on CUDA + blockmask = blockmask.to(dtype=torch.uint8) + nonzero_val, nonzero_sorted_rowidx = blockmask.sort(dim=0, stable=True, descending=True) + nonzero_unsorted_rowidx = nonzero_sorted_rowidx.argsort(dim=0) + last_nonzero_col_per_row = blockmask.sort(dim=-1, stable=True).indices[:, -1] + last_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[ + torch.arange(nrow, device=blockmask.device), last_nonzero_col_per_row + ] + first_nonzero_col_per_row = blockmask.sort(dim=-1, stable=True, descending=True).indices[:, 0] + first_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[ + torch.arange(nrow, device=blockmask.device), first_nonzero_col_per_row + ] + nonzero_idx = nonzero_sorted_rowidx * 4 + nonzero_idx[last_nonzero_col_per_row_after_sort, last_nonzero_col_per_row] += 2 + nonzero_idx[first_nonzero_col_per_row_after_sort, first_nonzero_col_per_row] += 1 + nonzero_idx[nonzero_val == 0] = -1 + return nonzero_idx.T.contiguous().to(dtype=torch.int32) + + +def _flash_blocksparse_attn_forward( + qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal, return_softmax +): + context, softmax_lse, *rest = flash_attn_cuda.fwd_block( + qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal, return_softmax, None + ) + # if context.isnan().any() or softmax_lse.isnan().any(): + # breakpoint() + S_dmask = rest[0] if return_softmax else None + return context, softmax_lse, S_dmask + + +def _flash_blocksparse_attn_backward( + dout, + qkv, + out, + S_dmask, + softmax_lse, + cu_seqlens, + blockmask, + dropout_p, + max_s, + softmax_scale, + causal, +): + dqkv, dp, softmax_d = flash_attn_cuda.bwd_block( + dout, + qkv, + out, + S_dmask, + softmax_lse, + cu_seqlens, + blockmask, + dropout_p, + softmax_scale, + max_s, + causal, + None, + ) + # if dqkv.isnan().any() or softmax_d.isnan().any(): + # breakpoint() + return dqkv + + +class FlashBlocksparseAttnFun(torch.autograd.Function): + @staticmethod + def forward(ctx, qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal): + # Save rng_state because the backward pass will regenerate the dropout mask + rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + context, softmax_lse, S_dmask = _flash_blocksparse_attn_forward( + qkv, + cu_seqlens, + blockmask, + dropout_p, + max_s, + softmax_scale, + causal=causal, + return_softmax=False, + ) + ctx.save_for_backward(qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state) + ctx.dropout_p = dropout_p + ctx.max_s = max_s + ctx.softmax_scale = softmax_scale + ctx.causal = causal + return context + + @staticmethod + def backward(ctx, dout): + qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state = ctx.saved_tensors + if rng_state is not None: + cur_rng_state = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(rng_state) + # S_dmask is None, temporarily use another tensor just to get it running + dqkv = _flash_blocksparse_attn_backward( + dout, + qkv, + context, + context, + softmax_lse, + cu_seqlens, + blockmask, + ctx.dropout_p, + ctx.max_s, + ctx.softmax_scale, + ctx.causal, + ) + if rng_state is not None: + torch.cuda.set_rng_state(cur_rng_state) + return dqkv, None, None, None, None, None, None, None + + +# We duplicate code to return both the output and the softmax for testing +# Returning both makes backward a bit slower, so we want to keep using the other version for speed. +class FlashBlocksparseAttnFunWithS(torch.autograd.Function): + @staticmethod + def forward(ctx, qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal): + # Save rng_state because the backward pass is gonna regenerate the dropout mask + rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None + if softmax_scale is None: + softmax_scale = qkv.shape[-1] ** (-0.5) + context, softmax_lse, S_dmask = _flash_blocksparse_attn_forward( + qkv, + cu_seqlens, + blockmask, + dropout_p, + max_s, + softmax_scale, + causal=causal, + return_softmax=True, + ) + ctx.save_for_backward(qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state) + ctx.dropout_p = dropout_p + ctx.max_s = max_s + ctx.softmax_scale = softmax_scale + ctx.causal = causal + return context, S_dmask, softmax_lse + + @staticmethod + def backward(ctx, dout, _dS_dmask_ignored, _dsoftmax_sum_ignored): + qkv, context, S_dmask, softmax_lse, cu_seqlens, blockmask, rng_state = ctx.saved_tensors + if rng_state is not None: + cur_rng_state = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(rng_state) + dqkv = _flash_blocksparse_attn_backward( + dout, + qkv, + context, + S_dmask, + softmax_lse, + cu_seqlens, + blockmask, + ctx.dropout_p, + ctx.max_s, + ctx.softmax_scale, + ctx.causal, + ) + if rng_state is not None: + torch.cuda.set_rng_state(cur_rng_state) + return dqkv, None, None, None, None, None, None + + +def flash_blocksparse_attn_func( + qkv, + cu_seqlens, + blockmask, + dropout_p, + max_s, + softmax_scale=None, + causal=False, + return_attn_probs=False, + convert_mask=True, +): + """dropout_p should be set to 0.0 during evaluation""" + func = FlashBlocksparseAttnFun if not return_attn_probs else FlashBlocksparseAttnFunWithS + if convert_mask: + blockmask = convert_blockmask(blockmask, causal=causal) + return func.apply(qkv, cu_seqlens, blockmask, dropout_p, max_s, softmax_scale, causal) diff --git a/lib/python3.12/site-packages/flash_attn/fused_softmax.py b/lib/python3.12/site-packages/flash_attn/fused_softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..382f94f092cd3999b2378dfc2fa165a7c08017e2 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/fused_softmax.py @@ -0,0 +1,201 @@ +# [2022-10-23] Copied from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_softmax.py +# for benchmarking. +# We added support for seqlen=2k and seqlen=4k + +# coding=utf-8 +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +from apex._autocast_utils import _cast_if_autocast_enabled +from apex.transformer.enums import AttnMaskType +from fused_softmax_lib import ( + scaled_masked_softmax_backward, + scaled_masked_softmax_forward, + scaled_masked_softmax_get_batch_per_block, + scaled_upper_triang_masked_softmax_backward, + scaled_upper_triang_masked_softmax_forward, +) + + +class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function): + """ + Fused operation which performs following three operations in sequence + 1. Scale the tensor. + 2. Apply upper triangular mask (typically used in gpt models). + 3. Perform softmax. + """ + + @staticmethod + def forward(ctx, inputs, scale): + scale_t = torch.tensor([scale]) + softmax_results = scaled_upper_triang_masked_softmax_forward(inputs, scale_t[0]) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + softmax_results, scale_t = ctx.saved_tensors + input_grads = scaled_upper_triang_masked_softmax_backward( + output_grads, softmax_results, scale_t[0] + ) + return input_grads, None + + +def scaled_upper_triang_masked_softmax(inputs, _, scale): + b, np, sq, sk = inputs.size() + assert sq == sk, "causal mask is only for self attention" + # Reshaping input to 3D tensor (attn_batches, sq, sk) + inputs = inputs.view(-1, sq, sk) + args = _cast_if_autocast_enabled(inputs, scale) + with torch.cuda.amp.autocast(enabled=False): + probs = ScaledUpperTriangMaskedSoftmax.apply(*args) + return probs.view(b, np, sq, sk) + + +# NOTE (mkozuki): `ScaledMaskedSoftmax` somehow doesn't work well with `torch.cuda.amp.custom_fwd`. +# Without `cast_inputs` kwarg, somehow inputs are not cast to dtype used in the autocast context. +# So I needed to manually write two `torch.autograd.Function` inheritances. +# Fused operation which performs following three operations in sequence +# 1. Scale the tensor. +# 2. Apply the mask. +# 3. Perform softmax. +class ScaledMaskedSoftmax(torch.autograd.Function): + @staticmethod + def forward(ctx, inputs, mask, scale): + scale_t = torch.tensor([scale]) + softmax_results = scaled_masked_softmax_forward(inputs, mask, scale_t[0]) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + softmax_results, scale_t = ctx.saved_tensors + input_grads = scaled_masked_softmax_backward(output_grads, softmax_results, scale_t[0]) + return input_grads, None, None + + +def scaled_masked_softmax(inputs, mask, scale): + # input is 4D tensor (b, np, sq, sk) + args = _cast_if_autocast_enabled(inputs, mask, scale) + with torch.cuda.amp.autocast(enabled=False): + return ScaledMaskedSoftmax.apply(*args) + + +class FusedScaleMaskSoftmax(torch.nn.Module): + """ + fused operation: scaling + mask + softmax + + Arguments: + input_in_fp16: flag to indicate if input in fp16 data format. + input_in_bf16: flag to indicate if input in bf16 data format. + attn_mask_type: attention mask type (pad or causal) + scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion + mask_func: mask function to be applied. + softmax_in_fp32: if true, softmax in performed at fp32 precision. + scale: scaling factor used in input tensor scaling. + """ + + def __init__( + self, + input_in_fp16, + input_in_bf16, + attn_mask_type, + scaled_masked_softmax_fusion, + mask_func, + softmax_in_fp32, + scale, + ): + super().__init__() + self.input_in_fp16 = input_in_fp16 + self.input_in_bf16 = input_in_bf16 + if self.input_in_fp16 and self.input_in_bf16: + raise RuntimeError("both fp16 and bf16 flags cannot be active at the same time.") + self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16 + self.attn_mask_type = attn_mask_type + self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion + self.mask_func = mask_func + self.softmax_in_fp32 = softmax_in_fp32 + self.scale = scale + + if not (self.scale is None or softmax_in_fp32): + raise RuntimeError("softmax should be in fp32 when scaled") + + if self.scaled_masked_softmax_fusion: + if self.attn_mask_type == AttnMaskType.causal: + self.fused_softmax_func = scaled_upper_triang_masked_softmax + elif self.attn_mask_type == AttnMaskType.padding: + self.fused_softmax_func = scaled_masked_softmax + else: + raise ValueError("Invalid attn_mask_type.") + + def forward(self, input, mask): + # [b, np, sq, sk] + assert input.dim() == 4 + + if self.is_kernel_available(mask, *input.size()): + return self.forward_fused_softmax(input, mask) + else: + return self.forward_torch_softmax(input, mask) + + def is_kernel_available(self, mask, b, np, sq, sk): + attn_batches = b * np + + if ( + self.scaled_masked_softmax_fusion # user want to fuse + and self.input_in_float16 # input must be fp16 + and ( + self.attn_mask_type == AttnMaskType.causal + or (self.attn_mask_type == AttnMaskType.padding and mask is not None) + ) + and 16 < sk <= 8192 # sk must be 16 ~ 8192 + and sq % 4 == 0 # sq must be divisor of 4 + and sk % 4 == 0 # sk must be divisor of 4 + and attn_batches % 4 == 0 # np * b must be divisor of 4 + ): + if 0 <= sk <= 8192: + batch_per_block = self.get_batch_per_block(sq, sk, b, np) + + if self.attn_mask_type == AttnMaskType.causal: + if attn_batches % batch_per_block == 0: + return True + else: + if sq % batch_per_block == 0: + return True + return False + + def forward_fused_softmax(self, input, mask): + # input.shape = [b, np, sq, sk] + scale = self.scale if self.scale is not None else 1.0 + return self.fused_softmax_func(input, mask, scale) + + def forward_torch_softmax(self, input, mask): + if self.input_in_float16 and self.softmax_in_fp32: + input = input.float() + + if self.scale is not None: + input = input * self.scale + mask_output = self.mask_func(input, mask) if mask is not None else input + probs = torch.nn.Softmax(dim=-1)(mask_output) + + if self.input_in_float16 and self.softmax_in_fp32: + if self.input_in_fp16: + probs = probs.half() + else: + probs = probs.bfloat16() + + return probs + + @staticmethod + def get_batch_per_block(sq, sk, b, np): + return scaled_masked_softmax_get_batch_per_block(sq, sk, b, np) diff --git a/lib/python3.12/site-packages/flash_attn/ops/__init__.py b/lib/python3.12/site-packages/flash_attn/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/flash_attn/ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94e294dc99e193e52b1aad02e8cddd09854cd810 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/__pycache__/activations.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/activations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83dcf2f2a73b2c9a785d1fab8f713ea3895ac607 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/activations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/__pycache__/fused_dense.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/fused_dense.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae832c09fb87e74f3c1712faeb12e4b4f0b57ee1 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/fused_dense.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/__pycache__/layer_norm.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/layer_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41d06a58bf78357d70277e51781911b60ced9ca2 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/layer_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/__pycache__/rms_norm.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/rms_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c5b3ce00dd1b3ca8bef47ae27f18724afca06ee Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/__pycache__/rms_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/activations.py b/lib/python3.12/site-packages/flash_attn/ops/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..7c09649fc41e12d5a360c5672825d8380bc7ec80 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/activations.py @@ -0,0 +1,135 @@ +# Copied from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/model/layers/activations.py +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 1/sqrt(2*pi)-> 0.3989423 +# 1/sqrt(2) -> 0.70710678 +# sqrt(2/pi) -> 0.79788456 + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.jit.script +def bias_gelu(y, bias): + x = bias + y + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.jit.script +def bias_gelu_back(g, y, bias): + """Assume that y has shape (B, D) and bias has shape (D)""" + x = bias + y + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + grad_y = ff * g + return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) + + +class GeLUFunction(torch.autograd.Function): + @staticmethod + # bias is an optional argument + def forward(ctx, input, bias): + ctx.save_for_backward(input, bias) + return bias_gelu(input, bias) + + @staticmethod + def backward(ctx, grad_output): + input, bias = ctx.saved_tensors + tmp = bias_gelu_back(grad_output, input, bias) + return tmp, tmp + + +bias_gelu_impl = GeLUFunction.apply + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.jit.script +def gelu_fwd(x): + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.jit.script +def gelu_bwd(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + return (ff * g).to(dtype=x.dtype) + + +class FastGeLUFunction(torch.autograd.Function): + @staticmethod + # bias is an optional argument + def forward(ctx, input): + ctx.save_for_backward(input) + return gelu_fwd(input) + + @staticmethod + def backward(ctx, grad_output): + (input,) = ctx.saved_tensors + tmp = gelu_bwd(grad_output, input) + return tmp + + +fast_gelu_impl = FastGeLUFunction.apply + + +@torch.jit.script +def relu_bwd(g, x): + return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) + + +@torch.jit.script +def sqrelu_fwd(x): + r = F.relu(x) + return (r * r).to(dtype=x.dtype) + + +@torch.jit.script +def sqrelu_bwd(g, x): + return (2.0 * g * F.relu(x)).to(dtype=x.dtype) + + +swiglu_fwd_codestring = """ +template T swiglu_fwd(T x, T y) { + return float(x) * float(y) / (1.0f + ::exp(-float(x))); +} +""" +swiglu_bwd_codestring = """ +template void swiglu_bwd(T x, T y, T g, T& dx, T& dy) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); + dy = float(x) * x_sigmoid * float(g); +} +""" +swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) +swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) + + +class SwiGLUFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, y): + ctx.save_for_backward(x, y) + return swiglu_fwd(x, y) + + @staticmethod + def backward(ctx, dout): + x, y = ctx.saved_tensors + return swiglu_bwd(x, y, dout) + +swiglu = SwiGLUFunction.apply diff --git a/lib/python3.12/site-packages/flash_attn/ops/fused_dense.py b/lib/python3.12/site-packages/flash_attn/ops/fused_dense.py new file mode 100644 index 0000000000000000000000000000000000000000..1e45b8e609812a1545781011141ec80f6dc3af0f --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/fused_dense.py @@ -0,0 +1,688 @@ +# Copyright (c) 2023, Tri Dao. +# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py +# We make it work with pytorch amp and with bfloat16. +# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py +from functools import partial +from typing import Optional + +# import fused_dense_cuda # from apex +import fused_dense_lib as fused_dense_cuda +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.cuda.amp import custom_bwd, custom_fwd +from torch.distributed import ProcessGroup + +from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd +from flash_attn.utils.distributed import ( + all_gather_raw, + all_reduce, + all_reduce_raw, + reduce_scatter, + reduce_scatter_raw, +) + + +class FusedDenseFunc(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True + ): + """ + If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel + with sequence parallelism: we do an all_gather_raw of x before doing the matmul. + """ + ctx.compute_weight_gradient = weight.requires_grad + ctx.return_residual = return_residual + ctx.process_group = process_group + ctx.sequence_parallel = sequence_parallel + + if torch.is_autocast_enabled(): + x = x.to(dtype=torch.get_autocast_gpu_dtype()) + x = x.contiguous() + if process_group is not None and sequence_parallel: + # We want to kick off the all_gather early, before weight dtype conversion + total_x, handle_x = all_gather_raw(x, process_group, async_op=True) + else: + total_x = x + + if torch.is_autocast_enabled(): + weight = weight.to(dtype=torch.get_autocast_gpu_dtype()) + bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None + weight = weight.contiguous() + if process_group is not None and sequence_parallel: + handle_x.wait() + batch_shape, n = total_x.shape[:-1], total_x.shape[-1] + batch_dim = batch_shape.numel() + # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 + if min(batch_dim, n, *weight.shape) > 65535 * 32: + raise RuntimeError("fused_dense only supports matrix dims <= 2M") + output = F.linear(total_x, weight, bias) + if ctx.compute_weight_gradient: + ctx.save_for_backward(x, weight) + else: + ctx.save_for_backward(weight) + return output if not return_residual else (output, x) + + @staticmethod + @custom_bwd + def backward(ctx, grad_output, *args): + grad_output = grad_output.contiguous() + if ctx.return_residual: + (grad_input,) = args + grad_input = grad_input.contiguous() + process_group = ctx.process_group + sequence_parallel = ctx.sequence_parallel + if ctx.compute_weight_gradient: + x, weight = ctx.saved_tensors + if process_group is not None and sequence_parallel: + total_x, handle_x = all_gather_raw(x, process_group, async_op=True) + else: + total_x = x + else: + (weight,) = ctx.saved_tensors + total_x = None + batch_shape = grad_output.shape[:-1] + batch_dim = batch_shape.numel() + grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) + if ctx.needs_input_grad[0]: + if not ctx.return_residual: + grad_input = F.linear(grad_output, weight.t()) + else: + grad_input = torch.addmm( + grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight + ) + grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) + if process_group is not None: + reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw + grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) + else: + grad_input = None + if ctx.needs_input_grad[1]: + assert ctx.compute_weight_gradient + if process_group is not None and sequence_parallel: + handle_x.wait() + grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad( + total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2] + ) + else: + grad_weight = None + grad_bias = grad_output if ctx.needs_input_grad[2] else None + if process_group is not None and ctx.needs_input_grad[0]: + handle_grad_input.wait() + return grad_input, grad_weight, grad_bias, None, None, None + + +def fused_dense_func( + x: Tensor, + weight: Tensor, + bias: Optional[Tensor] = None, + return_residual: bool = False, + process_group: Optional[ProcessGroup] = None, + sequence_parallel: bool = True, +): + dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( + x.dtype == torch.float32 and torch.is_autocast_enabled() + ) + if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible: + return FusedDenseFunc.apply( + x, weight, bias, return_residual, process_group, sequence_parallel + ) + else: + assert process_group is None + out = F.linear(x, weight, bias) + return out if not return_residual else (out, x) + + +class FusedDense(nn.Linear): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + return_residual: bool = False, + device=None, + dtype=None, + ) -> None: + super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) + self.return_residual = return_residual + + def forward(self, x, process_group=None): + """ + If process_group is not None, we're doing Tensor Parallel with sequence parallelism: + we do an all_gather of x before doing the matmul. + """ + return fused_dense_func( + x, + self.weight, + self.bias, + return_residual=self.return_residual, + process_group=process_group, + ) + + +class ColumnParallelLinear(nn.Linear): + def __init__( + self, + in_features: int, + out_features: int, + process_group: ProcessGroup, + bias: bool = True, + sequence_parallel=True, + multiple_of=1, + device=None, + dtype=None, + ) -> None: + world_size = torch.distributed.get_world_size(process_group) + if out_features % multiple_of: + raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}") + multiple = out_features // multiple_of + # We want to split @multiple across world_size, but it could be an uneven split + div = multiple // world_size + mod = multiple % world_size + # The first @mod ranks get @div + 1 copies, the rest get @div copies + local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) + super().__init__( + in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype + ) + self.process_group = process_group + self.sequence_parallel = sequence_parallel + + def forward(self, x): + # If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: + # we do an all_gather of x before doing the matmul. + # If not, then the input is already gathered. + return fused_dense_func( + x, + self.weight, + self.bias, + process_group=self.process_group, + sequence_parallel=self.sequence_parallel, + ) + + +class RowParallelLinear(nn.Linear): + def __init__( + self, + in_features: int, + out_features: int, + process_group: ProcessGroup, + bias: bool = True, + sequence_parallel=True, + multiple_of=1, + device=None, + dtype=None, + ) -> None: + world_size = torch.distributed.get_world_size(process_group) + rank = torch.distributed.get_rank(process_group) + if in_features % multiple_of: + raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}") + multiple = in_features // multiple_of + # We want to split @multiple across world_size, but it could be an uneven split + div = multiple // world_size + mod = multiple % world_size + # The first @mod ranks get @div + 1 copies, the rest get @div copies + local_multiple = div + int(torch.distributed.get_rank(process_group) < mod) + # Only rank 0 will have bias + super().__init__( + local_multiple * multiple_of, + out_features, + bias=bias and rank == 0, + device=device, + dtype=dtype, + ) + self.process_group = process_group + self.sequence_parallel = sequence_parallel + + def forward(self, x): + """ + We're doing Tensor Parallel with sequence parallelism: we do the matmul and then + a reduce_scatter of the result. + """ + out = fused_dense_func(x, self.weight, self.bias) + reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce + return reduce_fn(out, self.process_group) + + +class FusedMLPFunc(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, + x, + weight1, + bias1, + weight2, + bias2, + activation="gelu_approx", + save_pre_act=True, + return_residual=False, + checkpoint_lvl=0, + heuristic=0, + process_group=None, + sequence_parallel=True, + ): + """ + If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel + with sequence parallelism: we do an all_gather of x before doing the matmul. + If sequence_parallel=False, then the input is already gathered. + + checkpoint_lvl: + 0: no recomputation in the bwd + 1: recompute gelu_out / relu_out in the bwd + 2: recompute pre_act and gelu_out / relu_out in the bwd + """ + assert -1 <= heuristic <= 4 + assert activation in ["gelu_approx", "relu", "sqrelu"] + if activation == "sqrelu": + assert heuristic == -1 + if not save_pre_act: + checkpoint_lvl = 2 + assert checkpoint_lvl in [0, 1, 2] + ctx.return_residual = return_residual + ctx.process_group = process_group + ctx.sequence_parallel = sequence_parallel + ctx.checkpoint_lvl = checkpoint_lvl + ctx.activation = activation + ctx.heuristic = heuristic + + if torch.is_autocast_enabled(): + x = x.to(dtype=torch.get_autocast_gpu_dtype()) + x = x.contiguous() + if process_group is not None and sequence_parallel: + # We want to kick off the all_gather early, before weight dtype conversion + total_x, handle_x = all_gather_raw(x, process_group, async_op=True) + else: + total_x = x + + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_gpu_dtype() + weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]] + bias1 = bias1.to(dtype=dtype) if bias1 is not None else None + bias2 = bias2.to(dtype=dtype) if bias2 is not None else None + weight1 = weight1.contiguous() + bias1 = bias1.contiguous() if bias1 is not None else None + weight2 = weight2.contiguous() + bias2 = bias2.contiguous() if bias2 is not None else None + if process_group is not None and sequence_parallel: + handle_x.wait() + batch_shape, n = total_x.shape[:-1], total_x.shape[-1] + batch_dim = batch_shape.numel() + # https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174 + if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32: + raise RuntimeError("fused_dense only supports matrix dims <= 2M") + if heuristic == -1: + pre_act = F.linear(total_x, weight1, bias1) + activation_fn = ( + partial(F.gelu, approximate="tanh") + if activation == "gelu_approx" + else (sqrelu_fwd if activation == "sqrelu" else F.relu) + ) + with torch.jit.fuser("fuser2"): + output1 = activation_fn(pre_act) + # This is before adding bias1 + # pre_act = F.linear(total_x.reshape(batch_dim, n), weight1) + # with torch.jit.fuser('fuser2'): + # output1 = bias_gelu(pre_act, bias1) + else: + is_gelu = activation == "gelu_approx" + output1, *rest = fused_dense_cuda.linear_act_forward( + total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic + ) + if save_pre_act: + pre_act = rest[0] + output2 = F.linear(output1, weight2, bias2) + if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): + # For RELU the pre_act is very small (just a bit-mask) so we just save it + ctx.save_for_backward(x, weight1, weight2, pre_act, output1) + elif checkpoint_lvl == 1: + ctx.save_for_backward(x, weight1, weight2, pre_act) + elif checkpoint_lvl == 2: + ctx.save_for_backward(x, weight1, weight2, bias1) + output2 = output2.reshape(*batch_shape, output2.shape[-1]) + return output2 if not return_residual else (output2, x) + + @staticmethod + @custom_bwd + def backward(ctx, grad_output, *args): + grad_output = grad_output.contiguous() + checkpoint_lvl = ctx.checkpoint_lvl + activation = ctx.activation + activation_fn = ( + partial(F.gelu, approximate="tanh") + if activation == "gelu_approx" + else (sqrelu_fwd if activation == "sqrelu" else F.relu) + ) + if ctx.return_residual: + (grad_input,) = args + grad_input = grad_input.contiguous() + process_group = ctx.process_group + sequence_parallel = ctx.sequence_parallel + x, weight1, weight2, *rest = ctx.saved_tensors + if process_group is None or not sequence_parallel: + total_x = x + batch_shape = grad_output.shape[:-1] + batch_dim = batch_shape.numel() + if checkpoint_lvl in [0, 1]: + if process_group is not None and sequence_parallel: + total_x, handle_x = all_gather_raw(x, process_group, async_op=True) + if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"): + pre_act, output1 = rest + elif checkpoint_lvl == 1: + (pre_act,) = rest + with torch.jit.fuser("fuser2"): + output1 = activation_fn(pre_act) + elif checkpoint_lvl == 2: + (bias1,) = rest + if process_group is not None and sequence_parallel: + total_x, _ = all_gather_raw(x, process_group) + if ctx.heuristic == -1: + pre_act = F.linear(total_x, weight1, bias1) + with torch.jit.fuser("fuser2"): + output1 = activation_fn(pre_act) + else: + output1, pre_act = fused_dense_cuda.linear_act_forward( + total_x.reshape(batch_dim, total_x.shape[-1]), + weight1, + bias1, + activation == "gelu_approx", + True, + ctx.heuristic, + ) + + grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) + output1 = output1.reshape(batch_dim, output1.shape[-1]) + pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1]) + if ctx.needs_input_grad[3]: + grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad( + output1, grad_output, ctx.needs_input_grad[4] + ) + else: + grad_weight2 = None + grad_bias2 = grad_output if ctx.needs_input_grad[4] else None + if ctx.heuristic == -1: + # grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act) + grad_output1 = F.linear(grad_output, weight2.t()) + activation_grad_fn = ( + gelu_bwd + if activation == "gelu_approx" + else (sqrelu_bwd if activation == "sqrelu" else relu_bwd) + ) + with torch.jit.fuser("fuser2"): + grad_pre_act = activation_grad_fn(grad_output1, pre_act) + else: + # The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't + # just compute gelu/relu grad + grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad( + weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic + ) + if not ctx.needs_input_grad[2]: + grad_bias1 = None + if ctx.needs_input_grad[0]: + if not ctx.return_residual: + grad_input = F.linear(grad_pre_act, weight1.t()) + else: + grad_input = torch.addmm( + grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1 + ) + grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1]) + if process_group is not None: + reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw + grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True) + else: + grad_input = None + if ctx.heuristic == -1: + if ctx.needs_input_grad[1]: + if process_group is not None and sequence_parallel and checkpoint_lvl != 2: + handle_x.wait() + grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad( + total_x.reshape(batch_dim, total_x.shape[-1]), + grad_pre_act, + ctx.needs_input_grad[2], + ) + else: + grad_weight1 = None + grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None + else: + if ctx.needs_input_grad[1]: + if process_group is not None and sequence_parallel and checkpoint_lvl != 2: + handle_x.wait() + grad_weight1 = F.linear( + grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t() + ) + else: + grad_weight1 = None + if process_group is not None and ctx.needs_input_grad[0]: + handle_grad_input.wait() + return ( + grad_input, + grad_weight1, + grad_bias1, + grad_weight2, + grad_bias2, + None, + None, + None, + None, + None, + None, + None, + ) + + +def fused_mlp_func( + x: Tensor, + weight1: Tensor, + weight2: Tensor, + bias1: Optional[Tensor] = None, + bias2: Optional[Tensor] = None, + activation: str = "gelu_approx", + save_pre_act: bool = True, + return_residual: bool = False, + checkpoint_lvl: int = 0, + heuristic: int = 0, + process_group: Optional[ProcessGroup] = None, + sequence_parallel: bool = True, +): + assert activation in ["gelu_approx", "relu", "sqrelu"] + dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or ( + x.dtype == torch.float32 and torch.is_autocast_enabled() + ) + # If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu) + dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0) + if ( + x.is_cuda + and weight1.is_cuda + and weight2.is_cuda + and (bias1 is None or bias1.is_cuda) + and (bias2 is None or bias2.is_cuda) + and dtype_eligible + and dim_eligible + ): + return FusedMLPFunc.apply( + x, + weight1, + bias1, + weight2, + bias2, + activation, + save_pre_act, + return_residual, + checkpoint_lvl, + heuristic, + process_group, + sequence_parallel, + ) + else: + assert process_group is None + pre_act = F.linear(x, weight1, bias1) + activation_fn = ( + partial(F.gelu, approximate="tanh") + if activation == "gelu_approx" + else partial(F.relu, inplace=True) + ) + output1 = activation_fn(pre_act) + output2 = F.linear(output1, weight2, bias2) + return output2 if not return_residual else (output2, x) + + +class FusedMLP(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + bias1=True, + bias2=True, + activation="gelu_approx", + return_residual=False, + checkpoint_lvl=0, + heuristic="auto", + device=None, + dtype=None, + ): + """ + If process_group is not None, we're doing Tensor Parallel with sequence parallelism: + we do an all_gather of x before doing the matmul, gelu, then matmul. + Finally we do a reduce_scatter of the output. + + checkpoint_lvl (increasing lvl means slower but more memory saving): + 0: no recomputation in the bwd + 1: recompute gelu_out in the bwd + 2: recompute pre_act and gelu_out in the bwd + heuristic: + -1: don't fuse gemm + gelu (separate kernel) + 0..4: use this heuristic for the algo section in the fused gemm + gelu + 'auto': heuristic will be picked automatically: + For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. + For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. + For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation + is slower than the unfused version. + return_residual: whether to return the input x along with the output. This is for + performance reason: for post-norm architecture, returning the input allows us + to fuse the backward of nn.Linear with the residual connection. + """ + assert checkpoint_lvl in [0, 1, 2] + assert activation in ["gelu_approx", "relu", "sqrelu"] + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features * 4 + self.activation = activation + self.return_residual = return_residual + self.checkpoint_lvl = checkpoint_lvl + self.heuristic = heuristic if activation != "sqrelu" else -1 + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) + + def forward(self, x, process_group=None): + dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() + if self.heuristic == "auto": + if self.activation == "gelu_approx": + if torch.cuda.get_device_capability("cuda") == (9, 0): + heuristic = -1 + else: + cuda_ver = tuple(map(int, torch.version.cuda.split("."))) + heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) + else: + heuristic = 0 + else: + heuristic = self.heuristic + out = fused_mlp_func( + x, + self.fc1.weight, + self.fc2.weight, + self.fc1.bias, + self.fc2.bias, + activation=self.activation, + save_pre_act=self.training, + return_residual=self.return_residual, + checkpoint_lvl=self.checkpoint_lvl, + heuristic=heuristic, + process_group=process_group, + ) + if self.return_residual: + out, x = out + if process_group is not None: + out = reduce_scatter(out, process_group) + return out if not self.return_residual else (out, x) + + +class ParallelFusedMLP(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + activation="gelu_approx", + process_group: ProcessGroup = None, + bias1=True, + bias2=True, + sequence_parallel=True, + checkpoint_lvl=0, + heuristic="auto", + device=None, + dtype=None, + ): + """ + process_group is required. We're doing Tensor Parallel with sequence parallelism: + we do an all_gather of x before doing the matmul, gelu, then matmul. + Finally we do a reduce_scatter of the output. + + checkpoint_lvl (increasing lvl means slower but more memory saving): + 0: no recomputation in the bwd + 1: recompute gelu_out in the bwd + 2: recompute pre_act and gelu_out in the bwd + heuristic: + -1: don't fuse gemm + gelu (separate kernel) + 0..4: use this heuristic for the algo section in the fused gemm + gelu + 'auto': heuristic will be picked automatically: + For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf. + For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16. + """ + assert checkpoint_lvl in [0, 1, 2] + assert activation in ["gelu_approx", "relu", "sqrelu"] + assert process_group is not None + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features * 4 + self.activation = activation + self.process_group = process_group + self.sequence_parallel = sequence_parallel + self.checkpoint_lvl = checkpoint_lvl + self.heuristic = heuristic if activation != "sqrelu" else -1 + self.fc1 = ColumnParallelLinear( + in_features, hidden_features, process_group, bias=bias1, **factory_kwargs + ) + self.fc2 = RowParallelLinear( + hidden_features, out_features, process_group, bias=bias2, **factory_kwargs + ) + + def forward(self, x): + dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype() + if self.heuristic == "auto": + if self.activation == "gelu_approx": + cuda_ver = tuple(map(int, torch.version.cuda.split("."))) + heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1) + else: + heuristic = 0 + else: + heuristic = self.heuristic + out = fused_mlp_func( + x, + self.fc1.weight, + self.fc2.weight, + self.fc1.bias, + self.fc2.bias, + activation=self.activation, + save_pre_act=self.training, + checkpoint_lvl=self.checkpoint_lvl, + heuristic=heuristic, + process_group=self.process_group, + sequence_parallel=self.sequence_parallel, + ) + reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce + return reduce_fn(out, self.process_group) diff --git a/lib/python3.12/site-packages/flash_attn/ops/layer_norm.py b/lib/python3.12/site-packages/flash_attn/ops/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..4b6cd798fd02844ef9cd3897f8ab95e490e638bf --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/layer_norm.py @@ -0,0 +1,800 @@ +# Copyright (c) 2022, Tri Dao. +# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py + +import dropout_layer_norm +import torch +from torch.nn import init + + +def maybe_align(x, alignment_in_bytes=16): + """Assume that x already has last dim divisible by alignment_in_bytes""" + # TD [2023-07-04] I'm not 100% sure that clone will align the memory + # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440 + return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone() + + +def _dropout_add_layer_norm_forward( + x0, + residual, + gamma, + beta, + rowscale, + colscale, + dropout_p, + epsilon, + residual_in_fp32=False, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes""" + hidden_size = gamma.numel() + x0mat = x0.view((-1, hidden_size)) + residualmat = residual.view((-1, hidden_size)) if residual is not None else None + rowscale = rowscale.view(-1) if rowscale is not None else None + zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( + x0mat, + residualmat, + gamma, + beta, + rowscale, + colscale, + None, + None, + dropout_p, + epsilon, + 1.0, + 0, + None, + residual_in_fp32, + is_rms_norm, + ) + # dmask is None if dropout_p == 0.0 + # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype + return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma + + +def _dropout_add_layer_norm_backward( + dz, + dx, + x, + x0, + dmask, + mu, + rsigma, + gamma, + rowscale, + colscale, + dropout_p, + has_residual, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes + dx == None means that it was a post-norm architecture + (x = drop(x0) + residual was not returned in the fwd). + x0 must not be None if we have colscale. + """ + hidden_size = gamma.numel() + xmat = x.view((-1, hidden_size)) + dzmat = dz.view(xmat.shape) + dxmat = dx.view(xmat.shape) if dx is not None else None + x0mat = x0.view((-1, hidden_size)) if x0 is not None else None + rowscale = rowscale.view(-1) if rowscale is not None else None + if colscale is not None: + assert x0 is not None, "x0 is required to compute the gradient of colscale" + dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( + dzmat, + dxmat, + xmat, + x0mat, + dmask, + mu, + rsigma, + gamma, + rowscale, + colscale, + None, + None, + dropout_p, + 1.0, + 0, + has_residual, + is_rms_norm, + ) + # dresidualmat is None if not has_residual + if colscale is None: + return dx0mat, dresidualmat, dgamma, dbeta + else: + dcolscale = rest[0] + return dx0mat, dresidualmat, dgamma, dbeta, dcolscale + + +def _dropout_add_layer_norm_subset_forward( + x0, + residual, + gamma, + beta, + colscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + residual_in_fp32=False, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes""" + hidden_size = gamma.numel() + x0mat = x0.view((-1, hidden_size)) + residualmat = residual.view((-1, hidden_size)) if residual is not None else None + x0_subset = x0_subset.view(-1) if x0_subset is not None else None + out_subset = out_subset.view(-1) if out_subset is not None else None + zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd( + x0mat, + residualmat, + gamma, + beta, + None, + colscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + None, + residual_in_fp32, + is_rms_norm, + ) + # dmask is None if dropout_p == 0.0 + # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype + return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma + + +def _dropout_add_layer_norm_subset_backward( + dz, + dx, + x, + x0, + dmask, + mu, + rsigma, + gamma, + colscale, + x0_subset, + out_subset, + dropout_p, + rowscale_const, + x0_numrows, + has_residual, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes + dx == None means that it was a post-norm architecture + (x = drop(x0) + residual was not returned in the fwd). + x0 must not be None if we have colscale. + """ + hidden_size = gamma.numel() + xmat = x.view((-1, hidden_size)) + dzmat = dz.view(-1, hidden_size) + dxmat = dx.view(xmat.shape) if dx is not None else None + x0mat = x0.view((-1, hidden_size)) if x0 is not None else None + x0_subset = x0_subset.view(-1) if x0_subset is not None else None + out_subset = out_subset.view(-1) if out_subset is not None else None + if colscale is not None: + assert x0 is not None, "x0 is required to compute the gradient of colscale" + dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd( + dzmat, + dxmat, + xmat, + x0mat, + dmask, + mu, + rsigma, + gamma, + None, + colscale, + x0_subset, + out_subset, + dropout_p, + rowscale_const, + x0_numrows, + has_residual, + is_rms_norm, + ) + # dresidualmat is None if not has_residual + if colscale is None: + return dx0mat, dresidualmat, dgamma, dbeta + else: + dcolscale = rest[0] + return dx0mat, dresidualmat, dgamma, dbeta, dcolscale + + +def _dropout_add_layer_norm_parallel_residual_forward( + x0, + x1, + residual, + gamma0, + beta0, + gamma1, + beta1, + dropout_p, + epsilon, + residual_in_fp32=False, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes""" + hidden_size = gamma0.numel() + x0mat = x0.view((-1, hidden_size)) + x1mat = x1.view((-1, hidden_size)) if x1 is not None else None + residualmat = residual.view((-1, hidden_size)) if residual is not None else None + ( + z0mat, + z1mat, + xmat, + dmask0, + dmask1, + mu, + rsigma, + ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd( + x0mat, + x1mat, + residualmat, + gamma0, + beta0, + gamma1, + beta1, + dropout_p, + epsilon, + None, + residual_in_fp32, + is_rms_norm, + ) + # dmask0 and dmask1 are None if dropout_p == 0.0 + # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype + return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma + + +def _dropout_add_layer_norm_parallel_residual_backward( + dz0, + dz1, + dx, + x, + dmask0, + dmask1, + mu, + rsigma, + gamma0, + gamma1, + dropout_p, + has_x1, + has_residual, + is_rms_norm=False, +): + """Assume that arguments are contiguous and aligned to 16 bytes + dx == None means that it was a post-norm architecture + (x = drop(x0) + residual was not returned in the fwd). + """ + hidden_size = gamma0.numel() + xmat = x.view((-1, hidden_size)) + dz0mat = dz0.view(xmat.shape) + dz1mat = dz1.view(xmat.shape) if dz1 is not None else None + dxmat = dx.view(xmat.shape) if dx is not None else None + ( + dx0mat, + dx1mat, + dresidualmat, + dgamma0, + dbeta0, + dgamma1, + dbeta1, + *rest, + ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd( + dz0mat, + dz1mat, + dxmat, + xmat, + dmask0, + dmask1, + mu, + rsigma, + gamma0, + gamma1, + dropout_p, + has_x1, + has_residual, + is_rms_norm, + ) + # dresidualmat is None if not has_residual + return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1 + + +class DropoutAddLayerNormFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x0, + residual, + gamma, + beta, + rowscale, + colscale, + dropout_p, + epsilon, + residual_in_fp32=False, + prenorm=False, + is_rms_norm=False, + return_dmask=False, + ): + x0 = maybe_align(x0.contiguous(), 16) + residual = maybe_align(residual.contiguous(), 16) if residual is not None else None + gamma = maybe_align(gamma.contiguous(), 16) + beta = maybe_align(beta.contiguous(), 16) if beta is not None else None + rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None + colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None + zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward( + x0, + residual, + gamma, + beta, + rowscale, + colscale, + dropout_p, + epsilon, + residual_in_fp32, + is_rms_norm, + ) + # Only need to save x0 if we need to compute gradient wrt colscale + x0_saved = x0 if colscale is not None else None + ctx.save_for_backward( + xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale + ) + ctx.prenorm = prenorm + ctx.dropout_p = dropout_p + ctx.has_residual = residual is not None + ctx.is_rms_norm = is_rms_norm + ctx.has_beta = beta is not None + if not return_dmask: + return ( + zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape)) + ) + else: + dmask = ( + dmask.view(x0.shape) + if dropout_p > 0.0 + else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) + ) + ctx.mark_non_differentiable(dmask) + return ( + (zmat.view(x0.shape), dmask) + if not prenorm + else (zmat.view(x0.shape), xmat.view(x0.shape), dmask) + ) + + @staticmethod + def backward(ctx, dz, *args): + # assert dz.is_contiguous() + dz = maybe_align(dz.contiguous(), 16) # this happens! + dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None + x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors + # x0 is None if colscale is None + dropout_p = ctx.dropout_p + has_residual = ctx.has_residual + dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward( + dz, + dx, + x, + x0, + dmask, + mu, + rsigma, + gamma, + rowscale, + colscale, + dropout_p, + has_residual, + ctx.is_rms_norm, + ) + dx0 = dx0mat.view(x.shape) + dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None + dcolscale = rest[0] if colscale is not None else None + return ( + dx0, + dresidual, + dgamma, + dbeta if ctx.has_beta else None, + None, + dcolscale, + None, + None, + None, + None, + None, + None, + ) + + +class DropoutAddLayerNormSubsetFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x0, + residual, + gamma, + beta, + colscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + residual_in_fp32=False, + prenorm=False, + is_rms_norm=False, + return_dmask=False, + ): + x0 = maybe_align(x0.contiguous(), 16) + residual = maybe_align(residual.contiguous(), 16) if residual is not None else None + gamma = maybe_align(gamma.contiguous(), 16) + beta = maybe_align(beta.contiguous(), 16) if beta is not None else None + colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None + zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward( + x0, + residual, + gamma, + beta, + colscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + residual_in_fp32, + is_rms_norm, + ) + # Only need to save x0 if we need to compute gradient wrt colscale + x0_saved = x0 if colscale is not None else None + x_shape = (-1, *x0.shape[1:]) + ctx.save_for_backward( + xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset + ) + ctx.prenorm = prenorm + ctx.dropout_p = dropout_p + ctx.rowscale_const = rowscale_const + ctx.x0_numrows = x0.shape[:-1].numel() + ctx.has_residual = residual is not None + ctx.is_rms_norm = is_rms_norm + ctx.has_beta = beta is not None + z_shape = (-1, *x0.shape[1:]) + if not return_dmask: + return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape)) + else: + z = zmat.view(z_shape) + dmask = ( + dmask.view(x0.shape) + if dropout_p > 0.0 + else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) + ) + ctx.mark_non_differentiable(dmask) + return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask) + + @staticmethod + def backward(ctx, dz, *args): + # assert dz.is_contiguous() + dz = maybe_align(dz.contiguous(), 16) # this happens! + dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None + x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors + # x0 is None if colscale is None + dropout_p = ctx.dropout_p + has_residual = ctx.has_residual + dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward( + dz, + dx, + x, + x0, + dmask, + mu, + rsigma, + gamma, + colscale, + x0_subset, + out_subset, + dropout_p, + ctx.rowscale_const, + ctx.x0_numrows, + has_residual, + ctx.is_rms_norm, + ) + dx0 = dx0mat.view(-1, *x.shape[1:]) + dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None + dcolscale = rest[0] if colscale is not None else None + return ( + dx0, + dresidual, + dgamma, + dbeta if ctx.has_beta else None, + dcolscale, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x0, + x1, + residual, + gamma0, + beta0, + gamma1, + beta1, + dropout_p, + epsilon, + residual_in_fp32=False, + prenorm=False, + is_rms_norm=False, + return_dmask=False, + ): + x0 = maybe_align(x0.contiguous(), 16) + x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None + residual = maybe_align(residual.contiguous(), 16) if residual is not None else None + gamma0 = maybe_align(gamma0.contiguous(), 16) + beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None + gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None + beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None + ( + z0mat, + z1mat, + xmat, + dmask0, + dmask1, + mu, + rsigma, + ) = _dropout_add_layer_norm_parallel_residual_forward( + x0, + x1, + residual, + gamma0, + beta0, + gamma1, + beta1, + dropout_p, + epsilon, + residual_in_fp32, + is_rms_norm, + ) + ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma) + ctx.prenorm = prenorm + ctx.dropout_p = dropout_p + ctx.has_x1 = x1 is not None + ctx.has_residual = residual is not None + ctx.is_rms_norm = is_rms_norm + ctx.has_beta = beta0 is not None + z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None) + if not return_dmask: + return z if not prenorm else (*z, xmat.view(x0.shape)) + else: + dmask0 = ( + dmask0.view(x0.shape) + if dropout_p > 0.0 + else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) + ) + dmask1 = ( + dmask1.view(x0.shape) + if dropout_p > 0.0 and x1 is not None + else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device) + ) + ctx.mark_non_differentiable(dmask0) + ctx.mark_non_differentiable(dmask1) + return ( + (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1) + ) + + @staticmethod + def backward(ctx, dz0, dz1, *args): + dz0 = maybe_align(dz0.contiguous(), 16) # this happens! + dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None + dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None + x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors + dropout_p = ctx.dropout_p + has_x1 = ctx.has_x1 + has_residual = ctx.has_residual + ( + dx0mat, + dx1mat, + dresidualmat, + dgamma0, + dbeta0, + dgamma1, + dbeta1, + ) = _dropout_add_layer_norm_parallel_residual_backward( + dz0, + dz1, + dx, + x, + dmask0, + dmask1, + mu, + rsigma, + gamma0, + gamma1, + dropout_p, + has_x1, + has_residual, + ctx.is_rms_norm, + ) + dx0 = dx0mat.view(x.shape) + dx1 = dx1mat.view(x.shape) if dx1mat is not None else None + dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None + return ( + dx0, + dx1, + dresidual, + dgamma0, + dbeta0 if ctx.has_beta else None, + dgamma1, + dbeta1 if ctx.has_beta else None, + None, + None, + None, + None, + None, + None, + ) + + +def layer_norm(x, weight, bias, epsilon): + return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False) + + +def dropout_add_layer_norm( + x0, + residual, + weight, + bias, + dropout_p, + epsilon, + rowscale=None, + layerscale=None, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormFn.apply( + x0, + residual, + weight, + bias, + rowscale, + layerscale, + dropout_p, + epsilon, + residual_in_fp32, + prenorm, + False, + return_dropout_mask, + ) + + +def dropout_add_layer_norm_subset( + x0, + residual, + weight, + bias, + dropout_p, + epsilon, + layerscale=None, + x0_subset=None, + out_subset=None, + rowscale_const=1.0, + out_numrows=0, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormSubsetFn.apply( + x0, + residual, + weight, + bias, + layerscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + residual_in_fp32, + prenorm, + False, + return_dropout_mask, + ) + + +def dropout_add_layer_norm_parallel_residual( + x0, + x1, + residual, + weight0, + bias0, + weight1, + bias1, + dropout_p, + epsilon, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormParallelResidualFn.apply( + x0, + x1, + residual, + weight0, + bias0, + weight1, + bias1, + dropout_p, + epsilon, + residual_in_fp32, + prenorm, + False, + return_dropout_mask, + ) + + +class DropoutAddLayerNorm(torch.nn.Module): + def __init__( + self, + hidden_size, + prenorm=False, + p=0.0, + eps=1e-5, + residual_in_fp32=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.prenorm = prenorm + self.p = p + self.eps = eps + self.residual_in_fp32 = residual_in_fp32 + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.reset_parameters() + + def reset_parameters(self): + init.ones_(self.weight) + init.zeros_(self.bias) + + def forward(self, x0, residual=None): + return dropout_add_layer_norm( + x0, + residual, + self.weight, + self.bias, + self.p if self.training else 0.0, + self.eps, + prenorm=self.prenorm, + residual_in_fp32=self.residual_in_fp32, + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/rms_norm.py b/lib/python3.12/site-packages/flash_attn/ops/rms_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..068348d61290e3839dd082b540d898578ba1e8e2 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/rms_norm.py @@ -0,0 +1,174 @@ +# Copyright (c) 2022, Tri Dao. +# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py + +import torch +from torch.nn import init + +from flash_attn.ops.layer_norm import ( + DropoutAddLayerNormFn, + DropoutAddLayerNormParallelResidualFn, + DropoutAddLayerNormSubsetFn, +) + + +def rms_norm(x, weight, epsilon): + return DropoutAddLayerNormFn.apply( + x, None, weight, None, None, None, 0.0, epsilon, False, False, True + ) + + +def dropout_add_rms_norm( + x0, + residual, + weight, + bias, + dropout_p, + epsilon, + rowscale=None, + layerscale=None, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormFn.apply( + x0, + residual, + weight, + bias, + rowscale, + layerscale, + dropout_p, + epsilon, + residual_in_fp32, + prenorm, + True, + return_dropout_mask, + ) + + +def dropout_add_rms_norm_subset( + x0, + residual, + weight, + bias, + dropout_p, + epsilon, + layerscale=None, + x0_subset=None, + out_subset=None, + rowscale_const=1.0, + out_numrows=0, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormSubsetFn.apply( + x0, + residual, + weight, + bias, + layerscale, + x0_subset, + out_subset, + dropout_p, + epsilon, + rowscale_const, + out_numrows, + residual_in_fp32, + prenorm, + True, + return_dropout_mask, + ) + + +def dropout_add_rms_norm_parallel_residual( + x0, + x1, + residual, + weight0, + bias0, + weight1, + bias1, + dropout_p, + epsilon, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, +): + """residual_in_fp32 only has an effect if residual is None. + Otherwise residual dtype is residual.dtype. + """ + return DropoutAddLayerNormParallelResidualFn.apply( + x0, + x1, + residual, + weight0, + bias0, + weight1, + bias1, + dropout_p, + epsilon, + residual_in_fp32, + prenorm, + True, + return_dropout_mask, + ) + + +class RMSNorm(torch.nn.Module): + def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.reset_parameters() + + def reset_parameters(self): + init.ones_(self.weight) + + def forward(self, x): + return rms_norm(x, self.weight, self.eps) + + +class DropoutAddRMSNorm(torch.nn.Module): + def __init__( + self, + hidden_size, + prenorm=False, + p=0.0, + eps=1e-5, + residual_in_fp32=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.prenorm = prenorm + self.p = p + self.eps = eps + self.residual_in_fp32 = residual_in_fp32 + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.reset_parameters() + + def reset_parameters(self): + init.ones_(self.weight) + + def forward(self, x0, residual=None): + return dropout_add_rms_norm( + x0, + residual, + self.weight, + None, + self.p if self.training else 0.0, + self.eps, + prenorm=self.prenorm, + residual_in_fp32=self.residual_in_fp32, + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__init__.py b/lib/python3.12/site-packages/flash_attn/ops/triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/__init__.py @@ -0,0 +1 @@ + diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43d3eefbb087d9562d0eabae64c84e848ca792f7 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/cross_entropy.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/cross_entropy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46c225c6056213697d757602145e2dcf3770ab21 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/cross_entropy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/k_activations.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/k_activations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf8094b755df5bda0d1315a9da3bfdc2e456e527 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/k_activations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/layer_norm.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/layer_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fa09e716b04ad1c23318f216331956c922de218 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/layer_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/linear.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..123cc13bb29f0556cdcd7668d3406e829495e0f8 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/mlp.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/mlp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35f34b8a3a31bf927f54f229ffcea6cf517b7870 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/mlp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/rotary.cpython-312.pyc b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/rotary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b717f1899568b4f094a0bc34961480122e76a33 Binary files /dev/null and b/lib/python3.12/site-packages/flash_attn/ops/triton/__pycache__/rotary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/cross_entropy.py b/lib/python3.12/site-packages/flash_attn/ops/triton/cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..7b0315b9793a8935dc507ed1b71b6e9d44e3b689 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/cross_entropy.py @@ -0,0 +1,329 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Tuple, Optional, Union + +import torch +import torch.nn.functional as F + +import triton +import triton.language as tl + +# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for +# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent +# version of PyTorch. The following 2 lines are for backward compatibility with +# older PyTorch. +if "all_gather_into_tensor" not in dir(torch.distributed): + torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base + + +@triton.heuristics( + { + "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, + } +) +@triton.jit +def cross_entropy_fwd_kernel( + loss_ptr, # data ptrs + lse_ptr, + z_loss_ptr, + logits_ptr, + labels_ptr, + smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + logits_row_stride, # strides + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, + # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE + SPLIT: tl.constexpr, + PRECOMPUTED_LSE: tl.constexpr, # If LSE is already computed (also no smoothing and logit_scale == 1.0) +): + row_idx = tl.program_id(0) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + sum_logits = 0.0 # For smoothing + if not PRECOMPUTED_LSE: + # Statistics for online softmax + m_i = -float("inf") + l_i = 0.0 + for col_offset in range(0, n_cols, BLOCK_SIZE): + cols = col_offset + tl.arange(0, BLOCK_SIZE) + logits = tl.load(logits_ptr + cols, mask=cols < n_cols, other=-float("inf")).to( + tl.float32 + ) * logit_scale + if HAS_SMOOTHING: + sum_logits += tl.sum(tl.where(cols < n_cols, logits, 0.0)) + m_i_new = tl.maximum(m_i, tl.max(logits)) + l_i = tl.exp(m_i - m_i_new) * l_i + tl.sum(tl.exp(logits - m_i_new)) + m_i = m_i_new + lse = tl.log(l_i) + m_i + tl.store(lse_ptr + row_idx, lse) + else: + lse = tl.load(lse_ptr + row_idx) + label_idx = tl.load(labels_ptr + row_idx) + if label_idx == ignore_index: + loss = 0.0 + z_loss = 0.0 + else: + label_idx -= class_start_idx + if label_idx >= 0 and label_idx < n_cols: + logits_label = tl.load(logits_ptr + label_idx) * logit_scale + if HAS_SMOOTHING: + loss = ( + (lse if not SPLIT else 0.0) + - smoothing * sum_logits / total_classes + - (1 - smoothing) * logits_label + ) + else: + loss = (lse if not SPLIT else 0.0) - logits_label + else: + # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss + if HAS_SMOOTHING: + loss = smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) + else: + loss = 0.0 + if not SPLIT: + z_loss = lse_square_scale * lse * lse + loss += z_loss + else: + z_loss = 0.0 + tl.store(loss_ptr + row_idx, loss) + if not SPLIT: + tl.store(z_loss_ptr + row_idx, z_loss) + + +@triton.heuristics( + { + "HAS_SMOOTHING": lambda args: args["smoothing"] > 0.0, + } +) +@triton.jit +def cross_entropy_bwd_kernel( + dlogits_ptr, # data ptrs + dloss_ptr, + logits_ptr, + lse_ptr, + labels_ptr, + smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + logits_row_stride, # strides + dlogits_row_stride, + dloss_row_stride, + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, +): + row_idx = tl.program_id(0) + col_block_idx = tl.program_id(1) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) + col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + label_idx = tl.load(labels_ptr + row_idx) + if label_idx != ignore_index: + dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) + else: + dloss = 0.0 + logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( + tl.float32 + ) * logit_scale + lse = tl.load(lse_ptr + row_idx) + probs = tl.exp(logits - lse) + probs += 2.0 * lse_square_scale * lse * probs + label_idx -= class_start_idx + if HAS_SMOOTHING: + smooth_positive = 1.0 - smoothing + smooth_negative = smoothing / total_classes + probs = tl.where(col_offsets == label_idx, probs - smooth_positive, probs) - smooth_negative + else: + probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) + tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) + + +class CrossEntropyLoss(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + logits, + labels, + precomputed_lse=None, + smoothing=0.0, + logit_scale=1.0, + lse_square_scale=0.0, + ignore_index=-100, + inplace_backward=False, + process_group=None, + ): + # For some reason Triton generates wrong code when labels has dtype long and its address + # is not aligned to 16 bytes. The ld.global.b64 seems to load the wrong label index. + if labels.dtype == torch.long and labels.data_ptr() % 16 != 0: + labels = F.pad(labels, (0, 1))[..., :-1] + assert labels.data_ptr() % 16 == 0 + n_rows, n_cols = logits.shape + assert labels.shape == (n_rows,) + world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) + total_classes = world_size * n_cols + rank = 0 if process_group is None else torch.distributed.get_rank(process_group) + class_start_idx = rank * n_cols + use_precomputed_lse = precomputed_lse is not None and logit_scale == 1.0 and smoothing == 0.0 + + if logits.stride(-1) != 1: + logits = logits.contiguous() + MAX_BLOCK_SIZE = 16 * 1024 + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) + num_warps = ( + 4 + if BLOCK_SIZE < 2048 + else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) + ) + losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) + if use_precomputed_lse: + assert precomputed_lse.shape == (n_rows,) + lse = precomputed_lse.contiguous() + else: + lse = torch.empty(n_rows, dtype=torch.float, device=logits.device) + z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) + # Need this, otherwise Triton tries to launch from cuda:0 and we get + # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) + with torch.cuda.device(logits.device.index): + cross_entropy_fwd_kernel[(n_rows,)]( + losses, # data ptrs + lse, + z_losses, + logits, + labels, + smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, + n_cols, # shapes + logits.stride(0), # strides + BLOCK_SIZE=BLOCK_SIZE, # constants + SPLIT=world_size > 1, + PRECOMPUTED_LSE=use_precomputed_lse, + num_warps=num_warps, + ) + + if world_size > 1: + # If there's no smoothing, if labels are in the vocab of this partition, losses contains + # - predicted logit, and 0 otherwise. + # If there's smoothing=0.1, for labels in the vocab of this partition, losses contains + # -0.9 * predicted logit - 0.1 * sum logit / total_classes. + # For labels not in the vocab of this partition, losses contains + # -0.1 * sum logit / total_classes. + if world_size > 1: + lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) + torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) + handle_losses = torch.distributed.all_reduce( + losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True + ) + lse = torch.logsumexp(lse_allgather, dim=0) + handle_losses.wait() + # After the allreduce, if there's no smoothing, the total losses are - predicted_logit, + # we just have to add the (global) lse. + # If there's smoothing=0.1, the total losses are + # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. + # Again, we just have to add the (global) lse. + losses += lse + if lse_square_scale != 0.0: + z_losses = lse_square_scale * lse.square() + z_losses.masked_fill_(labels == ignore_index, 0.0) + losses += z_losses + else: + z_losses = torch.zeros_like(losses) + losses.masked_fill_(labels == ignore_index, 0.0) + + ctx.save_for_backward(logits, lse, labels) + ctx.mark_non_differentiable(z_losses) + ctx.smoothing = smoothing + ctx.logit_scale = logit_scale + ctx.lse_square_scale = lse_square_scale + ctx.ignore_index = ignore_index + ctx.total_classes = total_classes + ctx.class_start_idx = class_start_idx + ctx.inplace_backward = inplace_backward + return losses, z_losses + + @staticmethod + def backward(ctx, grad_losses, grad_z_losses): + del grad_z_losses # z_losses are only for logging. + + logits, lse, labels = ctx.saved_tensors + dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) + n_rows, n_cols = logits.shape + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) + num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) + grid = lambda META: (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa + # Need this, otherwise Triton tries to launch from cuda:0 and we get + # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) + with torch.cuda.device(logits.device.index): + cross_entropy_bwd_kernel[grid]( + dlogits, # data ptrs + grad_losses, + logits, + lse, + labels, + ctx.smoothing, + ctx.logit_scale, + ctx.lse_square_scale, + ctx.ignore_index, + ctx.total_classes, + ctx.class_start_idx, + n_cols, # shapes + logits.stride(0), # strides + dlogits.stride(0), + grad_losses.stride(0), + BLOCK_SIZE=BLOCK_SIZE, # constants + num_warps=num_warps, + ) + return dlogits, None, None, None, None, None, None, None, None, None + + +def cross_entropy_loss( + logits: torch.Tensor, + labels: torch.Tensor, + precomputed_lse: Optional[torch.Tensor] = None, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + ignore_index=-100, + inplace_backward: bool = False, + process_group=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Arguments: + logits: (batch, vocab_size) + labels: (batch,) + label_smoothing: float + logit_scale: float. Multiply logits by this scale before calculating the loss. + lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. + This is also referred to as "z-loss". + ignore_index: int. If labels == ignore_index, the loss is set to 0.0. + inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. + This saves memory. + process_group: if not None, we're doing Tensor Parallel: each process is responsible for + one part of the vocab. The loss will be aggregated across processes. + Returns: + losses: (batch,), float + z_losses: (batch,), float + """ + return CrossEntropyLoss.apply( + logits, + labels, + precomputed_lse, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + inplace_backward, + process_group, + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/k_activations.py b/lib/python3.12/site-packages/flash_attn/ops/triton/k_activations.py new file mode 100644 index 0000000000000000000000000000000000000000..efb83c358eb4a85d069ee340a3c83f418f9a805b --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/k_activations.py @@ -0,0 +1,162 @@ +# Adapted from https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_activations.py +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +import math +from enum import Enum +from typing import Optional + +import triton +import triton.language as tl + +_sqrt2pi = math.sqrt(2.0 / math.pi) +_sqrt1_2 = math.sqrt(1.0 / 2) +_gaussian_pdf_normalization = 1.0 / math.sqrt(2 * math.pi) + + +class Activation(str, Enum): + SquaredReLU = "squared_relu" + GeLU = "gelu" + GeLUApprox = "gelu_approx" + LeakyReLU = "leaky_relu" + ReLU = "relu" + + +def get_triton_activation_kernel(activation: Optional[Activation]): + return ( + { + Activation.ReLU: relu, + Activation.LeakyReLU: leaky_relu, + Activation.GeLU: gelu, + Activation.GeLUApprox: gelu_approx, + Activation.SquaredReLU: squared_relu, + }[activation] + if activation + else None + ) + + +def get_triton_activation_bwd_kernel(activation: Optional[Activation]): + return ( + { + Activation.ReLU: relu_grad, + Activation.LeakyReLU: leaky_relu_grad, + Activation.GeLU: gelu_grad, + Activation.GeLUApprox: gelu_approx_grad, + Activation.SquaredReLU: squared_relu_grad, + }[activation] + if activation + else None + ) + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def cosh(x): + exp_x = tl.exp(x) + return (exp_x + 1.0 / exp_x) * 0.5 + + +# a Triton implementation of the most used activations +# See for instance http://arxiv.org/abs/1606.08415 for an overview + +# ReLU +@triton.jit +def relu(x): + """ + ReLU_ activation function + + .. _ReLU: https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html + """ + zero = 0.0 + return tl.where(x >= 0, x, zero.to(x.dtype)) + + +@triton.jit +def relu_grad(x): + # ReLU is different from other activations + # in that it does not require the input to retrospectively compute its gradient + # here the input is the downstream gradient, and we return the upstream gradient directly + zero = 0.0 + one = 1.0 + return tl.where(x >= 0, one.to(x.dtype), zero.to(x.dtype)) + + +@triton.jit +def squared_relu(x): + """ + Squared ReLU activation, as proposed in the Primer_ paper. + + .. _Primer: https://arxiv.org/abs/2109.08668 + """ + x_ = relu(x) + return (x_ * x_).to(x.dtype) + + +@triton.jit +def squared_relu_grad(x): + return tl.where(x >= 0, 2.0 * x, 0.0) + + +# Leaky ReLU +@triton.jit +def leaky_relu(x): + """ + LeakyReLU_ activation + + .. _LeakyReLU: https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html + """ + scale = 0.01 + 0.0 + scale = scale.to(x.dtype) + return tl.where(x >= 0, x, scale * x) + + +@triton.jit +def leaky_relu_grad(x): + min_grad = 0.01 + max_grad = 1 + + min_grad = min_grad.to(x.dtype) + max_grad = max_grad.to(x.dtype) + + return tl.where(x >= 0, max_grad, min_grad) + + +@triton.jit +def gelu(x): + """Gaussian Error Linear Unit (GELU)""" + return x * 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) + + +@triton.jit +def gelu_grad(x): + cdf = 0.5 * (1.0 + tl.libdevice.erf(x * _sqrt1_2)) + pdf = tl.exp(-0.5 * x * x) * _gaussian_pdf_normalization + return cdf + x * pdf + + +@triton.jit +def gelu_approx(x): + """ + GeLU_ activation - Gaussian error linear unit, with tanh approximation + + .. _GeLU: https://arxiv.org/pdf/1606.08415.pdf + """ + return 0.5 * x * (1.0 + tanh(_sqrt2pi * x * (1.0 + 0.044715 * x * x))) + + +@triton.jit +def gelu_approx_grad(x): + # CREDITS: Fast implementation proposed in + # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/fused_bias_gelu.py#L30 + tanh_out = tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + return 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/layer_norm.py b/lib/python3.12/site-packages/flash_attn/ops/triton/layer_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..addffe1f18585b59a8c22f20d3708ffb46d6bf34 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/layer_norm.py @@ -0,0 +1,1112 @@ +# Copyright (c) 2024, Tri Dao. +# Implement dropout + residual + layer_norm / rms_norm. + +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +import math + +import torch +import torch.nn.functional as F +from torch.cuda.amp import custom_fwd, custom_bwd + +import triton +import triton.language as tl + + +def layer_norm_ref( + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + dropout_mask=None, + dropout_mask1=None, + upcast=False, +): + dtype = x.dtype + if upcast: + x = x.float() + weight = weight.float() + bias = bias.float() if bias is not None else None + residual = residual.float() if residual is not None else residual + x1 = x1.float() if x1 is not None else None + weight1 = weight1.float() if weight1 is not None else None + bias1 = bias1.float() if bias1 is not None else None + if x1 is not None: + assert rowscale is None, "rowscale is not supported with parallel LayerNorm" + if rowscale is not None: + x = x * rowscale[..., None] + if dropout_p > 0.0: + if dropout_mask is not None: + x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) + else: + x = F.dropout(x, p=dropout_p) + if x1 is not None: + if dropout_mask1 is not None: + x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) + else: + x1 = F.dropout(x1, p=dropout_p) + if x1 is not None: + x = x + x1 + if residual is not None: + x = (x + residual).to(x.dtype) + out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( + dtype + ) + if weight1 is None: + return out if not prenorm else (out, x) + else: + out1 = F.layer_norm( + x.to(weight1.dtype), x.shape[-1:], weight=weight1, bias=bias1, eps=eps + ).to(dtype) + return (out, out1) if not prenorm else (out, out1, x) + + +def rms_norm_ref( + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + dropout_mask=None, + dropout_mask1=None, + upcast=False, +): + dtype = x.dtype + if upcast: + x = x.float() + weight = weight.float() + bias = bias.float() if bias is not None else None + residual = residual.float() if residual is not None else residual + x1 = x1.float() if x1 is not None else None + weight1 = weight1.float() if weight1 is not None else None + bias1 = bias1.float() if bias1 is not None else None + if x1 is not None: + assert rowscale is None, "rowscale is not supported with parallel LayerNorm" + if rowscale is not None: + x = x * rowscale[..., None] + if dropout_p > 0.0: + if dropout_mask is not None: + x = x.masked_fill(~dropout_mask, 0.0) / (1.0 - dropout_p) + else: + x = F.dropout(x, p=dropout_p) + if x1 is not None: + if dropout_mask1 is not None: + x1 = x1.masked_fill(~dropout_mask1, 0.0) / (1.0 - dropout_p) + else: + x1 = F.dropout(x1, p=dropout_p) + if x1 is not None: + x = x + x1 + if residual is not None: + x = (x + residual).to(x.dtype) + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = ((x * rstd * weight) + bias if bias is not None else (x * rstd * weight)).to(dtype) + if weight1 is None: + return out if not prenorm else (out, x) + else: + out1 = ((x * rstd * weight1) + bias1 if bias1 is not None else (x * rstd * weight1)).to( + dtype + ) + return (out, out1) if not prenorm else (out, out1, x) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + triton.Config({}, num_warps=16), + triton.Config({}, num_warps=32), + ], + key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS"], +) +# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) +# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) +@triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) +@triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) +@triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) +@triton.jit +def _layer_norm_fwd_1pass_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + RESIDUAL, # pointer to the residual + X1, + W1, + B1, + Y1, + RESIDUAL_OUT, # pointer to the residual + ROWSCALE, + SEEDS, # Dropout seeds for each row + DROPOUT_MASK, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_res_row, + stride_res_out_row, + stride_x1_row, + stride_y1_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + dropout_p, # Dropout probability + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_DROPOUT: tl.constexpr, + STORE_DROPOUT_MASK: tl.constexpr, + HAS_ROWSCALE: tl.constexpr, + HAS_X1: tl.constexpr, + HAS_W1: tl.constexpr, + HAS_B1: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + X += row * stride_x_row + Y += row * stride_y_row + if HAS_RESIDUAL: + RESIDUAL += row * stride_res_row + if STORE_RESIDUAL_OUT: + RESIDUAL_OUT += row * stride_res_out_row + if HAS_X1: + X1 += row * stride_x1_row + if HAS_W1: + Y1 += row * stride_y1_row + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_ROWSCALE: + rowscale = tl.load(ROWSCALE + row).to(tl.float32) + x *= rowscale + if HAS_DROPOUT: + # Compute dropout mask + # 7 rounds is good enough, and reduces register pressure + keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p + x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) + if STORE_DROPOUT_MASK: + tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) + if HAS_X1: + x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_ROWSCALE: + rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) + x1 *= rowscale + if HAS_DROPOUT: + # Compute dropout mask + # 7 rounds is good enough, and reduces register pressure + keep_mask = ( + tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p + ) + x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) + if STORE_DROPOUT_MASK: + tl.store(DROPOUT_MASK + (M + row) * N + cols, keep_mask, mask=cols < N) + x += x1 + if HAS_RESIDUAL: + residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) + x += residual + if STORE_RESIDUAL_OUT: + tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + # Write output + tl.store(Y + cols, y, mask=mask) + if HAS_W1: + w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) + if HAS_B1: + b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) + y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 + tl.store(Y1 + cols, y1, mask=mask) + + +def _layer_norm_fwd( + x, + weight, + bias, + eps, + residual=None, + x1=None, + weight1=None, + bias1=None, + dropout_p=0.0, + rowscale=None, + out_dtype=None, + residual_dtype=None, + is_rms_norm=False, + return_dropout_mask=False, + out=None, + residual_out=None +): + if residual is not None: + residual_dtype = residual.dtype + M, N = x.shape + assert x.stride(-1) == 1 + if residual is not None: + assert residual.stride(-1) == 1 + assert residual.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + if x1 is not None: + assert x1.shape == x.shape + assert rowscale is None + assert x1.stride(-1) == 1 + if weight1 is not None: + assert weight1.shape == (N,) + assert weight1.stride(-1) == 1 + if bias1 is not None: + assert bias1.shape == (N,) + assert bias1.stride(-1) == 1 + if rowscale is not None: + assert rowscale.is_contiguous() + assert rowscale.shape == (M,) + # allocate output + if out is None: + out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + else: + assert out.shape == x.shape + assert out.stride(-1) == 1 + if weight1 is not None: + y1 = torch.empty_like(out) + assert y1.stride(-1) == 1 + else: + y1 = None + if ( + residual is not None + or (residual_dtype is not None and residual_dtype != x.dtype) + or dropout_p > 0.0 + or rowscale is not None + or x1 is not None + ): + if residual_out is None: + residual_out = torch.empty( + M, N, device=x.device, dtype=residual_dtype if residual_dtype is not None else x.dtype + ) + else: + assert residual_out.shape == x.shape + assert residual_out.stride(-1) == 1 + else: + residual_out = None + mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + if dropout_p > 0.0: + seeds = torch.randint( + 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 + ) + else: + seeds = None + if return_dropout_mask and dropout_p > 0.0: + dropout_mask = torch.empty(M if x1 is None else 2 * M, N, device=x.device, dtype=torch.bool) + else: + dropout_mask = None + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + with torch.cuda.device(x.device.index): + _layer_norm_fwd_1pass_kernel[(M,)]( + x, + out, + weight, + bias, + residual, + x1, + weight1, + bias1, + y1, + residual_out, + rowscale, + seeds, + dropout_mask, + mean, + rstd, + x.stride(0), + out.stride(0), + residual.stride(0) if residual is not None else 0, + residual_out.stride(0) if residual_out is not None else 0, + x1.stride(0) if x1 is not None else 0, + y1.stride(0) if y1 is not None else 0, + M, + N, + eps, + dropout_p, + is_rms_norm, + BLOCK_N, + residual is not None, + residual_out is not None, + bias is not None, + dropout_p > 0.0, + dropout_mask is not None, + rowscale is not None, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 + if dropout_mask is not None and x1 is not None: + dropout_mask, dropout_mask1 = dropout_mask.tensor_split(2, dim=0) + else: + dropout_mask1 = None + return ( + out, + y1, + mean, + rstd, + residual_out if residual_out is not None else x, + seeds, + dropout_mask, + dropout_mask1, + ) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + triton.Config({}, num_warps=16), + triton.Config({}, num_warps=32), + ], + key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS", "HAS_DROPOUT"], +) +# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) +# @triton.heuristics({"HAS_DRESIDUAL": lambda args: args["DRESIDUAL"] is not None}) +# @triton.heuristics({"STORE_DRESIDUAL": lambda args: args["DRESIDUAL_IN"] is not None}) +@triton.heuristics({"HAS_ROWSCALE": lambda args: args["ROWSCALE"] is not None}) +@triton.heuristics({"HAS_DY1": lambda args: args["DY1"] is not None}) +@triton.heuristics({"HAS_DX1": lambda args: args["DX1"] is not None}) +@triton.heuristics({"HAS_B1": lambda args: args["DB1"] is not None}) +@triton.heuristics({"RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None}) +@triton.jit +def _layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DRESIDUAL, + W1, + DY1, + DX1, + DW1, + DB1, + DRESIDUAL_IN, + ROWSCALE, + SEEDS, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dres_row, + stride_dy1_row, + stride_dx1_row, + stride_dres_in_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + dropout_p, + rows_per_program, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_DROPOUT: tl.constexpr, + HAS_ROWSCALE: tl.constexpr, + HAS_DY1: tl.constexpr, + HAS_DX1: tl.constexpr, + HAS_B1: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + row_start = row_block_id * rows_per_program + # Do not early exit if row_start >= M, because we need to write DW and DB + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += row_start * stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += row_start * stride_dres_in_row + DY += row_start * stride_dy_row + DX += row_start * stride_dx_row + if HAS_DY1: + DY1 += row_start * stride_dy1_row + if HAS_DX1: + DX1 += row_start * stride_dx1_row + if RECOMPUTE_OUTPUT: + Y += row_start * stride_y_row + w = tl.load(W + cols, mask=mask).to(tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_DY1: + w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_DY1: + dw1 = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_B1: + db1 = tl.zeros((BLOCK_N,), dtype=tl.float32) + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + if HAS_DY1: + dy1 = tl.load(DY1 + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.0) + if RECOMPUTE_OUTPUT: + y = xhat * w + b if HAS_BIAS else xhat * w + tl.store(Y + cols, y, mask=mask) + wdy = w * dy + dw += dy * xhat + if HAS_BIAS: + db += dy + if HAS_DY1: + wdy += w1 * dy1 + dw1 += dy1 * xhat + if HAS_B1: + db1 += dy1 + if not IS_RMS_NORM: + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + c1 = tl.sum(xhat * wdy, axis=0) / N + dx = (wdy - xhat * c1) * rstd + if HAS_DRESIDUAL: + dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) + dx += dres + # Write dx + if STORE_DRESIDUAL: + tl.store(DRESIDUAL_IN + cols, dx, mask=mask) + if HAS_DX1: + if HAS_DROPOUT: + keep_mask = ( + tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p + ) + dx1 = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) + else: + dx1 = dx + tl.store(DX1 + cols, dx1, mask=mask) + if HAS_DROPOUT: + keep_mask = tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p + dx = tl.where(keep_mask, dx / (1.0 - dropout_p), 0.0) + if HAS_ROWSCALE: + rowscale = tl.load(ROWSCALE + row).to(tl.float32) + dx *= rowscale + tl.store(DX + cols, dx, mask=mask) + + X += stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += stride_dres_in_row + if RECOMPUTE_OUTPUT: + Y += stride_y_row + DY += stride_dy_row + DX += stride_dx_row + if HAS_DY1: + DY1 += stride_dy1_row + if HAS_DX1: + DX1 += stride_dx1_row + tl.store(DW + row_block_id * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * N + cols, db, mask=mask) + if HAS_DY1: + tl.store(DW1 + row_block_id * N + cols, dw1, mask=mask) + if HAS_B1: + tl.store(DB1 + row_block_id * N + cols, db1, mask=mask) + + +def _layer_norm_bwd( + dy, + x, + weight, + bias, + eps, + mean, + rstd, + dresidual=None, + dy1=None, + weight1=None, + bias1=None, + seeds=None, + dropout_p=0.0, + rowscale=None, + has_residual=False, + has_x1=False, + is_rms_norm=False, + x_dtype=None, + recompute_output=False, +): + M, N = x.shape + assert x.stride(-1) == 1 + assert dy.stride(-1) == 1 + assert dy.shape == (M, N) + if dresidual is not None: + assert dresidual.stride(-1) == 1 + assert dresidual.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + if dy1 is not None: + assert weight1 is not None + assert dy1.shape == dy.shape + assert dy1.stride(-1) == 1 + if weight1 is not None: + assert weight1.shape == (N,) + assert weight1.stride(-1) == 1 + if bias1 is not None: + assert bias1.shape == (N,) + assert bias1.stride(-1) == 1 + if seeds is not None: + assert seeds.is_contiguous() + assert seeds.shape == (M if not has_x1 else M * 2,) + if rowscale is not None: + assert rowscale.is_contiguous() + assert rowscale.shape == (M,) + # allocate output + dx = ( + torch.empty_like(x) + if x_dtype is None + else torch.empty(M, N, dtype=x_dtype, device=x.device) + ) + dresidual_in = ( + torch.empty_like(x) + if has_residual + and (dx.dtype != x.dtype or dropout_p > 0.0 or rowscale is not None or has_x1) + else None + ) + dx1 = torch.empty_like(dx) if (has_x1 and dropout_p > 0.0) else None + y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None + if recompute_output: + assert weight1 is None, "recompute_output is not supported with parallel LayerNorm" + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count + _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) + _db = ( + torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) + if bias is not None + else None + ) + _dw1 = torch.empty_like(_dw) if weight1 is not None else None + _db1 = torch.empty_like(_db) if bias1 is not None else None + rows_per_program = math.ceil(M / sm_count) + grid = (sm_count,) + with torch.cuda.device(x.device.index): + _layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + y, + dy, + dx, + _dw, + _db, + dresidual, + weight1, + dy1, + dx1, + _dw1, + _db1, + dresidual_in, + rowscale, + seeds, + mean, + rstd, + x.stride(0), + 0 if not recompute_output else y.stride(0), + dy.stride(0), + dx.stride(0), + dresidual.stride(0) if dresidual is not None else 0, + dy1.stride(0) if dy1 is not None else 0, + dx1.stride(0) if dx1 is not None else 0, + dresidual_in.stride(0) if dresidual_in is not None else 0, + M, + N, + eps, + dropout_p, + rows_per_program, + is_rms_norm, + BLOCK_N, + dresidual is not None, + dresidual_in is not None, + bias is not None, + dropout_p > 0.0, + ) + dw = _dw.sum(0).to(weight.dtype) + db = _db.sum(0).to(bias.dtype) if bias is not None else None + dw1 = _dw1.sum(0).to(weight1.dtype) if weight1 is not None else None + db1 = _db1.sum(0).to(bias1.dtype) if bias1 is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype and dropout_p == 0.0 and rowscale is None: + dresidual_in = dx + if has_x1 and dropout_p == 0.0: + dx1 = dx + return ( + (dx, dw, db, dresidual_in, dx1, dw1, db1) + if not recompute_output + else (dx, dw, db, dresidual_in, dx1, dw1, db1, y) + ) + + +class LayerNormFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + return_dropout_mask=False, + out=None, + residual_out=None + ): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + if residual.stride(-1) != 1: + residual = residual.contiguous() + if x1 is not None: + assert x1.shape == x_shape_og + assert rowscale is None, "rowscale is not supported with parallel LayerNorm" + x1 = x1.reshape(-1, x1.shape[-1]) + if x1.stride(-1) != 1: + x1 = x1.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + if weight1 is not None: + weight1 = weight1.contiguous() + if bias1 is not None: + bias1 = bias1.contiguous() + if rowscale is not None: + rowscale = rowscale.reshape(-1).contiguous() + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + if out is not None: + out = out.reshape(-1, out.shape[-1]) + if residual_out is not None: + residual_out = residual_out.reshape(-1, residual_out.shape[-1]) + y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd( + x, + weight, + bias, + eps, + residual, + x1, + weight1, + bias1, + dropout_p=dropout_p, + rowscale=rowscale, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + return_dropout_mask=return_dropout_mask, + out=out, + residual_out=residual_out + ) + ctx.save_for_backward( + residual_out, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd + ) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.dropout_p = dropout_p + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.has_x1 = x1 is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + y = y.reshape(x_shape_og) + y1 = y1.reshape(x_shape_og) if y1 is not None else None + residual_out = residual_out.reshape(x_shape_og) if residual_out is not None else None + dropout_mask = dropout_mask.reshape(x_shape_og) if dropout_mask is not None else None + dropout_mask1 = dropout_mask1.reshape(x_shape_og) if dropout_mask1 is not None else None + if not return_dropout_mask: + if weight1 is None: + return y if not prenorm else (y, residual_out) + else: + return (y, y1) if not prenorm else (y, y1, residual_out) + else: + if weight1 is None: + return ( + (y, dropout_mask, dropout_mask1) + if not prenorm + else (y, residual_out, dropout_mask, dropout_mask1) + ) + else: + return ( + (y, y1, dropout_mask, dropout_mask1) + if not prenorm + else (y, y1, residual_out, dropout_mask, dropout_mask1) + ) + + @staticmethod + def backward(ctx, dy, *args): + x, weight, bias, weight1, bias1, rowscale, seeds, mean, rstd = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + if weight1 is not None: + dy1, args = args[0], args[1:] + dy1 = dy1.reshape(-1, dy1.shape[-1]) + if dy1.stride(-1) != 1: + dy1 = dy1.contiguous() + assert dy1.shape == x.shape + else: + dy1 = None + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + if dresidual.stride(-1) != 1: + dresidual = dresidual.contiguous() + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dw, db, dresidual_in, dx1, dw1, db1 = _layer_norm_bwd( + dy, + x, + weight, + bias, + ctx.eps, + mean, + rstd, + dresidual, + dy1, + weight1, + bias1, + seeds, + ctx.dropout_p, + rowscale, + ctx.has_residual, + ctx.has_x1, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + ) + return ( + dx.reshape(ctx.x_shape_og), + dw, + db, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + dx1.reshape(ctx.x_shape_og) if dx1 is not None else None, + dw1, + db1, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def layer_norm_fn( + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + return_dropout_mask=False, + out=None, + residual_out=None +): + return LayerNormFn.apply( + x, + weight, + bias, + residual, + x1, + weight1, + bias1, + eps, + dropout_p, + rowscale, + prenorm, + residual_in_fp32, + is_rms_norm, + return_dropout_mask, + out, + residual_out + ) + + +def rms_norm_fn( + x, + weight, + bias, + residual=None, + x1=None, + weight1=None, + bias1=None, + eps=1e-6, + dropout_p=0.0, + rowscale=None, + prenorm=False, + residual_in_fp32=False, + return_dropout_mask=False, + out=None, + residual_out=None +): + return LayerNormFn.apply( + x, + weight, + bias, + residual, + x1, + weight1, + bias1, + eps, + dropout_p, + rowscale, + prenorm, + residual_in_fp32, + True, + return_dropout_mask, + out, + residual_out + ) + + +class RMSNorm(torch.nn.Module): + + def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if dropout_p > 0.0: + self.drop = torch.nn.Dropout(dropout_p) + else: + self.drop = None + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return rms_norm_fn( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + dropout_p=self.drop.p if self.drop is not None and self.training else 0.0, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class LayerNormLinearFn(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + ): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + if residual.stride(-1) != 1: + residual = residual.contiguous() + norm_weight = norm_weight.contiguous() + if norm_bias is not None: + norm_bias = norm_bias.contiguous() + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + y, _, mean, rstd, residual_out, *rest = _layer_norm_fwd( + x, + norm_weight, + norm_bias, + eps, + residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = linear_weight.to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @custom_bwd + def backward(ctx, dout, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + if dresidual.stride(-1) != 1: + dresidual = dresidual.contiguous() + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, _, _, _, y = _layer_norm_bwd( + dy, + x, + norm_weight, + norm_bias, + ctx.eps, + mean, + rstd, + dresidual=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y) + return ( + dx.reshape(ctx.x_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +def layer_norm_linear_fn( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, +): + return LayerNormLinearFn.apply( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/linear.py b/lib/python3.12/site-packages/flash_attn/ops/triton/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..a8966dbc345ab0e593df0124451ee7be3dae131a --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/linear.py @@ -0,0 +1,594 @@ +# Adapted from https://github.com/ELS-RD/kernl/blob/main/src/kernl/implementations/linear_layer.py +# and https://github.com/openai/triton/blob/master/python/triton/ops/matmul.py +from typing import Optional + +import torch +import triton +import triton.language as tl +from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time + +from flash_attn.ops.triton.k_activations import ( + gelu, + gelu_approx, + gelu_approx_grad, + gelu_grad, + squared_relu, + squared_relu_grad, +) + +# CREDITS: Initially inspired by the Triton tutorial on matrix multiplications + + +def init_to_zero(name): + return lambda nargs: nargs[name].zero_() + + +def get_configs_io_bound(): + configs = [] + for num_stages in [2, 3, 4, 5, 6]: + for block_m in [16, 32]: + for block_k in [32, 64]: + for block_n in [32, 64, 128, 256]: + num_warps = 2 if block_n <= 64 else 4 + configs.append( + triton.Config( + { + "BLOCK_M": block_m, + "BLOCK_N": block_n, + "BLOCK_K": block_k, + "SPLIT_K": 1, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + # split_k not used + # for split_k in [2, 4, 8, 16]: + # configs.append(triton.Config( + # {'BLOCK_M': block_m, 'BLOCK_N': block_n, 'BLOCK_K': block_k, 'SPLIT_K': split_k}, + # num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero('C'))) + return configs + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 + ), + # good for int8 + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 + ), + ] + + get_configs_io_bound(), + key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], + prune_configs_by={ + "early_config_prune": early_config_prune, + "perf_model": estimate_matmul_time, + "top_k": 10, + }, +) +@triton.heuristics( + { + "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, + } +) +@triton.jit +def kernel_fwd( + C, # Pointers to matrices + ACT_INPUT, + A, + B, + bias, + # Matrix dimensions + M, + N, + K, + CACHE_KEY_M, + CACHE_KEY_N, + CACHE_KEY_K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. stride_am is how much to increase a_ptr + # by to get the element one row down (A has M rows) + stride_cm, + # stride_cn, # Assume that stride_cn == 1 + stride_am, + stride_ak, + stride_bn, + stride_bk, + # Meta-parameters + BLOCK_M: tl.constexpr, + GROUP_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # split k not used, not performant with activation, kept because early_config_prune is expecting it + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + A_ROWMAJOR: tl.constexpr, + B_COLMAJOR: tl.constexpr, + BIAS: tl.constexpr, + SAVE_ACT_INPUT: tl.constexpr, + ACTIVATION: tl.constexpr, +): + + """ + Kernel for computing Out = activation(A x W + C) + - Input has shape (M, K) + - Weight has shape (K, N) + - Bias has shape (N,) + - Output has shape (M, N) + - ActInputs (optional) has shape (M, N) + 'ActInputs' optionally saves the A x W + C intermediate for backward computations + This kernel will consolidate over K + """ + + pid = tl.program_id(axis=0) + + grid_m = (M + BLOCK_M - 1) // BLOCK_M + grid_n = (N + BLOCK_N - 1) // BLOCK_N + # re-order program ID for better L2 performance + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // (group_size) + + # now compute the block that each program will go through + # rm (resp. rn) denotes a range of indices + # for rows (resp. col) of C + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + # trick to avoid masking on M and N axis + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = tl.arange(0, BLOCK_K) + + if A_ROWMAJOR: + A = A + (ram[:, None] * stride_am + rk[None, :]) + else: + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + if B_COLMAJOR: + B = B + (rk[:, None] + rbn[None, :] * stride_bn) + else: + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k in range(K, 0, -BLOCK_K): + if EVEN_K: + a = tl.load(A) + b = tl.load(B) + else: + a = tl.load(A, mask=rk[None, :] < k, other=0.0) + b = tl.load(B, mask=rk[:, None] < k, other=0.0) + acc += tl.dot(a, b) + + if A_ROWMAJOR: + A += BLOCK_K + else: + A += BLOCK_K * stride_ak + if B_COLMAJOR: + B += BLOCK_K + else: + B += BLOCK_K * stride_bk + + # Putting bias after the matmul (instead of before) is faster, idk why + if BIAS: + bias = tl.load(bias + rn, mask=rn < N, other=0.0).to(tl.float32) + acc += bias[None, :] + + # optional: save the activation inputs + if SAVE_ACT_INPUT: + # act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] * stride_cn + act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] + tl.store(act_in_ptrs, acc) + + # optional: fused activation (while the data is in shared memory) + if ACTIVATION == "gelu": + acc = gelu(acc) + elif ACTIVATION == "gelu_approx": + acc = gelu_approx(acc) + elif ACTIVATION == "squared_relu": + acc = squared_relu(acc) + # rematerialize rm and rn to save registers + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + # write back result + # C = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn + C = C + rm[:, None] * stride_cm + rn[None, :] + mask = (rm < M)[:, None] & (rn < N)[None, :] + tl.store(C, acc) + + +def triton_linear_act( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: str = "id", + save_act_input: bool = False, +) -> torch.Tensor: + """ + Compute e = activation(x @ weight.T + bias). + This wrapper kicks the `kernel_fwd` Triton kernel + :param x: input tensor + :param weight: weight matrix + :param bias: an optional bias tensor + :param activation: Activation name. Needs to be a Triton kernel. + :param act_input: an optional tensor to save the activation inputs (for backward) + :return: result tensor + """ + # if torch.is_autocast_enabled(): + # dtype = torch.get_autocast_gpu_dtype() + # x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]] + + assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] + + batch_shape, n = x.shape[:-1], x.shape[-1] + batch_dim = batch_shape.numel() + x_reshaped = x.reshape(batch_dim, n) + + if x_reshaped.stride(0) > 1 and x_reshaped.stride(1) > 1: + x_reshaped = x_reshaped.contiguous() + if weight.stride(0) > 1 and weight.stride(1) > 1: + weight = weight.contiguous() + bias = bias.contiguous() if bias is not None else None + + assert ( + x.dtype == weight.dtype + ), f"Input and weight must have the same dtype, got {x.dtype} and {weight.dtype}" + if bias is not None: + assert ( + x.dtype == bias.dtype + ), f"Input and bias must have the same dtype, got {x.dtype} and {bias.dtype}" + assert ( + x_reshaped.shape[1] == weight.shape[1] + ), f"Incompatible dimensions: {x_reshaped.shape} - {weight.shape}" + + assert ( + bias is None or bias.shape[0] == weight.shape[0] + ), "Incompatible dimensions in between weight and bias" + + M, K = x_reshaped.shape + N, K = weight.shape + + output = torch.empty((M, N), device=x.device, dtype=x.dtype) + act_input = torch.empty_like(output) if save_act_input else None + + # 1D launch kernel where each block gets its own program. + grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa + + kernel_fwd[grid]( + output, + act_input, + x_reshaped, + weight, # data ptrs + bias if bias is not None else x, # auto skip bias if not present + M, # shapes + N, + K, + M // 32, # key for triton cache (limit number of compilations) + N // 32, + K // 32, + stride_cm=output.stride(0), # strides + # stride_cn=output.stride(1), + stride_am=x_reshaped.stride(0), + stride_ak=x_reshaped.stride(1), + stride_bk=weight.stride(1), + stride_bn=weight.stride(0), + BIAS=bias is not None, # optional fused bias + SAVE_ACT_INPUT=save_act_input, # optional save activation inputs + ACTIVATION=activation, # optional fused activation + A_ROWMAJOR=x_reshaped.stride(1) == 1, + B_COLMAJOR=weight.stride(1) == 1, + GROUP_M=8, # speed optimization: group the programs + ) + + if not save_act_input: + return output.reshape(*batch_shape, output.shape[-1]) + else: + return ( + output.reshape(*batch_shape, output.shape[-1]), + act_input.reshape(*batch_shape, act_input.shape[-1]), + ) + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2 + ), + # good for int8 + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2 + ), + ] + + get_configs_io_bound(), + key=["CACHE_KEY_M", "CACHE_KEY_N", "CACHE_KEY_K"], + prune_configs_by={ + "early_config_prune": early_config_prune, + "perf_model": estimate_matmul_time, + "top_k": 10, + }, +) +@triton.heuristics( + { + "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, + } +) +@triton.jit +def kernel_bwd( + C, # Pointers to matrices + ACT_INPUT, + A, + B, + # Matrix dimensions + M, + N, + K, + CACHE_KEY_M, + CACHE_KEY_N, + CACHE_KEY_K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. stride_am is how much to increase a_ptr + # by to get the element one row down (A has M rows) + stride_cm, + # stride_cn, # Assume that stride_cn == 1 + stride_am, + stride_ak, + stride_bk, + stride_bn, + # Meta-parameters + BLOCK_M: tl.constexpr, + GROUP_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + # split k not used, not performant with activation, kept because early_config_prune is expecting it + SPLIT_K: tl.constexpr, + EVEN_K: tl.constexpr, + ACTIVATION: tl.constexpr, +): + + """ + Kernel for computing Out = activation(A x W + C) + - Input has shape (M, K) + - Weight has shape (K, N) + - Output has shape (M, N) + - ActInputs (optional) has shape (M, N) + 'ActInputs' optionally saves the A x W + C intermediate for backward computations + This kernel will consolidate over K + """ + + pid = tl.program_id(axis=0) + + grid_m = (M + BLOCK_M - 1) // BLOCK_M + grid_n = (N + BLOCK_N - 1) // BLOCK_N + # re-order program ID for better L2 performance + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // (group_size) + + # now compute the block that each program will go through + # rm (resp. rn) denotes a range of indices + # for rows (resp. col) of C + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + # trick to avoid masking on M and N axis + ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) + rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) + rk = tl.arange(0, BLOCK_K) + + A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) + B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k in range(K, 0, -BLOCK_K): + if EVEN_K: + a = tl.load(A) + b = tl.load(B) + else: + a = tl.load(A, mask=rk[None, :] < k, other=0.0) + b = tl.load(B, mask=rk[:, None] < k, other=0.0) + acc += tl.dot(a, b) + + A += BLOCK_K * stride_ak + B += BLOCK_K * stride_bk + + # optional: fused activation (while the data is in shared memory) + if ACTIVATION != "id": + act_in_ptrs = ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :] + act_input = tl.load(act_in_ptrs).to(acc.dtype) + if ACTIVATION == "gelu": + acc *= gelu_grad(act_input) + elif ACTIVATION == "gelu_approx": + acc *= gelu_approx_grad(act_input) + elif ACTIVATION == "squared_relu": + acc *= squared_relu_grad(act_input) + + # rematerialize rm and rn to save registers + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + # write back result + C = C + rm[:, None] * stride_cm + rn[None, :] + mask = (rm < M)[:, None] & (rn < N)[None, :] + tl.store(C, acc, mask=mask) + + +def triton_dgrad_act( + grad_output: torch.Tensor, + weight: torch.Tensor, + activation: str = "id", + act_input: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + Compute e = activation(grad_output @ weight + bias). + This wrapper kicks the `kernel_fwd` Triton kernel + :param grad_output: input tensor + :param weight: weight matrix + :param activation: Activation name. Needs to be a Triton kernel. + :param act_input: an optional tensor to save the activation inputs (for backward) + :return: result tensor + """ + assert activation in ["id", "gelu", "gelu_approx", "squared_relu"] + + batch_shape, n = grad_output.shape[:-1], grad_output.shape[-1] + batch_dim = batch_shape.numel() + grad_output_reshaped = grad_output.reshape(batch_dim, n) + + if grad_output_reshaped.stride(0) > 1 and grad_output_reshaped.stride(1) > 1: + grad_output_reshaped = grad_output_reshaped.contiguous() + if weight.stride(0) > 1 and weight.stride(1) > 1: + weight = weight.contiguous() + + assert ( + grad_output.dtype == weight.dtype + ), f"grad_output and weight must have the same dtype, got {grad_output.dtype} and {weight.dtype}" + assert ( + grad_output_reshaped.shape[1] == weight.shape[0] + ), f"Incompatible dimensions: {grad_output_reshaped.shape} - {weight.shape}" + if activation != "id": + assert act_input is not None, f"act_input is required for activation {activation}" + + # M, N, K in bwd are different from M, N, K in fwd + M, K = grad_output_reshaped.shape + K, N = weight.shape + + grad_input = torch.empty((M, N), device=grad_output.device, dtype=grad_output.dtype) + + # 1D launch kernel where each block gets its own program. + grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) # noqa + + kernel_bwd[grid]( + grad_input, + act_input, + grad_output_reshaped, + weight, # data ptrs + M, # shapes + N, + K, + M // 32, # key for triton cache (limit number of compilations) + N // 32, + K // 32, + stride_cm=grad_input.stride(0), # strides + # stride_cn=grad_input.stride(1), + stride_am=grad_output_reshaped.stride(0), + stride_ak=grad_output_reshaped.stride(1), + stride_bk=weight.stride(0), + stride_bn=weight.stride(1), + ACTIVATION=activation, # optional fused activation + GROUP_M=8, # speed optimization: group the programs + ) + + return grad_input.reshape(*batch_shape, grad_input.shape[-1]) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/mlp.py b/lib/python3.12/site-packages/flash_attn/ops/triton/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..b795310f1c8afc8203124597bb6ca70f1af7ed11 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/mlp.py @@ -0,0 +1,149 @@ +# The triton fused matmul + sqrelu is faster for fp16 but slower for bf16, compared +# to naive implementation. +import fused_dense_lib as fused_dense_cuda +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.cuda.amp import custom_bwd, custom_fwd + +from flash_attn.ops.activations import sqrelu_bwd, sqrelu_fwd +from flash_attn.ops.triton.linear import triton_dgrad_act, triton_linear_act + + +class FusedDenseSqreluDenseFunc(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0): + """checkpoint_lvl: + 0: no recomputation in the bwd + 1: recompute gelu_out in the bwd + 2: recompute act_input and gelu_out in the bwd + """ + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_gpu_dtype() + x, weight1, bias1, weight2, bias2 = [ + a.to(dtype=dtype) for a in [x, weight1, bias1, weight2, bias2] + ] + is_bf16 = x.dtype == torch.bfloat16 + assert checkpoint_lvl in [0, 1, 2] + x = x.contiguous() + weight1 = weight1.contiguous() + bias1 = bias1.contiguous() + weight2 = weight2.contiguous() + bias2 = bias2.contiguous() + batch_shape, n = x.shape[:-1], x.shape[-1] + batch_dim = batch_shape.numel() + if is_bf16: + act_input = fused_dense_cuda.linear_bias_forward( + x.reshape(batch_dim, n), weight1, bias1 + ) + output1 = sqrelu_fwd(act_input) + else: + save_act_input = checkpoint_lvl != 2 + result = triton_linear_act( + x.reshape(batch_dim, n), + weight1, + bias1, + activation="squared_relu", + save_act_input=save_act_input, + ) + if save_act_input: + output1, act_input = result + else: + output1 = result + output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2) + ctx.checkpoint_lvl = checkpoint_lvl + if checkpoint_lvl == 0: + ctx.save_for_backward(x, weight1, bias1, weight2, act_input, output1) + elif checkpoint_lvl == 1: + ctx.save_for_backward(x, weight1, bias1, weight2, act_input) + elif checkpoint_lvl == 2: + ctx.save_for_backward(x, weight1, bias1, weight2) + return output2.reshape(*batch_shape, output2.shape[-1]) + + @staticmethod + @custom_bwd + def backward(ctx, grad_output): + grad_output = grad_output.contiguous() + checkpoint_lvl = ctx.checkpoint_lvl + x, weight1, bias1, weight2, *rest = ctx.saved_tensors + batch_shape, n = x.shape[:-1], x.shape[-1] + batch_dim = batch_shape.numel() + is_bf16 = x.dtype == torch.bfloat16 + if checkpoint_lvl == 0: + act_input, output1 = rest + elif checkpoint_lvl == 1: + (act_input,) = rest + output1 = sqrelu_fwd(act_input) + elif checkpoint_lvl == 2: + if is_bf16: + act_input = fused_dense_cuda.linear_bias_forward( + x.reshape(batch_dim, n), weight1, bias1 + ) + output1 = sqrelu_fwd(act_input) + else: + output1, act_input = triton_linear_act( + x.reshape(batch_dim, n), + weight1, + bias1, + activation="squared_relu", + save_act_input=True, + ) + + if is_bf16: + grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) + grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) + grad_output1 = grad_output @ weight2 + grad_act_input = sqrelu_bwd(grad_output1, act_input) + grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( + x.reshape(batch_dim, n), weight1, grad_act_input + ) + else: + grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1]) + grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output) + grad_act_input = triton_dgrad_act( + grad_output, weight2, activation="squared_relu", act_input=act_input + ) + grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward( + x.reshape(batch_dim, n), weight1, grad_act_input + ) + return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None + + +fused_dense_sqrelu_dense_function = FusedDenseSqreluDenseFunc.apply + + +class FusedDenseSqreluDense(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + bias1=True, + bias2=True, + checkpoint_lvl=0, + device=None, + dtype=None, + ): + """ + checkpoint_lvl (increasing lvl means slower but more memory saving): + 0: no recomputation in the bwd + 1: recompute gelu_out in the bwd + 2: recompute gelu_in and gelu_out in the bwd + """ + assert checkpoint_lvl in [0, 1, 2] + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features * 4 + assert bias1 == True, "DenseSqreluDense module without bias is currently not supported" + assert bias2 == True, "DenseSqreluDense module without bias is currently not supported" + self.checkpoint_lvl = checkpoint_lvl + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) + + def forward(self, x): + assert x.is_cuda + return fused_dense_sqrelu_dense_function( + x, self.fc1.weight, self.fc1.bias, self.fc2.weight, self.fc2.bias, self.checkpoint_lvl + ) diff --git a/lib/python3.12/site-packages/flash_attn/ops/triton/rotary.py b/lib/python3.12/site-packages/flash_attn/ops/triton/rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..0ee56d647738cb1c74e2b0a341e7ae7bef85f966 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/ops/triton/rotary.py @@ -0,0 +1,227 @@ +# Copyright (c) 2023, Tri Dao. + +from typing import Optional, Union + +import torch + +import triton +import triton.language as tl + + +@triton.jit +def rotary_kernel( + OUT, # Pointers to matrices + X, + COS, + SIN, + CU_SEQLENS, + SEQLEN_OFFSETS, # this could be int or a pointer + # Matrix dimensions + seqlen, + rotary_dim, + seqlen_ro, + # strides + stride_out_batch, + stride_out_seqlen, + stride_out_nheads, + stride_out_headdim, + stride_x_batch, + stride_x_seqlen, + stride_x_nheads, + stride_x_headdim, + # Meta-parameters + BLOCK_K: tl.constexpr, + IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, + IS_VARLEN: tl.constexpr, + INTERLEAVED: tl.constexpr, + CONJUGATE: tl.constexpr, + BLOCK_M: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_batch = tl.program_id(axis=1) + pid_head = tl.program_id(axis=2) + rotary_dim_half = rotary_dim // 2 + + if not IS_VARLEN: + X = X + pid_batch * stride_x_batch + pid_head * stride_x_nheads + OUT = OUT + pid_batch * stride_out_batch + pid_head * stride_out_nheads + else: + start_idx = tl.load(CU_SEQLENS + pid_batch) + seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx + X = X + start_idx * stride_x_seqlen + pid_head * stride_x_nheads + OUT = OUT + start_idx * stride_out_seqlen + pid_head * stride_out_nheads + + if pid_m * BLOCK_M >= seqlen: + return + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + if not IS_SEQLEN_OFFSETS_TENSOR: + rm_cs = rm + SEQLEN_OFFSETS + else: + rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) + rk = tl.arange(0, BLOCK_K) + rk_half = tl.arange(0, BLOCK_K // 2) + + if not INTERLEAVED: + # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT + X = X + (rm[:, None] * stride_x_seqlen + rk_half[None, :] * stride_x_headdim) + COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :]) + SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :]) + cos = tl.load( + COS, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=1.0 + ).to(tl.float32) + sin = tl.load( + SIN, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=0.0 + ).to(tl.float32) + x0 = tl.load( + X, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), other=0.0 + ).to(tl.float32) + x1 = tl.load( + X + rotary_dim_half * stride_x_headdim, + mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), + other=0.0, + ).to(tl.float32) + if CONJUGATE: + sin = -sin + o0 = x0 * cos - x1 * sin + o1 = x0 * sin + x1 * cos + # write back result + OUT = OUT + (rm[:, None] * stride_out_seqlen + rk_half[None, :] * stride_out_headdim) + tl.store(OUT, o0, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half)) + tl.store( + OUT + rotary_dim_half * stride_out_headdim, + o1, + mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), + ) + else: + # We don't want to load X[0, 2, 4, ...] and X[1, 3, 5, ...] separately since both are slow. + # Instead, we load x0 = X[0, 1, 2, 3, ...] and x1 = X[1, 0, 3, 2, ...]. + # Loading x0 will be fast but x1 will be slow. + # Then we load cos = COS[0, 0, 1, 1, ...] and sin = SIN[0, 0, 1, 1, ...]. + # Then we do the calculation and use tl.where to pick put the right outputs for the even + # and for the odd indices. + rk_swap = rk + ((rk + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ... + rk_repeat = tl.arange(0, BLOCK_K) // 2 + X0 = X + (rm[:, None] * stride_x_seqlen + rk[None, :] * stride_x_headdim) + X1 = X + (rm[:, None] * stride_x_seqlen + rk_swap[None, :] * stride_x_headdim) + COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :]) + SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :]) + cos = tl.load( + COS, + mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half), + other=1.0, + ).to(tl.float32) + sin = tl.load( + SIN, + mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half), + other=0.0, + ).to(tl.float32) + x0 = tl.load(X0, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim), other=0.0).to( + tl.float32 + ) + x1 = tl.load( + X1, mask=(rm[:, None] < seqlen) & (rk_swap[None, :] < rotary_dim), other=0.0 + ).to(tl.float32) + if CONJUGATE: + sin = -sin + x0_cos = x0 * cos + x1_sin = x1 * sin + out = tl.where(rk[None, :] % 2 == 0, x0_cos - x1_sin, x0_cos + x1_sin) + OUT = OUT + (rm[:, None] * stride_out_seqlen + rk[None, :] * stride_out_headdim) + tl.store(OUT, out, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim)) + + +def apply_rotary( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + seqlen_offsets: Union[int, torch.Tensor] = 0, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + interleaved=False, + inplace=False, + conjugate=False, +) -> torch.Tensor: + """ + Arguments: + x: (batch, seqlen, nheads, headdim) if cu_seqlens is None + else (total_seqlen, nheads, headdim). + cos: (seqlen_ro, rotary_dim / 2) + sin: (seqlen_ro, rotary_dim / 2) + seqlen_offsets: integer or integer tensor of size (batch,) + cu_seqlens: (batch + 1,) or None + max_seqlen: int + Returns: + y: (batch, seqlen, nheads, headdim) + """ + is_varlen = cu_seqlens is not None + if not is_varlen: + batch, seqlen, nheads, headdim = x.shape + else: + assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" + total_seqlen, nheads, headdim = x.shape + batch_p_1 = cu_seqlens.shape[0] + batch = batch_p_1 - 1 + seqlen = max_seqlen + seqlen_ro, rotary_dim = cos.shape + assert sin.shape == cos.shape + rotary_dim *= 2 + assert rotary_dim <= headdim, "rotary_dim must be <= headdim" + assert headdim <= 256, "Only support headdim <= 256" + assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen" + + assert ( + cos.dtype == sin.dtype + ), f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}" + assert ( + x.dtype == cos.dtype + ), f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}" + + cos, sin = cos.contiguous(), sin.contiguous() + if isinstance(seqlen_offsets, torch.Tensor): + assert seqlen_offsets.shape == (batch,) + assert seqlen_offsets.dtype in [torch.int32, torch.int64] + seqlen_offsets = seqlen_offsets.contiguous() + else: + assert seqlen_offsets + seqlen <= seqlen_ro + + output = torch.empty_like(x) if not inplace else x + if rotary_dim < headdim and not inplace: + output[..., rotary_dim:].copy_(x[..., rotary_dim:]) + + BLOCK_K = ( + 32 + if rotary_dim <= 32 + else (64 if rotary_dim <= 64 else (128 if rotary_dim <= 128 else 256)) + ) + grid = lambda META: (triton.cdiv(seqlen, META["BLOCK_M"]), batch, nheads) # noqa + BLOCK_M = 4 if interleaved else (8 if rotary_dim <= 128 else 4) + + # Need this, otherwise Triton tries to launch from cuda:0 and we get + # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?) + with torch.cuda.device(x.device.index): + rotary_kernel[grid]( + output, # data ptrs + x, + cos, + sin, + cu_seqlens, + seqlen_offsets, + seqlen, # shapes + rotary_dim, + seqlen_ro, + output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 + output.stride(-3), # seqlen_stride or total_seqlen_stride + output.stride(-2), # nheads_stride + output.stride(-1), # headdim_stride + x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0 + x.stride(-3), # seqlen stride or total_seqlen_stride + x.stride(-2), # nheads stride + x.stride(-1), # headdim stride + BLOCK_K, + isinstance(seqlen_offsets, torch.Tensor), + is_varlen, + interleaved, + conjugate, + BLOCK_M, + ) + return output diff --git a/lib/python3.12/site-packages/flash_attn/utils/__init__.py b/lib/python3.12/site-packages/flash_attn/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/flash_attn/utils/benchmark.py b/lib/python3.12/site-packages/flash_attn/utils/benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..15b30405f209921189b75f7307814876350e7317 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/utils/benchmark.py @@ -0,0 +1,268 @@ +# Copyright (c) 2023, Tri Dao. +""" Useful functions for writing test code. """ + +import torch +import torch.utils.benchmark as benchmark + + +def benchmark_forward( + fn, *inputs, repeats=10, desc="", verbose=True, amp=False, amp_dtype=torch.float16, **kwinputs +): + """Use Pytorch Benchmark on the forward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward pass") + + def amp_wrapper(*inputs, **kwinputs): + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + fn(*inputs, **kwinputs) + + t = benchmark.Timer( + stmt="fn_amp(*inputs, **kwinputs)", + globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_backward( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(*inputs, y, grad): + # Set .grad to None to avoid extra operation of gradient accumulation + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(*inputs, y=y, grad=grad)", + globals={"f": f, "inputs": inputs, "y": y, "grad": grad}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_combined( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward + Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(grad, *inputs, **kwinputs): + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(grad, *inputs, **kwinputs)", + globals={"f": f, "fn": fn, "inputs": inputs, "grad": grad, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_fwd_bwd( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def benchmark_all( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_combined( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def pytorch_profiler( + fn, + *inputs, + trace_filename=None, + backward=False, + amp=False, + amp_dtype=torch.float16, + cpu=False, + verbose=True, + **kwinputs, +): + """Wrap benchmark functions in Pytorch profiler to see CUDA information.""" + if backward: + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + g = torch.randn_like(out) + for _ in range(30): # Warm up + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + # Backward should be done outside autocast + if backward: + out.backward(g, retain_graph=True) + activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [ + torch.profiler.ProfilerActivity.CUDA + ] + with torch.profiler.profile( + activities=activities, + record_shapes=True, + # profile_memory=True, + with_stack=True, + ) as prof: + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + if backward: + out.backward(g, retain_graph=True) + if verbose: + # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50)) + print(prof.key_averages().table(row_limit=50)) + if trace_filename is not None: + prof.export_chrome_trace(trace_filename) + + +def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + fn(*inputs, **kwinputs) + torch.cuda.synchronize() + mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000) + if verbose: + print(f"{desc} max memory: {mem}GB") + torch.cuda.empty_cache() + return mem diff --git a/lib/python3.12/site-packages/flash_attn/utils/distributed.py b/lib/python3.12/site-packages/flash_attn/utils/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..74c55279645cd0fd687584bc1b7374c8c3c73e56 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/utils/distributed.py @@ -0,0 +1,144 @@ +from typing import Optional + +import torch +from torch import Tensor +from torch.distributed import ProcessGroup + +# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for +# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent +# version of PyTorch. The following 4 lines are for backward compatibility with +# older PyTorch. +if "all_gather_into_tensor" not in dir(torch.distributed): + torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base +if "reduce_scatter_tensor" not in dir(torch.distributed): + torch.distributed.reduce_scatter_tensor = torch.distributed._reduce_scatter_base + + +# Raw operation, does not support autograd, but does support async +def all_gather_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + world_size = torch.distributed.get_world_size(process_group) + output = torch.empty( + world_size * input_.shape[0], *input_.shape[1:], dtype=input_.dtype, device=input_.device + ) + handle = torch.distributed.all_gather_into_tensor( + output, input_.contiguous(), group=process_group, async_op=async_op + ) + return output, handle + + +# Raw operation, does not support autograd, but does support async +def reduce_scatter_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + world_size = torch.distributed.get_world_size(process_group) + assert input_.shape[0] % world_size == 0 + output = torch.empty( + input_.shape[0] // world_size, *input_.shape[1:], dtype=input_.dtype, device=input_.device + ) + handle = torch.distributed.reduce_scatter_tensor( + output, input_.contiguous(), group=process_group, async_op=async_op + ) + return output, handle + + +# Raw operation, does not support autograd, but does support async +def all_reduce_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + input_ = input_.contiguous() + handle = torch.distributed.all_reduce(input_, group=process_group, async_op=async_op) + return input_, handle + + +class AllGatherFunc(torch.autograd.Function): + """Gather the input from sequence parallel region and concatenate.""" + + @staticmethod + def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor: + ctx.process_group = process_group + output, _ = all_gather_raw(input_, process_group) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + grad_input, _ = reduce_scatter_raw(grad_output, ctx.process_group) + return grad_input, None + + +# Supports autograd, but does not support async +all_gather = AllGatherFunc.apply + + +class ReduceScatterFunc(torch.autograd.Function): + """Reduce scatter the input from the sequence parallel region and concatenate.""" + + @staticmethod + def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor: + ctx.process_group = process_group + output, _ = reduce_scatter_raw(input_, process_group) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + grad_input, _ = all_gather_raw(grad_output, ctx.process_group) + return grad_input, None + + +# Supports autograd, but does not support async +reduce_scatter = ReduceScatterFunc.apply + + +class AllReduceFunc(torch.autograd.Function): + """Gather the input from sequence parallel region and concatenate.""" + + @staticmethod + def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor: + ctx.process_group = process_group + output, _ = all_reduce_raw(input_, process_group) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + return grad_output, None + + +# Supports autograd, but does not support async +all_reduce = AllReduceFunc.apply + + +def sync_shared_params(model: torch.nn.Module, process_group: ProcessGroup): + # We want to iterate over parameters with _shared_params=True in the same order, + # as different ranks might have different number of parameters (e.g., only rank 0 has bias). + pamams_shared = { + name: p for name, p in model.named_parameters() if getattr(p, "_shared_params", False) + } + for _, p in sorted(pamams_shared.items()): + with torch.no_grad(): + # Broadcast needs src to be global rank, not group rank + torch.distributed.broadcast( + p, src=torch.distributed.get_global_rank(process_group, 0), group=process_group + ) + + +# Ref: https://github.com/NVIDIA/Megatron-LM/blob/52e636888cccc41e931251c417a7181fc36de926/megatron/optimizer/optimizer.py#L256 +def allreduce_sequence_parallel_grad(model: torch.nn.Module, process_group: ProcessGroup): + # We want to iterate over parameters with _sequence_parallel=True in the same order, + # as different ranks might have different number of parameters (e.g., only rank 0 has bias). + params_seqparallel = { + name: p for name, p in model.named_parameters() if getattr(p, "_sequence_parallel", False) + } + grads = [p.grad for _, p in sorted(params_seqparallel.items())] + if grads: + with torch.no_grad(): + coalesced = torch._utils._flatten_dense_tensors(grads) + torch.distributed.all_reduce(coalesced, group=process_group) + for buf, synced in zip(grads, torch._utils._unflatten_dense_tensors(coalesced, grads)): + buf.copy_(synced) + + +def get_dim_for_local_rank(dim: int, world_size: int, local_rank: int, multiple_of: int = 1) -> int: + """Get the dim for the local rank derived from splitting dim on world_size processes. + + The split may not be even across the world_size processes. + """ + multiple = dim // multiple_of + div = multiple // world_size + mod = multiple % world_size + local_multiple = div + int(local_rank < mod) + return local_multiple * multiple_of diff --git a/lib/python3.12/site-packages/flash_attn/utils/generation.py b/lib/python3.12/site-packages/flash_attn/utils/generation.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9120c386596f25b544391af10dc479cf00c822 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/utils/generation.py @@ -0,0 +1,740 @@ +# Copyright (c) 2023, Tri Dao. +# Adapted from https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/text_generation/forward_step.py#L31 +import gc +import time +from collections import namedtuple +from dataclasses import dataclass, field +from functools import partial +from typing import Callable, Optional, Sequence, Union + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch import Tensor +from torch.profiler import ProfilerActivity, profile, record_function + +try: + from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput +except ImportError: + GreedySearchDecoderOnlyOutput = namedtuple("GreedySearchDecoderOnlyOutput", ["sequences", "scores"]) + SampleDecoderOnlyOutput = namedtuple("SampleDecoderOnlyOutput", ["sequences", "scores"]) + + +@dataclass +class InferenceParams: + """Inference parameters that are passed to the main model in order + to efficienly calculate and store the context during inference.""" + + max_seqlen: int + max_batch_size: int + seqlen_offset: int = 0 + batch_size_offset: int = 0 + key_value_memory_dict: dict = field(default_factory=dict) + lengths_per_sample: Optional[Tensor] = None + + def reset(self, max_seqlen, max_batch_size): + self.max_seqlen = max_seqlen + self.max_batch_size = max_batch_size + self.seqlen_offset = 0 + if self.lengths_per_sample is not None: + self.lengths_per_sample.zero_() + + +# https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/text_generation/sampling.py +# https://github.com/huggingface/transformers/blob/a44985b41cfa2de48a5e1de7f1f93b7483da25d1/src/transformers/generation/logits_process.py#L231 +def modify_logits_for_top_k_filtering(logits, top_k): + """Set the logits for none top-k values to -inf. Done in-place.""" + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits.masked_fill_(indices_to_remove, float("-Inf")) + + +# https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/text_generation/sampling.py +# https://github.com/huggingface/transformers/blob/a44985b41cfa2de48a5e1de7f1f93b7483da25d1/src/transformers/generation/logits_process.py#L170 +def modify_logits_for_top_p_filtering(logits, top_p): + """Set the logits for none top-p values to -inf. Done in-place.""" + if top_p <= 0.0 or top_p >= 1.0: + return + # First sort and calculate cumulative sum of probabilities. + sorted_logits, sorted_indices = torch.sort(logits, descending=False) + cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + # Remove tokens with cumulative top_p above the threshold (token with 0 are kept) + sorted_indices_to_remove = cumulative_probs <= (1 - top_p) + # scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter( + 1, sorted_indices, sorted_indices_to_remove + ) + logits.masked_fill_(indices_to_remove, float("-inf")) + + +def sample(logits, top_k=1, top_p=0.0, temperature=1.0): + """Sample from top-k logits. + Arguments: + logits: Tensor of shape (batch_size, vocab_size) + """ + if top_k == 1: # Short-circuit for greedy decoding + return logits.argmax(dim=-1) + else: + if top_p > 0.0: + assert top_p <= 1.0, "top-p should be in (0, 1]." + if top_k > 0: + top_k = min(top_k, logits.size(-1)) # Safety check + logits_top, indices = torch.topk(logits, top_k, dim=-1) + if temperature != 1.0: + logits_top /= temperature + modify_logits_for_top_p_filtering(logits_top, top_p) + return indices[ + torch.arange(indices.shape[0], device=indices.device), + torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze(dim=-1), + ] + else: + # Clone so that when we modify for top_p we don't change the original logits + logits_top = logits / temperature if temperature != 1.0 else logits.clone() + modify_logits_for_top_p_filtering(logits_top, top_p) + return torch.multinomial(torch.softmax(logits_top, dim=-1), num_samples=1).squeeze( + dim=-1 + ) + + +@torch.inference_mode() +def decode( + input_ids, + model, + max_length, + top_k=1, + top_p=0.0, + temperature=1.0, + eos_token_id=None, + teacher_outputs=None, + vocab_size=None, + tensor_parallel=1, + cg=False, + enable_timing=False, +): + """Decoding, either greedy or with top-k or top-p sampling. + If top-k = 0, don't limit the number of candidates (pure sampling). + Top-k and top-p can be used together. If top_k > 0 and top_p > 0, then top-k is applied first, + then top-p. + We assume that all sequences in the same batch have the same length. + + Arguments: + input_ids: (batch, seq_len) + max_length: int + teacher_outputs (optional): (batch, seq_len). If provided, instead of sampling from the + logits, the next token is taken from the teacher_outputs. Useful for testing. + Returns: GreedySearchDecoderOnlyOutput or SampleDecoderOnlyOutput, with the following fields: + sequences: (batch, max_length) + scores: tuples of (batch, vocab_size) + """ + batch_size, seqlen_og = input_ids.shape + teacher_output_len = teacher_outputs.shape[1] if teacher_outputs is not None else 0 + if cg: + if not hasattr(model, "_decoding_cache"): + model._decoding_cache = None + model._decoding_cache = update_graph_cache( + model, + model._decoding_cache, + batch_size, + seqlen_og, + max_length, + tensor_parallel=tensor_parallel, + ) + inference_params = model._decoding_cache.inference_params + inference_params.reset(max_length, batch_size) + else: + inference_params = InferenceParams(max_seqlen=max_length, max_batch_size=batch_size) + + def get_logits(input_ids, inference_params): + decoding = inference_params.seqlen_offset > 0 + if decoding: + position_ids = torch.full( + (batch_size, 1), + inference_params.seqlen_offset, + dtype=torch.long, + device=input_ids.device, + ) + else: + position_ids = None + if not cg or not decoding: + logits = model( + input_ids, + position_ids=position_ids, + inference_params=inference_params, + num_last_tokens=1, + ).logits.squeeze(dim=1) + else: + logits = model._decoding_cache.run( + input_ids, position_ids, inference_params.seqlen_offset + ).squeeze(dim=1) + return logits[..., :vocab_size] if vocab_size is not None else logits + + def sample_tokens(logits, inference_params): + if teacher_outputs is None or teacher_output_len <= inference_params.seqlen_offset: + token = sample(logits, top_k=top_k, top_p=top_p, temperature=temperature) + else: + token = teacher_outputs[:, inference_params.seqlen_offset] + # return rearrange(token, "b -> b 1") + return token.unsqueeze(1) + + def should_stop(current_token, inference_params): + if inference_params.seqlen_offset == 0: + return False + if eos_token_id is not None and (current_token == eos_token_id).all(): + return True + if inference_params.seqlen_offset >= max_length - 1: + return True + return False + + start = torch.cuda.Event(enable_timing=enable_timing) + end = torch.cuda.Event(enable_timing=enable_timing) + + if enable_timing: + if tensor_parallel > 1: + torch.distributed.barrier() + start.record() + scores, sequences = [], [input_ids] + while not should_stop(sequences[-1], inference_params): + scores.append(get_logits(sequences[-1], inference_params)) + inference_params.seqlen_offset += sequences[-1].shape[1] + sequences.append(sample_tokens(scores[-1], inference_params)) + if enable_timing: + end.record() + if tensor_parallel > 1: + torch.distributed.barrier() + torch.cuda.synchronize() + print(f"Prompt processing + decoding time: {(start.elapsed_time(end)):.0f}ms") + output_cls = GreedySearchDecoderOnlyOutput if top_k == 1 else SampleDecoderOnlyOutput + return output_cls(sequences=torch.cat(sequences, dim=1), scores=tuple(scores)) + + +def sample_speculative(logits, logits_draft, tokens_draft, top_k=1, top_p=0.0, temperature=1.0): + """Algorithm 1 from [1] + [1] Fast Inference from Transformers via Speculative Decoding + Yaniv Leviathan, Matan Kalman, Yossi Matias + https://arxiv.org/abs/2211.17192 + + Arguments: + logits: Tensor of shape (batch_size, seqlen + 1, vocab_size) + logits_draft: Tensor of shape (batch_size, seqlen, vocab_size) + tokens_draft: Tensor of shape (batch_size, seqlen) + Return: + tokens: Tensor of shape (batch_size, seqlen + 1) + num_generated_tokens: Tensor of shape (batch_size), with value in [1, seqlen + 1]. + For each sequence in the batch, the number of valid tokens that were sampled by + speculative sampling. + """ + batch, seqlen_p_1, vocab_size = logits.shape + seqlen = seqlen_p_1 - 1 + assert logits_draft.shape == (batch, seqlen, vocab_size) + assert tokens_draft.shape == (batch, seqlen) + assert tokens_draft.dtype in [torch.int64, torch.int32] + # TODO: if top_k = 1 we can simplify things and only work with indices + if top_p > 0.0: + assert top_p <= 1.0, "top-p should be in (0, 1]." + # Clone so that when we modify for top_p we don't change the original logits + logits = logits / temperature if temperature != 1.0 else logits.clone() + logits_draft = logits_draft / temperature if temperature != 1.0 else logits_draft.clone() + if top_k > 0: + top_k = min(top_k, logits.size(-1)) # Safety check + modify_logits_for_top_k_filtering(logits, top_k) + modify_logits_for_top_k_filtering(logits_draft, top_k) + modify_logits_for_top_p_filtering(logits, top_p) + modify_logits_for_top_p_filtering(logits_draft, top_p) + probs = torch.softmax(logits, dim=-1) + probs_draft = torch.softmax(logits_draft, dim=-1) + gather = lambda probs, tokens: rearrange( + probs.gather(dim=-1, index=rearrange(tokens, "... -> ... 1")), "... 1 -> ..." + ) + # (batch, seqlen) + accepted = torch.rand(batch, seqlen, device=probs.device) * gather( + probs_draft, tokens_draft + ) <= gather(probs[:, :-1], tokens_draft) + accepted_all = accepted.all(dim=-1) + # (batch,) + first_rejected_idx = torch.where(accepted_all, seqlen, accepted.int().argmin(dim=-1)) + probs_diff = torch.clamp(probs[:, :-1] - probs_draft, min=0.0) + # torch.multinomial can deal with unnormalized probabilities + # probs_diff /= probs_diff.sum(dim=-1, keepdim=True) + resample_probs = torch.cat([probs_diff, probs[:, -1:]], dim=1) + resample_probs = rearrange( + resample_probs.gather(dim=1, index=repeat(first_rejected_idx, "b -> b 1 d", d=vocab_size)), + "b 1 d -> b d", + ) + resample = torch.multinomial(resample_probs, num_samples=1).squeeze(dim=-1) # (batch,) + tokens = F.pad(tokens_draft, (0, 1)) + tokens[:, first_rejected_idx] = resample + return tokens, first_rejected_idx + 1 + + +@torch.inference_mode() +def decode_speculative( + input_ids, + model, + model_draft, + max_length, + speculative_lookahead=3, + top_k=1, + top_p=0.0, + temperature=1.0, + eos_token_id=None, + vocab_size=None, + tensor_parallel=1, + cg=False, + enable_timing=False, + debug=False, +): + """ + TD: WIP, for my own understanding, lightly tested. Only support batch_size == 1 for now. + + Speculative decoding, either greedy or with top-k or top-p sampling. + If top-k = 0, don't limit the number of candidates (pure sampling). + Top-k and top-p can be used together. If top_k > 0 and top_p > 0, then top-k is applied first, + then top-p. + We assume that all sequences in the same batch have the same length. + + Arguments: + input_ids: (batch, seq_len) + max_length: int + Returns: GreedySearchDecoderOnlyOutput or SampleDecoderOnlyOutput, with the following fields: + sequences: (batch, max_length) + scores: tuples of (batch, vocab_size) + """ + batch_size, seqlen_og = input_ids.shape + assert batch_size == 1, "Speculative decoding implementation only supports batch_size=1" + assert eos_token_id is None, "Speculative decoding implementation doesn't support eos_token_id" + if cg: + if not hasattr(model_draft, "_decoding_cache"): + model_draft._decoding_cache = None + model_draft._decoding_cache = update_graph_cache( + model_draft, + model_draft._decoding_cache, + batch_size, + seqlen_og, + max_length, + # draft model needs to process either 1 or 2 tokens at a time + decoding_seqlens=(1, 2), + tensor_parallel=tensor_parallel, + ) + inference_params_draft = model_draft._decoding_cache.inference_params + inference_params_draft.reset(max_length, batch_size) + if not hasattr(model, "_decoding_cache"): + model._decoding_cache = None + model._decoding_cache = update_graph_cache( + model, + model._decoding_cache, + batch_size, + seqlen_og, + max_length, + decoding_seqlens=range(1, speculative_lookahead + 2), + tensor_parallel=tensor_parallel, + ) + inference_params = model._decoding_cache.inference_params + inference_params.reset(max_length, batch_size) + else: + inference_params_draft = InferenceParams(max_seqlen=max_length, max_batch_size=batch_size) + inference_params = InferenceParams(max_seqlen=max_length, max_batch_size=batch_size) + + def get_logits(input_ids, inference_params, model, num_last_tokens=1, cg=False): + decoding = inference_params.seqlen_offset > 0 + if decoding: + seqlen = input_ids.shape[1] + # if inference_params.lengths_per_sample is None: + # TODO: in the case of batched decoding where each sequence has a different length, + # we need to compute the position_ids for each sequence using lengths_per_sample + if True: + cache_seqlens = torch.full( + (input_ids.shape[0],), + inference_params.seqlen_offset, + dtype=torch.int32, + device=input_ids.device, + ) + else: + cache_seqlens = inference_params.lengths_per_sample + position_ids = cache_seqlens[:, None] + torch.arange( + seqlen, dtype=torch.long, device=input_ids.device + ) + else: + position_ids = None + if not cg or not decoding: + logits = model( + input_ids, + position_ids=position_ids, + inference_params=inference_params, + num_last_tokens=num_last_tokens, + ).logits + else: + # NOTE: careful, CUDA graph is set to have num_last_tokens=input_ids.shape[1]. + # This might not be compatible the num_last_tokens used here. + assert num_last_tokens <= input_ids.shape[1] + logits = model._decoding_cache.run( + input_ids, position_ids, inference_params.seqlen_offset + )[:, -num_last_tokens:] + return logits[..., :vocab_size] if vocab_size is not None else logits + + def sample_tokens(input_ids, get_logits_fn, inference_params, sample_fn, num_tokens=1): + """Sample `num_tokens` tokens from the model, given the previous logits. + Also return the logits of the sampled tokens. + Arguments: + input_ids: (batch, seqlen) + Return: + tokens: (batch, num_tokens) + scores: (batch, num_tokens), which contains @previous_logits and the logits of the next + (num_tokens - 1) tokens. The logits of the last token isn't computed. + """ + assert num_tokens >= 1 + sequences, scores = [input_ids], [] + for i in range(num_tokens): + scores.append(get_logits_fn(sequences[-1], inference_params)[:, -1]) + inference_params.seqlen_offset += sequences[-1].shape[1] + sequences.append(sample_fn(scores[-1]).unsqueeze(1)) + return torch.cat(sequences[1:], dim=1), torch.stack(scores, dim=1) + + sampling_kwargs = dict(top_k=top_k, top_p=top_p, temperature=temperature) + sample_fn = partial(sample, **sampling_kwargs) + get_logits_main = partial(get_logits, model=model, cg=cg) + get_logits_draft = partial(get_logits, model=model_draft, cg=cg) + sample_tokens_main = partial( + sample_tokens, + get_logits_fn=get_logits_main, + sample_fn=sample_fn, + inference_params=inference_params, + ) + sample_tokens_draft = partial( + sample_tokens, + get_logits_fn=get_logits_draft, + sample_fn=sample_fn, + inference_params=inference_params_draft, + ) + + if debug: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("gpt2") + if enable_timing: + if tensor_parallel > 1: + torch.distributed.barrier() + torch.cuda.synchronize() + start = time.time() + + sequences, scores = [input_ids], [] + num_main_model_calls = 0 + num_draft_tokens = 0 + num_accepted_tokens_history = [] + if seqlen_og >= max_length - 1: + # Don't do speculative sampling, just sample 1 token from the model + tokens, scores_new = sample_tokens_main(input_ids, num_tokens=1) + sequences.append(tokens) + scores.append(scores_new) + else: + # Sample from draft model, which produces @n_spec_tokens, and @model + # will then use to produce between 1 and 1 + @n_spec_tokens tokens. + # We want seqlen_og + 1 + @n_spec_tokens to be <= @max_length. + n_spec_tokens = min(speculative_lookahead, max_length - seqlen_og - 1) + tokens_draft, scores_draft = sample_tokens_draft(input_ids, num_tokens=n_spec_tokens) + num_draft_tokens += n_spec_tokens + if debug: + scores_draft_ref = model_draft( + torch.cat([input_ids, tokens_draft], dim=1), num_last_tokens=n_spec_tokens + 1 + ).logits + print((scores_draft - scores_draft_ref[:, :-1]).abs().max()) + + # Evaluate the draft tokens with the model + logits = get_logits_main( + torch.cat([input_ids, tokens_draft], dim=1), + inference_params, + num_last_tokens=n_spec_tokens + 1, + ) + num_main_model_calls += 1 + if debug: + logits_ref = model( + torch.cat([input_ids, tokens_draft], dim=1), num_last_tokens=n_spec_tokens + 1 + ).logits + print((logits - logits_ref).abs().max()) + # breakpoint() + tokens, num_generated_tokens = sample_speculative( + logits, scores_draft, tokens_draft, **sampling_kwargs + ) + num_accepted_tokens_history.append(num_generated_tokens - 1) + if debug: + print(tokens) + print(num_generated_tokens) + # breakpoint() + # TODO: we're using the fact that batch_size == 1 + # TODO: check eos_token_id + sequences.append(tokens[:1, : num_generated_tokens[0]]) + scores.append(logits[:1, : num_generated_tokens[0]]) + # Note that @model has not evaluated the last sampled token yet, so we'll need to pass + # that in the next time we call @model. + num_generated = num_generated_tokens[0].item() + inference_params.seqlen_offset = seqlen_og + num_generated - 1 + inference_params_draft.seqlen_offset = ( + inference_params.seqlen_offset - 1 + if num_generated > 1 + else inference_params.seqlen_offset + ) + if debug: + cur_ids = torch.cat([input_ids, sequences[-1]], dim=1) + scores_ref = model(cur_ids, num_last_tokens=num_generated_tokens[0].item() + 1).logits + print((scores[-1] - scores_ref[:, :-1]).abs().max()) + # breakpoint() + + while True: + # seqlen_offset is total length generated - 1 + if inference_params.seqlen_offset >= max_length - 1: + break + if inference_params.seqlen_offset >= max_length - 2: + # Don't do speculative sampling, just sample 1 token from the model + tokens, scores_new = sample_tokens_main(sequences[-1][:, -1:], num_tokens=1) + sequences.append(tokens) + scores.append(scores_new) + break + # Sample from draft model + n_spec_tokens = min( + speculative_lookahead, max_length - inference_params_draft.seqlen_offset - 2 + ) + # If the main model accepts all the draft tokens, plus it samples one new token, + # then at the next iteration the draft model need to evaluate the logits of the last draft + # token and the logits of the newly sampled token. So here we pass in the last 2 tokens + # of sequences[-1]. + # This exception is when the main model rejects all the draft tokens, in which case we + # will only have 1 token to pass in. + tokens_draft, scores_draft = sample_tokens_draft( + sequences[-1][:, -2:], num_tokens=n_spec_tokens + ) + num_draft_tokens += n_spec_tokens + if debug: + scores_draft_ref = model_draft( + torch.cat([cur_ids, tokens_draft], dim=1), num_last_tokens=n_spec_tokens + 1 + ).logits + print((scores_draft - scores_draft_ref[:, :-1]).abs().max()) + # breakpoint() + # Evaluate the draft tokens with the model + logits = get_logits_main( + torch.cat([sequences[-1][:, -1:], tokens_draft], dim=1), + inference_params, + num_last_tokens=n_spec_tokens + 1, + ) # (batch, n_spec_tokens + 1, vocab_size) + num_main_model_calls += 1 + if debug: + logits_ref = model( + torch.cat([cur_ids, tokens_draft], dim=1), num_last_tokens=n_spec_tokens + 1 + ).logits + print((logits - logits_ref).abs().max()) + # breakpoint() + tokens, num_generated_tokens = sample_speculative( + logits, scores_draft, tokens_draft, **sampling_kwargs + ) + num_accepted_tokens_history.append(num_generated_tokens - 1) + if debug: + print(tokens) + print(num_generated_tokens) + # breakpoint() + sequences.append(tokens[:1, : num_generated_tokens[0]]) + scores.append(logits[:1, : num_generated_tokens[0]]) + # We've evaluated 1 token from sequences[-1][:, -1:] above, plus + # num_generated_tokens[0].item() - 1 tokens from the draft model. + num_generated = num_generated_tokens[0].item() + inference_params.seqlen_offset += num_generated + inference_params_draft.seqlen_offset = ( + inference_params.seqlen_offset - 1 + if num_generated > 1 + else inference_params.seqlen_offset + ) + if debug: + cur_ids = torch.cat([cur_ids, sequences[-1]], dim=1) + scores_ref = model(cur_ids, num_last_tokens=num_generated_tokens[0].item() + 1).logits + print((scores[-1] - scores_ref[:, :-1]).abs().max()) + # breakpoint() + + if enable_timing: + if tensor_parallel > 1: + torch.distributed.barrier() + torch.cuda.synchronize() + print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms") + print(f"Number of calls to main model: {num_main_model_calls}") + print( + f"Acceptance rate: {torch.cat(num_accepted_tokens_history).sum().item() / num_draft_tokens * 100:.2f}%" + ) + sequences = torch.cat(sequences, dim=1) + scores = torch.cat(scores, dim=1) + if debug: + scores_ref = model(sequences).logits + print((scores - scores_ref[:, seqlen_og - 1 : -1]).abs().max()) + output_cls = GreedySearchDecoderOnlyOutput if top_k == 1 else SampleDecoderOnlyOutput + return output_cls(sequences=sequences, scores=scores) + + +class GenerationMixin: + def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs): + raise NotImplementedError + + def generate( + self, + input_ids, + max_length, + top_k=1, + top_p=0.0, + temperature=1.0, + return_dict_in_generate=False, + output_scores=False, + **kwargs, + ): + output = decode( + input_ids, self, max_length, top_k=top_k, top_p=top_p, temperature=temperature, **kwargs + ) + if not output_scores: + output.scores = None + return output if return_dict_in_generate else output.sequences + + +def allocate_inference_cache( + max_batch_size, + max_seqlen, + nheads, + headdim, + layers: Union[int, Sequence], + device, + dtype=torch.float16, +): + assert dtype in [torch.float16, torch.bfloat16, torch.float32] + kv_cache_shape = (max_batch_size, max_seqlen, 2, nheads, headdim) + if isinstance(layers, int): + layers = range(layers) + return {i: torch.empty(kv_cache_shape, device=device, dtype=dtype) for i in layers} + + +@dataclass +class DecodingCGCache: + max_batch_size: int = 0 + max_seqlen: int = 0 + device = None + dtype = None + callables: dict = field(default_factory=dict) + mempool = None + inference_params: Optional[InferenceParams] = None + run: Optional[Callable] = None + + +@torch.inference_mode() +def update_graph_cache( + model, + cache, + batch_size, + seqlen_og, + max_seqlen, + decoding_seqlens=(1,), + tensor_parallel=1, + dtype=None, + n_warmups=2, +): + if cache is None: + cache = DecodingCGCache() + param_example = next(iter(model.parameters())) + device = param_example.device + if dtype is None: + dtype = param_example.dtype + if ( + (device, dtype) != (cache.device, cache.dtype) + or batch_size > cache.max_batch_size + or max_seqlen > cache.max_seqlen + ): # Invalidate the cache + cache.callables = {} + cache.mempool = None + cache.inference_params = None + gc.collect() + cache.device, cache.dtype = device, dtype + cache.max_batch_size, cache.max_seqlen = batch_size, max_seqlen + if hasattr(model, "allocate_inference_cache"): + inf_cache = model.allocate_inference_cache(batch_size, max_seqlen, dtype) + else: + headdim = getattr( + model.config, + "head_dim", + model.config.hidden_size // model.config.num_attention_heads, + ) + inf_cache = allocate_inference_cache( + batch_size, + max_seqlen, + model.config.num_attention_heads // tensor_parallel, + headdim, + model.config.num_hidden_layers, + device, + dtype, + ) + lengths_per_sample = torch.full((batch_size,), seqlen_og, dtype=torch.int32, device=device) + cache.inference_params = InferenceParams( + max_seqlen=max_seqlen, + max_batch_size=batch_size, + seqlen_offset=seqlen_og, + key_value_memory_dict=inf_cache, + lengths_per_sample=lengths_per_sample, + ) + cache.mempool = torch.cuda.graphs.graph_pool_handle() + for decoding_seqlen in decoding_seqlens: + if (batch_size, decoding_seqlen) not in cache.callables: + cache.callables[batch_size, decoding_seqlen] = capture_graph( + model, + cache.inference_params, + batch_size, + max_seqlen, + decoding_seqlen=decoding_seqlen, + mempool=cache.mempool, + n_warmups=n_warmups, + ) + + def dispatch(input_ids, position_ids, seqlen): + batch_size, decoding_seqlen = input_ids.shape[:2] + return cache.callables[batch_size, decoding_seqlen](input_ids, position_ids, seqlen) + + cache.run = dispatch + cache.inference_params.seqlen_offset = 0 # Reset so it's not confusing + return cache + + +def capture_graph( + model, inference_params, batch_size, max_seqlen, decoding_seqlen=1, mempool=None, n_warmups=2 +): + device = next(iter(model.parameters())).device + input_ids = torch.full((batch_size, decoding_seqlen), 0, dtype=torch.long, device=device) + position_ids = torch.full((batch_size, decoding_seqlen), 0, dtype=torch.long, device=device) + seqlen_offset_og = inference_params.seqlen_offset + inference_params.seqlen_offset = max_seqlen - decoding_seqlen + inference_params.lengths_per_sample[:] = inference_params.seqlen_offset + + # Warmup before capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(n_warmups): + logits = model( + input_ids, + position_ids=position_ids, + inference_params=inference_params, + num_last_tokens=decoding_seqlen, + ).logits + s.synchronize() + # This might be needed for correctness if we run with NCCL_GRAPH_MIXING_SUPPORT=0, + # which requires that graph launch and non-captured launch to not overlap (I think, + # that's how I interpret the documentation). I'm not sure if this is required. + if torch.distributed.is_initialized(): + torch.distributed.barrier() + torch.cuda.current_stream().wait_stream(s) + # Captures the graph + # To allow capture, automatically sets a side stream as the current stream in the context + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=mempool): + logits = model( + input_ids, + position_ids=position_ids, + inference_params=inference_params, + num_last_tokens=decoding_seqlen, + ).logits + + def run(new_input_ids, new_position_ids, seqlen): + inference_params.lengths_per_sample[:] = seqlen + input_ids.copy_(new_input_ids) + position_ids.copy_(new_position_ids) + graph.replay() + return logits.clone() + + inference_params.seqlen_offset = seqlen_offset_og + return run diff --git a/lib/python3.12/site-packages/flash_attn/utils/pretrained.py b/lib/python3.12/site-packages/flash_attn/utils/pretrained.py new file mode 100644 index 0000000000000000000000000000000000000000..40e76bd2692335c7f474f6b6479be67eb95f8d20 --- /dev/null +++ b/lib/python3.12/site-packages/flash_attn/utils/pretrained.py @@ -0,0 +1,79 @@ +import os +from functools import partial + +import torch +from safetensors.torch import load_file as safe_load_file +from transformers.utils import ( + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, +) +from transformers.utils.hub import cached_file, get_checkpoint_shard_files + + +def state_dict_from_pretrained(model_name, device=None, dtype=None): + # If not fp32, then we don't want to load directly to the GPU + mapped_device = "cpu" if dtype not in [torch.float32, None] else device + is_sharded = False + load_safe = False + resolved_archive_file = None + + weights_path = os.path.join(model_name, WEIGHTS_NAME) + weights_index_path = os.path.join(model_name, WEIGHTS_INDEX_NAME) + safe_weights_path = os.path.join(model_name, SAFE_WEIGHTS_NAME) + safe_weights_index_path = os.path.join(model_name, SAFE_WEIGHTS_INDEX_NAME) + + if os.path.isfile(weights_path): + resolved_archive_file = cached_file( + model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False + ) + elif os.path.isfile(weights_index_path): + resolved_archive_file = cached_file( + model_name, WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False + ) + is_sharded = True + elif os.path.isfile(safe_weights_path): + resolved_archive_file = cached_file( + model_name, SAFE_WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False + ) + load_safe = True + elif os.path.isfile(safe_weights_index_path): + resolved_archive_file = cached_file( + model_name, SAFE_WEIGHTS_INDEX_NAME, _raise_exceptions_for_missing_entries=False + ) + is_sharded = True + load_safe = True + else: # Try loading from HF hub instead of from local files + resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, + _raise_exceptions_for_missing_entries=False) + if resolved_archive_file is None: + resolved_archive_file = cached_file(model_name, WEIGHTS_INDEX_NAME, + _raise_exceptions_for_missing_entries=False) + if resolved_archive_file is not None: + is_sharded = True + + if resolved_archive_file is None: + raise EnvironmentError(f"Model name {model_name} was not found.") + + if load_safe: + loader = partial(safe_load_file, device=mapped_device) + else: + loader = partial(torch.load, map_location=mapped_device) + + if is_sharded: + # resolved_archive_file becomes a list of files that point to the different + # checkpoint shards in this case. + resolved_archive_file, sharded_metadata = get_checkpoint_shard_files( + model_name, resolved_archive_file + ) + state_dict = {} + for sharded_file in resolved_archive_file: + state_dict.update(loader(sharded_file)) + else: + state_dict = loader(resolved_archive_file) + # Convert dtype before moving to GPU to save memory + if dtype is not None: + state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()} + state_dict = {k: v.to(device=device) for k, v in state_dict.items()} + return state_dict diff --git a/lib/python3.12/site-packages/liger_kernel/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d27de3efc374bab8ee27170c22525684d2c1603 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/__pycache__/env_report.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/__pycache__/env_report.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b8fce447f7e976c396cead6fe686713ec0216f8 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/__pycache__/env_report.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30d6d1d51a6bfa42d3a418dcaacad9b949a1da4b Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d88857f3cd29d3eb27b9b7449ac755f59f72947 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/auto_model.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/auto_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f19e7ca51c13266bb9e334d0541dd3e560039822 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/auto_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/cross_entropy.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/cross_entropy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9537bf273ba1f6a470773cd824a6d6996c76839e Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/cross_entropy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/dyt.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/dyt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1a4172821960ac6db69b9b38fe3dd1454f1068b Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/dyt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fsdp.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fsdp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e34c1c4b27da86660228561556a6db6515ad33a Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fsdp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/functional.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/functional.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adcb8faa5b9c92f3a064f9db364b29d431bf779d Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/functional.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_cross_entropy.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_cross_entropy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ecc2b79fe64224b313f51dfa16b9a7f1a83b094 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_cross_entropy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_jsd.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_jsd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c99a59cd0537e5610901dc8815256fba505dbfe7 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/fused_linear_jsd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/geglu.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/geglu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11e7ad954ee3f04ac80b256e08362de4a57baad3 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/geglu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/gema3_rms.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/gema3_rms.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6566bbbb6ac3fcca2d5b8f9ef123100c5156df8f Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/gema3_rms.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/group_norm.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/group_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4e75cea4e42d91d7b4eb9ce8d124c142462e74a Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/group_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/grpo_loss.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/grpo_loss.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9c4c424f1245d41f1d74f5823da58da4143a2f1 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/grpo_loss.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/jsd.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/jsd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e265148da3a37cedf67435cc941817b982f47882 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/jsd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/kl_div.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/kl_div.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa2b2eb30ee57676d76b2a6a7ecd0488262d423d Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/kl_div.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/layer_norm.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/layer_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2af339d83d874c893301cd479792874bddd93f8 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/layer_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/monkey_patch.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/monkey_patch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79b9361800640b018a50fc10c85aeb227a860f9c Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/monkey_patch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/qwen2vl_mrope.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/qwen2vl_mrope.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efcba2b9f08dcef71085722dea079786fe4b36b5 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/qwen2vl_mrope.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rms_norm.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rms_norm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f3c88d3fd7c1bd5a16a44379a29206f00efd65a Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rms_norm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rope.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rope.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d817bc49230fc808ae911db38f1101651cfcd338 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/rope.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/sparsemax.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/sparsemax.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2ef0be6e118b8f6c2592d326ce56ce9857832f8 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/sparsemax.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/swiglu.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/swiglu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5d47e63151dd503a94b3c3502479fb1e8510b581 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/swiglu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/trainer_integration.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/trainer_integration.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a27b5f5862a7b4310828b3173b5b01015fbf22d Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/trainer_integration.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/tvd.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/tvd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77318c5e63c919a4dc75d1ffb9c1ba221f1dbd7b Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/__pycache__/tvd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..990cf2f13b749b62d3098996c5882811f5c91ef0 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72c4d5b6a3b2b413901c2baf98510bb0eb72a8d8 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma2.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f7b9ec59e58512cfe22d600e0a31ad41a76128c Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma3.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94b742a431b029555bb8fb09f02117e110cd55a5 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/gemma3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/glm4.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/glm4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11a11e7720387d2720bd7b4be2f47451409f43aa Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/glm4.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llama.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b5d904781237aa00007b92d6ad3988850c13b6d Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llama.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llava.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llava.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ae472a9920e5aae55b9054ac77343051e0d7811 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/llava.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/loss_utils.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/loss_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d562211a1c533d79bf6bce0a64ffdb9682f296d1 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/loss_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mistral.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mistral.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43b7045042b09a6a78a6ec65594db175ce0864eb Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mistral.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mixtral.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mixtral.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfd44029c40c185a9c2ffac1dc3ee407be99fdb6 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mixtral.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mllama.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mllama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..493b474216dace0e680e45631bfbb64d6f8305e2 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/mllama.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/olmo2.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/olmo2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68c7602a1503b18b9f74dd3e5e04c7b1c2095146 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/olmo2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/paligemma.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/paligemma.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9490b393a438102b9ba058cd7c07be3c81bbbb39 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/paligemma.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/phi3.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/phi3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b53fcfbe01e802bc9c3d790b66d53f5748e68df0 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/phi3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1132bf276f16a4bfd28eb589a9df121b2e469ce7 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_5_vl.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_5_vl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93c6845a429263c407062bec9ef188fc3d846263 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_5_vl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_vl.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_vl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58026182753b62753928d8eb4ceca468d48c4c42 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen2_vl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a8b038cf1b79de6b1b6341d27f7e3fbd75834eb Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3_moe.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8deaf8d960d16cadd233d5da04904218ca4f45a0 Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/transformers/model/__pycache__/qwen3_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/llava.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/llava.py new file mode 100644 index 0000000000000000000000000000000000000000..de9f6377a0f4f47cff9c5af4dd40df9636d4c575 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/llava.py @@ -0,0 +1,361 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast +from transformers.utils import is_torchdynamo_compiling +from transformers.utils.deprecation import deprecate_kwarg + +from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss + + +def lce_forward_deprecated( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + vision_feature_layer: Optional[int] = None, + vision_feature_select_strategy: Optional[str] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, +) -> Union[Tuple, LlavaCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + num_logits_to_keep (`int`, *optional*): + Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, LlavaForConditionalGeneration + + >>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf") + >>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf") + + >>> prompt = "USER: \nWhat's the content of the image? ASSISTANT:" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_new_tokens=15) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "USER: \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed" + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + vision_feature_layer = ( + vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer + ) + vision_feature_select_strategy = ( + vision_feature_select_strategy + if vision_feature_select_strategy is not None + else self.config.vision_feature_select_strategy + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and input_ids.shape[1] != 1: + image_outputs = self.vision_tower(pixel_values, output_hidden_states=True) + # this is not memory efficient at all (output_hidden_states=True) will save all the hidden stated. + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] + + if vision_feature_select_strategy == "default": + selected_image_feature = selected_image_feature[:, 1:] + elif vision_feature_select_strategy == "full": + selected_image_feature = selected_image_feature + else: + raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") + + image_features = self.multi_modal_projector(selected_image_feature) + inputs_embeds = inputs_embeds.to(image_features.dtype) + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, inputs_embeds, input_ids, attention_mask, labels + ) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat((extended_attention_mask, attention_mask[:, -target_length:]), dim=1) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + # TODO: @raushan retain only the new behavior after v4.47 + elif image_features is not None: + n_image_tokens = (input_ids == self.config.image_token_index).sum().item() + n_image_features = image_features.shape[0] * image_features.shape[1] + + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + special_image_mask = ( + (input_ids == self.config.image_token_index).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + ) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model.model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = outputs[0] + + loss = None + logits = None + + if self.training and (labels is not None): + # Shift so that tokens < n predict n + if attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -(hidden_states.shape[1] - 1) :].to(hidden_states.device) + shift_hidden_states = hidden_states[..., :-1, :][ + shift_attention_mask.to(hidden_states.device) != 0 + ].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_hidden_states = hidden_states[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce(self.language_model.lm_head.weight, shift_hidden_states, shift_labels) + + if not return_dict: + # NOTE: This part has not been tested. + output = outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") +def lce_forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + vision_feature_layer: Optional[int] = None, + vision_feature_select_strategy: Optional[str] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + image_sizes: torch.Tensor = None, + **lm_kwargs, +) -> Union[Tuple, LlavaCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, LlavaForConditionalGeneration + + >>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf") + >>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf") + + >>> prompt = "USER: \nWhat's the content of the image? ASSISTANT:" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_new_tokens=15) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "USER: \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed" + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + vision_feature_layer = ( + vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer + ) + vision_feature_select_strategy = ( + vision_feature_select_strategy + if vision_feature_select_strategy is not None + else self.config.vision_feature_select_strategy + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + image_sizes=image_sizes, + ) + + special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel(): + n_image_tokens = (input_ids == self.config.image_token_index).sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model.model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + **lm_kwargs, + ) + hidden_states = outputs[0] + + loss = None + logits = None + + if self.training and (labels is not None): + # Shift so that tokens < n predict n + if attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -(hidden_states.shape[1] - 1) :].to(hidden_states.device) + shift_hidden_states = hidden_states[..., :-1, :][ + shift_attention_mask.to(hidden_states.device) != 0 + ].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() + else: + shift_hidden_states = hidden_states[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce( + self.language_model.lm_head.weight, + shift_hidden_states.view(-1, shift_hidden_states.size(-1)), + shift_labels.view(-1).to(shift_hidden_states.device), + ) + + if not return_dict: + # NOTE: This part has not been tested. + output = outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/loss_utils.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/loss_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..279fa160df0f3992f0327b1f1cf2719788304544 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/loss_utils.py @@ -0,0 +1,64 @@ +from typing import Optional + +import torch +import torch.nn as nn + +import liger_kernel.transformers.functional as F + + +def fixed_fused_linear_cross_entropy( + hidden_states: torch.Tensor, + lm_head_weight: torch.Tensor, + target: torch.Tensor, + num_items_in_batch: Optional[int] = None, + ignore_index: int = -100, + final_logit_softcapping: Optional[float] = None, + **kwargs, +): + reduction = "sum" if num_items_in_batch is not None else "mean" + loss = F.liger_fused_linear_cross_entropy( + hidden_states, + lm_head_weight, + target, + reduction=reduction, + ignore_index=ignore_index, + softcap=final_logit_softcapping, + ) + if reduction == "sum": + loss = loss / num_items_in_batch + + return loss + + +def LigerForCausalLMLoss( + hidden_states, + lm_head_weight, + labels, + hidden_size: int, + num_items_in_batch: Optional[int] = None, + ignore_index: int = -100, + shift_labels: Optional[torch.Tensor] = None, + final_logit_softcapping: Optional[float] = None, + **kwargs, +): + # Skip upcast since intermediate values for the loss are all fp32 in kernel + if shift_labels is None: + # Shift so that token < n predict n + labels = nn.functional.pad(labels, (0, 1), value=ignore_index) + shift_labels = labels[..., 1:].contiguous() + + # Flatten the tokens + hidden_states = hidden_states.view(-1, hidden_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(hidden_states.device) + loss = fixed_fused_linear_cross_entropy( + hidden_states, + lm_head_weight, + shift_labels, + num_items_in_batch, + ignore_index, + final_logit_softcapping, + **kwargs, + ) + return loss diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/mixtral.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/mixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..3132821943043514cfc61cf73db6053addb92282 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/mixtral.py @@ -0,0 +1,269 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from torch.nn import CrossEntropyLoss +from transformers.modeling_outputs import MoeCausalLMOutputWithPast +from transformers.models.mixtral.modeling_mixtral import load_balancing_loss_func +from transformers.utils.deprecation import deprecate_kwarg + +from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss +from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss + + +def lce_forward_deprecated( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, +) -> Union[Tuple, MoeCausalLMOutputWithPast]: + r""" + Copy paste Mixtral's forward from transformers v4.44.2 but replace torch cross entropy with liger fused linear cross entropy + + + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + loss = None + if self.training and (labels is not None): + shift_hidden_states = hidden_states[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size) + shift_labels = shift_labels.view(-1) + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels) + elif labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.weight, shift_labels) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) + + +@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") +# Ignore copy +def lce_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **loss_kwargs, +) -> Union[Tuple, MoeCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + kept_hidden_states = hidden_states[:, slice_indices, :] + + shift_labels = loss_kwargs.pop("shift_labels", None) + logits = None + loss = None + # if in training mode, don't materialize logits + if self.training and (labels is not None or shift_labels is not None): + loss = LigerForCausalLMLoss( + hidden_states=kept_hidden_states, + lm_head_weight=self.lm_head.weight, + labels=labels, + shift_labels=shift_labels, + hidden_size=self.config.hidden_size, + **loss_kwargs, + ) + + else: # if in inference mode materialize logits + logits = self.lm_head(kept_hidden_states) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func( + outputs.router_logits if return_dict else outputs[-1], + self.num_experts, + self.num_experts_per_tok, + attention_mask, + ) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=outputs.router_logits, + ) diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/mllama.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..3cddf8bf723a4fb5035975af47b1e5cdabdac3b2 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/mllama.py @@ -0,0 +1,249 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.utils.deprecation import deprecate_kwarg + +from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss +from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss + + +def lce_forward_deprecated( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + cross_attention_states: Optional[torch.LongTensor] = None, + cross_attention_mask: Optional[torch.LongTensor] = None, + full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, +) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Copy paste mllama forward but replace torch cross entropy with liger fused linear cross entropy + + + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + num_logits_to_keep (`int`, *optional*): + Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + Returns: + Example: + ```python + >>> from transformers import AutoTokenizer, MllamaForCausalLM + >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision") + >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision") + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + cross_attention_states=cross_attention_states, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + + loss = None + logits = None + + if self.training and (labels is not None): + kept_hidden_states = hidden_states[:, -num_logits_to_keep:, :] + + shift_hidden_states = kept_hidden_states[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + # flatten tokens + shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size) + shift_labels = shift_labels.view(-1) + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels) + + else: + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") +def lce_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + cross_attention_states: Optional[torch.LongTensor] = None, + cross_attention_mask: Optional[torch.LongTensor] = None, + full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **loss_kwargs, +) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MllamaForCausalLM + + >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision") + >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision") + + >>> prompt = "If I had to write a haiku, it would be:" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) + >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + >>> print(result) + If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. + I love the idea of snowflakes gently falling, each one + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + cross_attention_states=cross_attention_states, + attention_mask=attention_mask, + position_ids=position_ids, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + kept_hidden_states = hidden_states[:, slice_indices, :] + + shift_labels = loss_kwargs.pop("shift_labels", None) + logits = None + loss = None + # if in training mode, don't materialize logits + if self.training and (labels is not None or shift_labels is not None): + loss = LigerForCausalLMLoss( + hidden_states=kept_hidden_states, + lm_head_weight=self.lm_head.weight, + labels=labels, + shift_labels=shift_labels, + hidden_size=self.config.hidden_size, + **loss_kwargs, + ) + + else: # if in inference mode materialize logits + logits = self.lm_head(kept_hidden_states) + if labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.vocab_size, + **loss_kwargs, + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/paligemma.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/paligemma.py new file mode 100644 index 0000000000000000000000000000000000000000..db72a745717786e9b8beb70510c029a419b758b5 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/paligemma.py @@ -0,0 +1,389 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache +from transformers.models.paligemma.modeling_paligemma import PaliGemmaCausalLMOutputWithPast +from transformers.utils import is_torchdynamo_compiling +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss + +logger = logging.get_logger(__name__) + + +def lce_forward_deprecated( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, +) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/PaliGemma-test-224px-hf") + >>> processor = AutoProcessor.from_pretrained("google/PaliGemma-test-224px-hf") + + >>> prompt = "answer en Where is the cow standing?" + >>> url = "https://huggingface.co/gv-hf/PaliGemma-test-224px-hf/resolve/main/cow_beach_1.png" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(text=prompt, images=image, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_length=30) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "answer en Where is the cow standing?\nbeach" + ```""" + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # the attention mask is turned 4d after, we keep track of the original one + input_attention_mask = attention_mask + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and input_ids.shape[1] != 1: + image_outputs = self.vision_tower(pixel_values.to(inputs_embeds.dtype)) + selected_image_feature = image_outputs.last_hidden_state + image_features = self.multi_modal_projector(selected_image_feature) + + if cache_position is None: + cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + inputs_embeds, attention_mask, labels, position_ids = self._merge_input_ids_with_image_features( + image_features, inputs_embeds, input_ids, attention_mask, labels, token_type_ids, cache_position + ) + + else: + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + if past_key_values is not None and pixel_values is not None and input_ids.shape[1] == 1: + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + # TODO @molbap this will only work for dynamic cache. + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where(first_layer_past_key_value.float().sum(-2) == 0) + + # Get the target length + target_seqlen = cache_position[-1] + 1 + extended_attention_mask = torch.ones( + (attention_mask.shape[0], target_seqlen - attention_mask.shape[1] + 1), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + # Filter out only the tokens that can be un-attended, this can happen + # if one uses PaliGemma+ Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + attention_mask = attention_mask.to(inputs_embeds.dtype) + outputs = self.language_model.model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + + loss = None + logits = None + + if self.training and (labels is not None): + shift_hidden_states = hidden_states[..., :-1, :] + shift_labels = labels[..., 1:] + + hidden_device = shift_hidden_states.device + + if attention_mask is not None: + # we use the input attention mask to shift the hidden_states and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -shift_hidden_states.shape[1] :].to(hidden_device) + shift_hidden_states = shift_hidden_states[shift_attention_mask.to(hidden_device) != 0].contiguous() + shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous() + else: + shift_hidden_states = shift_hidden_states.contiguous() + shift_labels = shift_labels.contiguous() + + # Flatten hidden state + shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size) + shift_labels = shift_labels.view(-1).to(hidden_device) + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce(self.language_model.lm_head.weight, shift_hidden_states, shift_labels) + + else: + logits = self.language_model.lm_head(hidden_states) + if labels is not None: + shift_logits = logits[..., :-1, :] + shift_labels = labels[..., 1:] + if input_attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + shift_attention_mask = input_attention_mask[..., 1:] + shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous() + else: + shift_logits = shift_logits.contiguous() + shift_labels = shift_labels.contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + + flat_logits = shift_logits.view(-1, self.config.vocab_size) + flat_labels = shift_labels.view(-1).to(shift_logits.device) + loss = loss_fct(flat_logits, flat_labels) + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return PaliGemmaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") +def lce_forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + token_type_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **lm_kwargs, +) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, PaliGemmaForConditionalGeneration + + >>> model = PaliGemmaForConditionalGeneration.from_pretrained("google/PaliGemma-test-224px-hf") + >>> processor = AutoProcessor.from_pretrained("google/PaliGemma-test-224px-hf") + + >>> prompt = "answer en Where is the cow standing?" + >>> url = "https://huggingface.co/gv-hf/PaliGemma-test-224px-hf/resolve/main/cow_beach_1.png" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_length=30) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "answer en Where is the cow standing?\nbeach" + ```""" + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if pixel_values is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" + ) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + is_training = token_type_ids is not None and labels is not None + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + 1 # Paligemma positions are 1-indexed + + # Merge text and images + if pixel_values is not None: + image_features = self.get_image_features(pixel_values) + + special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel(): + image_tokens_in_text = torch.sum(input_ids == self.config.image_token_index) + raise ValueError( + f"Number of images does not match number of special image tokens in the input text. " + f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} " + "tokens from image embeddings." + ) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + # mask out pad-token-ids in labels for BC + if labels is not None and self.pad_token_id in labels: + logger.warning_once( + "`labels` contains `pad_token_id` which will be masked with `config.ignore_index`. " + "You have to mask out `pad_token_id` when preparing `labels`, this behavior will be removed in v.4.46.", + ) + labels = torch.where(input_ids == self.pad_token_id, self.config.ignore_index, labels) + + causal_mask = self._update_causal_mask( + attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training + ) + + outputs = self.language_model.model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + logits_to_keep=logits_to_keep, + **lm_kwargs, + ) + + hidden_states = outputs[0] + + loss = None + logits = None + + if self.training and (labels is not None): + shift_hidden_states = hidden_states[..., :-1, :] + shift_labels = labels[..., 1:] + + hidden_device = shift_hidden_states.device + + if attention_mask is not None: + # we use the input attention mask to shift the hidden_states and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -shift_hidden_states.shape[1] :].to(hidden_device) + shift_hidden_states = shift_hidden_states[shift_attention_mask.to(hidden_device) != 0].contiguous() + shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous() + else: + shift_hidden_states = shift_hidden_states.contiguous() + shift_labels = shift_labels.contiguous() + + # Flatten hidden state + shift_hidden_states = shift_hidden_states.view(-1, self.config.text_config.hidden_size) + shift_labels = shift_labels.view(-1).to(hidden_device) + + lce = LigerFusedLinearCrossEntropyLoss() + loss = lce(self.language_model.lm_head.weight, shift_hidden_states, shift_labels) + else: + logits = self.language_model.lm_head(hidden_states) + if labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + shift_logits = logits[..., :-1, :] + shift_labels = labels[..., 1:] + if attention_mask is not None: + # we use the input attention mask to shift the logits and labels, because it is 2D. + # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft + shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device) + shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous() + shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous() + else: + shift_logits = shift_logits.contiguous() + shift_labels = shift_labels.contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + + flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size) + flat_labels = shift_labels.view(-1).to(shift_logits.device) + loss = loss_fct(flat_logits, flat_labels) + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return PaliGemmaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) diff --git a/lib/python3.12/site-packages/liger_kernel/transformers/model/qwen2_vl.py b/lib/python3.12/site-packages/liger_kernel/transformers/model/qwen2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..cbbef944db4e30d0c231276d4d276d60758933a9 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/transformers/model/qwen2_vl.py @@ -0,0 +1,204 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +from packaging import version +from torch.nn import CrossEntropyLoss +from transformers import __version__ as transformers_version +from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLCausalLMOutputWithPast + +from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss + + +def lce_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + **loss_kwargs, +) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]: + r""" + Copy paste Qwen2VL's forward but replace torch cross entropy with liger fused linear cross entropy + + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Qwen2VLForConditionalGeneration + + >>> model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": "What is shown in this image?"}, + ], + }, + ] + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos]) + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.get_dtype()) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + image_mask = ( + (input_ids == self.config.image_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + video_mask = ( + (input_ids == self.config.video_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + if version.parse(transformers_version) > version.parse("4.46.3"): + # NOTE: this bug fix for qwen2-vl is not applied until transformers 4.47.0 + # https://github.com/huggingface/transformers/issues/33401 + # While correct, this breaks equivalence with past versions of Qwen2-VL from + # transformers and leads to failed tests or users noticing differences in results. + # TODO: remove above conditional when liger drops support for transformers<4.47.0 + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # calculate RoPE index once per generation in the pre-fill stage only + if (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None: + position_ids, rope_deltas = self.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, attention_mask + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0 + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + + shift_labels = loss_kwargs.pop("shift_labels", None) + loss = None + logits = None + + if self.training and (labels is not None or shift_labels is not None): + loss = LigerForCausalLMLoss( + hidden_states=hidden_states, + lm_head_weight=self.lm_head.weight, + labels=labels, + shift_labels=shift_labels, + hidden_size=self.config.hidden_size, + **loss_kwargs, + ) + else: + logits = self.lm_head(hidden_states) + if labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=rope_deltas, + ) diff --git a/lib/python3.12/site-packages/liger_kernel/triton/__init__.py b/lib/python3.12/site-packages/liger_kernel/triton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d373966a9bb488ce448ec6989c919967382fe8c7 --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/triton/__init__.py @@ -0,0 +1 @@ +from liger_kernel.triton.monkey_patch import apply_liger_triton_cache_manager # noqa: F401 diff --git a/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f93f2cdcd5b868c78bf8c5a2c8bd312a24900c2d Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/monkey_patch.cpython-312.pyc b/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/monkey_patch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1b3c4e2da6e3ab839e51669a1f8048305b5934e Binary files /dev/null and b/lib/python3.12/site-packages/liger_kernel/triton/__pycache__/monkey_patch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/liger_kernel/triton/monkey_patch.py b/lib/python3.12/site-packages/liger_kernel/triton/monkey_patch.py new file mode 100644 index 0000000000000000000000000000000000000000..bac4a6a0d6a8fc74b56562bdba3c659e175c39ca --- /dev/null +++ b/lib/python3.12/site-packages/liger_kernel/triton/monkey_patch.py @@ -0,0 +1,40 @@ +import os +import random + +from triton.runtime.cache import FileCacheManager + + +class LigerTritonFileCacheManager(FileCacheManager): + def put(self, data, filename, binary=True) -> str: + if not self.cache_dir: + raise RuntimeError("Could not create or locate cache dir") + binary = isinstance(data, bytes) + if not binary: + data = str(data) + assert self.lock_path is not None + filepath = self._make_path(filename) + # Random ID to avoid any collisions + rnd_id = random.randint(0, 1000000) + # we use the PID incase a bunch of these around so we can see what PID made it + pid = os.getpid() + # use temp dir to be robust against program interruptions + temp_dir = os.path.join(self.cache_dir, f"tmp.pid_{pid}_{rnd_id}") + os.makedirs(temp_dir, exist_ok=True) + temp_path = os.path.join(temp_dir, filename) + + mode = "wb" if binary else "w" + with open(temp_path, mode) as f: + f.write(data) + # Replace is guaranteed to be atomic on POSIX systems if it succeeds + # so filepath cannot see a partial write + os.replace(temp_path, filepath) + os.removedirs(temp_dir) + return filepath + + +def apply_liger_triton_cache_manager(): + """ + Experimental feature to get around transient FileNotFoundError in triton compilation. + For more details please see https://github.com/triton-lang/triton/pull/4295 + """ + os.environ["TRITON_CACHE_MANAGER"] = "liger_kernel.triton.monkey_patch:LigerTritonFileCacheManager" diff --git a/lib/python3.12/site-packages/multidict/__init__.py b/lib/python3.12/site-packages/multidict/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3db9e199b191c665d5ca56e08b2521bb0f92874 --- /dev/null +++ b/lib/python3.12/site-packages/multidict/__init__.py @@ -0,0 +1,60 @@ +""" +Multidict implementation. + +HTTP Headers and URL query string require specific data structure: +multidict. It behaves mostly like a dict but it can have +several values for the same key. +""" + +from typing import TYPE_CHECKING + +from ._abc import MultiMapping, MutableMultiMapping +from ._compat import USE_EXTENSIONS + +__all__ = ( + "CIMultiDict", + "CIMultiDictProxy", + "MultiDict", + "MultiDictProxy", + "MultiMapping", + "MutableMultiMapping", + "getversion", + "istr", + "upstr", +) + +__version__ = "6.7.1" + + +if TYPE_CHECKING or not USE_EXTENSIONS: + from ._multidict_py import ( + CIMultiDict, + CIMultiDictProxy, + MultiDict, + MultiDictProxy, + getversion, + istr, + ) +else: + from collections.abc import ItemsView, KeysView, ValuesView + + from ._multidict import ( + CIMultiDict, + CIMultiDictProxy, + MultiDict, + MultiDictProxy, + _ItemsView, + _KeysView, + _ValuesView, + getversion, + istr, + ) + + MultiMapping.register(MultiDictProxy) + MutableMultiMapping.register(MultiDict) + KeysView.register(_KeysView) + ItemsView.register(_ItemsView) + ValuesView.register(_ValuesView) + + +upstr = istr diff --git a/lib/python3.12/site-packages/multidict/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/multidict/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f53680cd3a519d4fbacb9797bada4d88d54c7323 Binary files /dev/null and b/lib/python3.12/site-packages/multidict/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/multidict/__pycache__/_abc.cpython-312.pyc b/lib/python3.12/site-packages/multidict/__pycache__/_abc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8abaf97c47143e815dc2a37efe3f23da7dd36231 Binary files /dev/null and b/lib/python3.12/site-packages/multidict/__pycache__/_abc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/multidict/__pycache__/_compat.cpython-312.pyc b/lib/python3.12/site-packages/multidict/__pycache__/_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ee8441a7a917d2ca7b89671860304482f646526 Binary files /dev/null and b/lib/python3.12/site-packages/multidict/__pycache__/_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/multidict/__pycache__/_multidict_py.cpython-312.pyc b/lib/python3.12/site-packages/multidict/__pycache__/_multidict_py.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2f9d7461d40d9131de149aa7bcc0ff22ad52408 Binary files /dev/null and b/lib/python3.12/site-packages/multidict/__pycache__/_multidict_py.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/multidict/_abc.py b/lib/python3.12/site-packages/multidict/_abc.py new file mode 100644 index 0000000000000000000000000000000000000000..54253e9e779915aa8741313563359a8d9d87ec16 --- /dev/null +++ b/lib/python3.12/site-packages/multidict/_abc.py @@ -0,0 +1,73 @@ +import abc +from collections.abc import Iterable, Mapping, MutableMapping +from typing import TYPE_CHECKING, Protocol, TypeVar, Union, overload + +if TYPE_CHECKING: + from ._multidict_py import istr +else: + istr = str + +_V = TypeVar("_V") +_V_co = TypeVar("_V_co", covariant=True) +_T = TypeVar("_T") + + +class SupportsKeys(Protocol[_V_co]): + def keys(self) -> Iterable[str]: ... + def __getitem__(self, key: str, /) -> _V_co: ... + + +class SupportsIKeys(Protocol[_V_co]): + def keys(self) -> Iterable[istr]: ... + def __getitem__(self, key: istr, /) -> _V_co: ... + + +MDArg = Union[SupportsKeys[_V], SupportsIKeys[_V], Iterable[tuple[str, _V]], None] + + +class MultiMapping(Mapping[str, _V_co]): + @overload + def getall(self, key: str) -> list[_V_co]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V_co], _T]: ... + @abc.abstractmethod + def getall(self, key: str, default: _T = ...) -> Union[list[_V_co], _T]: + """Return all values for key.""" + + @overload + def getone(self, key: str) -> _V_co: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V_co, _T]: ... + @abc.abstractmethod + def getone(self, key: str, default: _T = ...) -> Union[_V_co, _T]: + """Return first value for key.""" + + +class MutableMultiMapping(MultiMapping[_V], MutableMapping[str, _V]): + @abc.abstractmethod + def add(self, key: str, value: _V) -> None: + """Add value to list.""" + + @abc.abstractmethod + def extend(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Add everything from arg and kwargs to the mapping.""" + + @abc.abstractmethod + def merge(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Merge into the mapping, adding non-existing keys.""" + + @overload + def popone(self, key: str) -> _V: ... + @overload + def popone(self, key: str, default: _T) -> Union[_V, _T]: ... + @abc.abstractmethod + def popone(self, key: str, default: _T = ...) -> Union[_V, _T]: + """Remove specified key and return the corresponding value.""" + + @overload + def popall(self, key: str) -> list[_V]: ... + @overload + def popall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + @abc.abstractmethod + def popall(self, key: str, default: _T = ...) -> Union[list[_V], _T]: + """Remove all occurrences of key and return the list of corresponding values.""" diff --git a/lib/python3.12/site-packages/multidict/_compat.py b/lib/python3.12/site-packages/multidict/_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..264d327e6a920503a72cc9e6174b26b63677faf5 --- /dev/null +++ b/lib/python3.12/site-packages/multidict/_compat.py @@ -0,0 +1,15 @@ +import os +import platform + +NO_EXTENSIONS = bool(os.environ.get("MULTIDICT_NO_EXTENSIONS")) + +PYPY = platform.python_implementation() == "PyPy" + +USE_EXTENSIONS = not NO_EXTENSIONS and not PYPY + +if USE_EXTENSIONS: + try: + from . import _multidict # type: ignore[attr-defined] # noqa: F401 + except ImportError: # pragma: no cover + # FIXME: Refactor for coverage. See #837. + USE_EXTENSIONS = False diff --git a/lib/python3.12/site-packages/multidict/_multidict_py.py b/lib/python3.12/site-packages/multidict/_multidict_py.py new file mode 100644 index 0000000000000000000000000000000000000000..6b68d52eeff981cc8082e3426ba037ae1aedb2c1 --- /dev/null +++ b/lib/python3.12/site-packages/multidict/_multidict_py.py @@ -0,0 +1,1242 @@ +import enum +import functools +import reprlib +import sys +from array import array +from collections.abc import ( + ItemsView, + Iterable, + Iterator, + KeysView, + Mapping, + ValuesView, +) +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + NoReturn, + Optional, + TypeVar, + Union, + cast, + overload, +) + +from ._abc import MDArg, MultiMapping, MutableMultiMapping, SupportsKeys + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + + +class istr(str): + """Case insensitive str.""" + + __is_istr__ = True + __istr_identity__: Optional[str] = None + + +_V = TypeVar("_V") +_T = TypeVar("_T") + +_SENTINEL = enum.Enum("_SENTINEL", "sentinel") +sentinel = _SENTINEL.sentinel + +_version = array("Q", [0]) + + +class _Iter(Generic[_T]): + __slots__ = ("_size", "_iter") + + def __init__(self, size: int, iterator: Iterator[_T]): + self._size = size + self._iter = iterator + + def __iter__(self) -> Self: + return self + + def __next__(self) -> _T: + return next(self._iter) + + def __length_hint__(self) -> int: + return self._size + + +class _ViewBase(Generic[_V]): + def __init__( + self, + md: "MultiDict[_V]", + ): + self._md = md + + def __len__(self) -> int: + return len(self._md) + + +class _ItemsView(_ViewBase[_V], ItemsView[str, _V]): + def __contains__(self, item: object) -> bool: + if not isinstance(item, (tuple, list)) or len(item) != 2: + return False + key, value = item + try: + identity = self._md._identity(key) + except TypeError: + return False + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and value == e.value: + return True + return False + + def __iter__(self) -> _Iter[tuple[str, _V]]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[tuple[str, _V]]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield self._md._key(e.key), e.value + + @reprlib.recursive_repr() + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(f"'{e.key}': {e.value!r}") + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + def _parse_item( + self, arg: Union[tuple[str, _V], _T] + ) -> Optional[tuple[int, str, str, _V]]: + if not isinstance(arg, tuple): + return None + if len(arg) != 2: + return None + try: + identity = self._md._identity(arg[0]) + return (hash(identity), identity, arg[0], arg[1]) + except TypeError: + return None + + def _tmp_set(self, it: Iterable[_T]) -> set[tuple[str, _V]]: + tmp = set() + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + else: + tmp.add((item[1], item[3])) + return tmp + + def __and__(self, other: Iterable[Any]) -> set[tuple[str, _V]]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + e.hash = -1 + if e.identity == identity and e.value == value: + ret.add((e.key, e.value)) + self._md._keys.restore_hash(hash_) + return ret + + def __rand__(self, other: Iterable[_T]) -> set[_T]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: + ret.add(arg) + break + return ret + + def __or__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + ret: set[Union[tuple[str, _V], _T]] = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item: Optional[tuple[int, str, str, _V]] = self._parse_item(arg) + if item is None: + ret.add(arg) + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + break + else: + ret.add(arg) + return ret + + def __ror__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + try: + ret: set[Union[tuple[str, _V], _T]] = set(other) + except TypeError: + return NotImplemented + tmp = self._tmp_set(ret) + + for e in self._md._keys.iter_entries(): + if (e.identity, e.value) not in tmp: + ret.add((e.key, e.value)) + return ret + + def __sub__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + ret: set[Union[tuple[str, _V], _T]] = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + tmp = self._tmp_set(it) + + for e in self._md._keys.iter_entries(): + if (e.identity, e.value) not in tmp: + ret.add((e.key, e.value)) + + return ret + + def __rsub__(self, other: Iterable[_T]) -> set[_T]: + ret: set[_T] = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + ret.add(arg) + continue + + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + break + else: + ret.add(arg) + return ret + + def __xor__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + try: + rgt = set(other) + except TypeError: + return NotImplemented + ret: set[Union[tuple[str, _V], _T]] = self - rgt + ret |= rgt - self + return ret + + __rxor__ = __xor__ + + def isdisjoint(self, other: Iterable[tuple[str, _V]]) -> bool: + for arg in other: + item = self._parse_item(arg) + if item is None: + continue + + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + return False + return True + + +class _ValuesView(_ViewBase[_V], ValuesView[_V]): + def __contains__(self, value: object) -> bool: + for e in self._md._keys.iter_entries(): + if e.value == value: + return True + return False + + def __iter__(self) -> _Iter[_V]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[_V]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield e.value + + @reprlib.recursive_repr() + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(repr(e.value)) + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + +class _KeysView(_ViewBase[_V], KeysView[str]): + def __contains__(self, key: object) -> bool: + if not isinstance(key, str): + return False + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return True + return False + + def __iter__(self) -> _Iter[str]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[str]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield self._md._key(e.key) + + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(f"'{e.key}'") + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + def __and__(self, other: Iterable[object]) -> set[str]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + ret.add(e.key) + break + return ret + + def __rand__(self, other: Iterable[_T]) -> set[_T]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + if key in self._md: + ret.add(key) + return cast(set[_T], ret) + + def __or__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + ret: set[Union[str, _T]] = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + ret.add(key) + continue + if key not in self._md: + ret.add(key) + return ret + + def __ror__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + try: + ret: set[Union[str, _T]] = set(other) + except TypeError: + return NotImplemented + + tmp = set() + for key in ret: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + tmp.add(identity) + + for e in self._md._keys.iter_entries(): + if e.identity not in tmp: + ret.add(e.key) + return ret + + def __sub__(self, other: Iterable[object]) -> set[str]: + ret = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + ret.discard(e.key) + break + return ret + + def __rsub__(self, other: Iterable[_T]) -> set[_T]: + try: + ret: set[_T] = set(other) + except TypeError: + return NotImplemented + for key in other: + if not isinstance(key, str): + continue + if key in self._md: + ret.discard(key) # type: ignore[arg-type] + return ret + + def __xor__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + try: + rgt = set(other) + except TypeError: + return NotImplemented + ret: set[Union[str, _T]] = self - rgt # type: ignore[assignment] + ret |= rgt - self + return ret + + __rxor__ = __xor__ + + def isdisjoint(self, other: Iterable[object]) -> bool: + for key in other: + if not isinstance(key, str): + continue + if key in self._md: + return False + return True + + +class _CSMixin: + _ci: ClassVar[bool] = False + + def _key(self, key: str) -> str: + return key + + def _identity(self, key: str) -> str: + if isinstance(key, str): + return key + else: + raise TypeError("MultiDict keys should be either str or subclasses of str") + + +class _CIMixin: + _ci: ClassVar[bool] = True + + def _key(self, key: str) -> str: + if type(key) is istr: + return key + else: + return istr(key) + + def _identity(self, key: str) -> str: + if isinstance(key, istr): + ret = key.__istr_identity__ + if ret is None: + ret = key.lower() + key.__istr_identity__ = ret + return ret + if isinstance(key, str): + return key.lower() + else: + raise TypeError("MultiDict keys should be either str or subclasses of str") + + +def estimate_log2_keysize(n: int) -> int: + # 7 == HT_MINSIZE - 1 + return (((n * 3 + 1) // 2) | 7).bit_length() + + +@dataclass +class _Entry(Generic[_V]): + hash: int + identity: str + key: str + value: _V + + +@dataclass +class _HtKeys(Generic[_V]): # type: ignore[misc] + LOG_MINSIZE: ClassVar[int] = 3 + MINSIZE: ClassVar[int] = 8 + PREALLOCATED_INDICES: ClassVar[dict[int, array]] = { # type: ignore[type-arg] + log2_size: array( + "b" if log2_size < 8 else "h", (-1 for i in range(1 << log2_size)) + ) + for log2_size in range(3, 10) + } + + log2_size: int + usable: int + + indices: array # type: ignore[type-arg] # in py3.9 array is not generic + entries: list[Optional[_Entry[_V]]] + + @functools.cached_property + def nslots(self) -> int: + return 1 << self.log2_size + + @functools.cached_property + def mask(self) -> int: + return self.nslots - 1 + + if sys.implementation.name != "pypy": + + def __sizeof__(self) -> int: + return ( + object.__sizeof__(self) + + sys.getsizeof(self.indices) + + sys.getsizeof(self.entries) + ) + + @classmethod + def new(cls, log2_size: int, entries: list[Optional[_Entry[_V]]]) -> Self: + size = 1 << log2_size + usable = (size << 1) // 3 + if log2_size < 10: + indices = cls.PREALLOCATED_INDICES[log2_size].__copy__() + elif log2_size < 16: + indices = array("h", (-1 for i in range(size))) + elif log2_size < 32: + indices = array("l", (-1 for i in range(size))) + else: # pragma: no cover # don't test huge multidicts + indices = array("q", (-1 for i in range(size))) + ret = cls( + log2_size=log2_size, + usable=usable, + indices=indices, + entries=entries, + ) + return ret + + def clone(self) -> "_HtKeys[_V]": + entries = [ + _Entry(e.hash, e.identity, e.key, e.value) if e is not None else None + for e in self.entries + ] + + return _HtKeys( + log2_size=self.log2_size, + usable=self.usable, + indices=self.indices.__copy__(), + entries=entries, + ) + + def build_indices(self, update: bool) -> None: + mask = self.mask + indices = self.indices + for idx, e in enumerate(self.entries): + assert e is not None + hash_ = e.hash + if update: + if hash_ == -1: + hash_ = hash(e.identity) + else: + assert hash_ != -1 + i = hash_ & mask + perturb = hash_ & sys.maxsize + while indices[i] != -1: + perturb >>= 5 + i = mask & (i * 5 + perturb + 1) + indices[i] = idx + + def find_empty_slot(self, hash_: int) -> int: + mask = self.mask + indices = self.indices + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + return i + + def iter_hash(self, hash_: int) -> Iterator[tuple[int, int, _Entry[_V]]]: + mask = self.mask + indices = self.indices + entries = self.entries + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + if ix != -2: + e = entries[ix] + if e.hash == hash_: + yield i, ix, e + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + + def del_idx(self, hash_: int, idx: int) -> None: + mask = self.mask + indices = self.indices + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != idx: + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + indices[i] = -2 + + def iter_entries(self) -> Iterator[_Entry[_V]]: + return filter(None, self.entries) + + def restore_hash(self, hash_: int) -> None: + mask = self.mask + indices = self.indices + entries = self.entries + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + if ix != -2: + entry = entries[ix] + if entry.hash == -1: + entry.hash = hash_ + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + + +class MultiDict(_CSMixin, MutableMultiMapping[_V]): + """Dictionary with the support for duplicate keys.""" + + __slots__ = ("_keys", "_used", "_version") + + def __init__(self, arg: MDArg[_V] = None, /, **kwargs: _V): + self._used = 0 + v = _version + v[0] += 1 + self._version = v[0] + if not kwargs: + md = None + if isinstance(arg, MultiDictProxy): + md = arg._md + elif isinstance(arg, MultiDict): + md = arg + if md is not None and md._ci is self._ci: + self._from_md(md) + return + + it = self._parse_args(arg, kwargs) + log2_size = estimate_log2_keysize(cast(int, next(it))) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in init + log2_size = 17 + self._keys: _HtKeys[_V] = _HtKeys.new(log2_size, []) + self._extend_items(cast(Iterator[_Entry[_V]], it)) + + def _from_md(self, md: "MultiDict[_V]") -> None: + # Copy everything as-is without compacting the new multidict, + # otherwise it requires reindexing + self._keys = md._keys.clone() + self._used = md._used + + @overload + def getall(self, key: str) -> list[_V]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def getall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Return a list of all values matching the key.""" + identity = self._identity(key) + hash_ = hash(identity) + res = [] + restore = [] + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + res.append(e.value) + e.hash = -1 + restore.append(idx) + + if res: + entries = self._keys.entries + for idx in restore: + entries[idx].hash = hash_ # type: ignore[union-attr] + return res + if not res and default is not sentinel: + return default + raise KeyError("Key not found: %r" % key) + + @overload + def getone(self, key: str) -> _V: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V, _T]: ... + def getone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Get first value matching the key. + + Raises KeyError if the key is not found and no default is provided. + """ + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return e.value + if default is not sentinel: + return default + raise KeyError("Key not found: %r" % key) + + # Mapping interface # + + def __getitem__(self, key: str) -> _V: + return self.getone(key) + + @overload + def get(self, key: str, /) -> Union[_V, None]: ... + @overload + def get(self, key: str, /, default: _T) -> Union[_V, _T]: ... + def get(self, key: str, default: Union[_T, None] = None) -> Union[_V, _T, None]: + """Get first value matching the key. + + If the key is not found, returns the default (or None if no default is provided) + """ + return self.getone(key, default) + + def __iter__(self) -> Iterator[str]: + return iter(self.keys()) + + def __len__(self) -> int: + return self._used + + def keys(self) -> KeysView[str]: + """Return a new view of the dictionary's keys.""" + return _KeysView(self) + + def items(self) -> ItemsView[str, _V]: + """Return a new view of the dictionary's items *(key, value) pairs).""" + return _ItemsView(self) + + def values(self) -> _ValuesView[_V]: + """Return a new view of the dictionary's values.""" + return _ValuesView(self) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Mapping): + return NotImplemented + if isinstance(other, MultiDictProxy): + return self == other._md + if isinstance(other, MultiDict): + lft = self._keys + rht = other._keys + if self._used != other._used: + return False + for e1, e2 in zip(lft.iter_entries(), rht.iter_entries()): + if e1.identity != e2.identity or e1.value != e2.value: + return False + return True + if self._used != len(other): + return False + for k, v in self.items(): + nv = other.get(k, sentinel) + if v != nv: + return False + return True + + def __contains__(self, key: object) -> bool: + if not isinstance(key, str): + return False + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return True + return False + + @reprlib.recursive_repr() + def __repr__(self) -> str: + body = ", ".join(f"'{e.key}': {e.value!r}" for e in self._keys.iter_entries()) + return f"<{self.__class__.__name__}({body})>" + + if sys.implementation.name != "pypy": + + def __sizeof__(self) -> int: + return object.__sizeof__(self) + sys.getsizeof(self._keys) + + def __reduce__(self) -> tuple[type[Self], tuple[list[tuple[str, _V]]]]: + return (self.__class__, (list(self.items()),)) + + def add(self, key: str, value: _V) -> None: + identity = self._identity(key) + hash_ = hash(identity) + self._add_with_hash(_Entry(hash_, identity, key, value)) + self._incr_version() + + def copy(self) -> Self: + """Return a copy of itself.""" + cls = self.__class__ + return cls(self) + + __copy__ = copy + + def extend(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Extend current MultiDict with more values. + + This method must be used instead of update. + """ + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + self._resize(estimate_log2_keysize(newsize), False) + self._extend_items(cast(Iterator[_Entry[_V]], it)) + + def _parse_args( + self, + arg: MDArg[_V], + kwargs: Mapping[str, _V], + ) -> Iterator[Union[int, _Entry[_V]]]: + identity_func = self._identity + if arg: + if isinstance(arg, MultiDictProxy): + arg = arg._md + if isinstance(arg, MultiDict): + yield len(arg) + len(kwargs) + if self._ci is not arg._ci: + for e in arg._keys.iter_entries(): + identity = identity_func(e.key) + yield _Entry(hash(identity), identity, e.key, e.value) + else: + for e in arg._keys.iter_entries(): + yield _Entry(e.hash, e.identity, e.key, e.value) + if kwargs: + for key, value in kwargs.items(): + identity = identity_func(key) + yield _Entry(hash(identity), identity, key, value) + else: + if hasattr(arg, "keys"): + arg = cast(SupportsKeys[_V], arg) + arg = [(k, arg[k]) for k in arg.keys()] + if kwargs: + arg = list(arg) + arg.extend(list(kwargs.items())) + try: + yield len(arg) + len(kwargs) # type: ignore[arg-type] + except TypeError: + yield 0 + for pos, item in enumerate(arg): + if not len(item) == 2: + raise ValueError( + f"multidict update sequence element #{pos}" + f"has length {len(item)}; 2 is required" + ) + identity = identity_func(item[0]) + yield _Entry(hash(identity), identity, item[0], item[1]) + else: + yield len(kwargs) + for key, value in kwargs.items(): + identity = identity_func(key) + yield _Entry(hash(identity), identity, key, value) + + def _extend_items(self, items: Iterable[_Entry[_V]]) -> None: + for e in items: + self._add_with_hash(e) + self._incr_version() + + def clear(self) -> None: + """Remove all items from MultiDict.""" + self._used = 0 + self._keys = _HtKeys.new(_HtKeys.LOG_MINSIZE, []) + self._incr_version() + + # Mapping interface # + + def __setitem__(self, key: str, value: _V) -> None: + identity = self._identity(key) + hash_ = hash(identity) + found = False + + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + if not found: + e.key = key + e.value = value + e.hash = -1 + found = True + self._incr_version() + elif e.hash != -1: # pragma: no branch + self._del_at(slot, idx) + + if not found: + self._add_with_hash(_Entry(hash_, identity, key, value)) + else: + self._keys.restore_hash(hash_) + + def __delitem__(self, key: str) -> None: + found = False + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + self._del_at(slot, idx) + found = True + if not found: + raise KeyError(key) + else: + self._incr_version() + + @overload + def setdefault( + self: "MultiDict[Union[_T, None]]", key: str, default: None = None + ) -> Union[_T, None]: ... + @overload + def setdefault(self, key: str, default: _V) -> _V: ... + def setdefault(self, key: str, default: Union[_V, None] = None) -> Union[_V, None]: # type: ignore[misc] + """Return value for key, set value to default if key is not present.""" + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return e.value + self.add(key, default) # type: ignore[arg-type] + return default + + @overload + def popone(self, key: str) -> _V: ... + @overload + def popone(self, key: str, default: _T) -> Union[_V, _T]: ... + def popone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Remove specified key and return the corresponding value. + + If key is not found, d is returned if given, otherwise + KeyError is raised. + + """ + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + value = e.value + self._del_at(slot, idx) + self._incr_version() + return value + if default is sentinel: + raise KeyError(key) + else: + return default + + # Type checking will inherit signature for pop() if we don't confuse it here. + if not TYPE_CHECKING: + pop = popone + + @overload + def popall(self, key: str) -> list[_V]: ... + @overload + def popall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def popall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Remove all occurrences of key and return the list of corresponding + values. + + If key is not found, default is returned if given, otherwise + KeyError is raised. + + """ + found = False + identity = self._identity(key) + hash_ = hash(identity) + ret = [] + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + found = True + ret.append(e.value) + self._del_at(slot, idx) + self._incr_version() + + if not found: + if default is sentinel: + raise KeyError(key) + else: + return default + else: + return ret + + def popitem(self) -> tuple[str, _V]: + """Remove and return an arbitrary (key, value) pair.""" + if self._used <= 0: + raise KeyError("empty multidict") + + pos = len(self._keys.entries) - 1 + entry = self._keys.entries.pop() + + while entry is None: + pos -= 1 + entry = self._keys.entries.pop() + + ret = self._key(entry.key), entry.value + self._keys.del_idx(entry.hash, pos) + self._used -= 1 + self._incr_version() + return ret + + def update(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Update the dictionary, overwriting existing keys.""" + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + log2_size = estimate_log2_keysize(newsize) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in update, + # duplicate keys could reduce the resulting anount of entries + log2_size = 17 + if log2_size > self._keys.log2_size: + self._resize(log2_size, False) + try: + self._update_items(cast(Iterator[_Entry[_V]], it)) + finally: + self._post_update() + + def _update_items(self, items: Iterator[_Entry[_V]]) -> None: + for entry in items: + found = False + hash_ = entry.hash + identity = entry.identity + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + if not found: + found = True + e.key = entry.key + e.value = entry.value + e.hash = -1 + else: + self._del_at_for_upd(e) + if not found: + self._add_with_hash_for_upd(entry) + + def _post_update(self) -> None: + keys = self._keys + indices = keys.indices + entries = keys.entries + for slot in range(keys.nslots): + idx = indices[slot] + if idx >= 0: + e2 = entries[idx] + assert e2 is not None + if e2.key is None: + entries[idx] = None + indices[slot] = -2 + self._used -= 1 + if e2.hash == -1: + e2.hash = hash(e2.identity) + + self._incr_version() + + def merge(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Merge into the dictionary, adding non-existing keys.""" + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + log2_size = estimate_log2_keysize(newsize) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in update, + # duplicate keys could reduce the resulting anount of entries + log2_size = 17 + if log2_size > self._keys.log2_size: + self._resize(log2_size, False) + try: + self._merge_items(cast(Iterator[_Entry[_V]], it)) + finally: + self._post_update() + + def _merge_items(self, items: Iterator[_Entry[_V]]) -> None: + for entry in items: + hash_ = entry.hash + identity = entry.identity + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + break + else: + self._add_with_hash_for_upd(entry) + + def _incr_version(self) -> None: + v = _version + v[0] += 1 + self._version = v[0] + + def _resize(self, log2_newsize: int, update: bool) -> None: + oldkeys = self._keys + newentries = self._used + + if len(oldkeys.entries) == newentries: + entries = oldkeys.entries + else: + entries = [e for e in oldkeys.entries if e is not None] + newkeys: _HtKeys[_V] = _HtKeys.new(log2_newsize, entries) + newkeys.usable -= newentries + newkeys.build_indices(update) + self._keys = newkeys + + def _add_with_hash(self, entry: _Entry[_V]) -> None: + if self._keys.usable <= 0: + self._resize((self._used * 3 | _HtKeys.MINSIZE - 1).bit_length(), False) + keys = self._keys + slot = keys.find_empty_slot(entry.hash) + keys.indices[slot] = len(keys.entries) + keys.entries.append(entry) + self._incr_version() + self._used += 1 + keys.usable -= 1 + + def _add_with_hash_for_upd(self, entry: _Entry[_V]) -> None: + if self._keys.usable <= 0: + self._resize((self._used * 3 | _HtKeys.MINSIZE - 1).bit_length(), True) + keys = self._keys + slot = keys.find_empty_slot(entry.hash) + keys.indices[slot] = len(keys.entries) + entry.hash = -1 + keys.entries.append(entry) + self._incr_version() + self._used += 1 + keys.usable -= 1 + + def _del_at(self, slot: int, idx: int) -> None: + self._keys.entries[idx] = None + self._keys.indices[slot] = -2 + self._used -= 1 + + def _del_at_for_upd(self, entry: _Entry[_V]) -> None: + entry.key = None # type: ignore[assignment] + entry.value = None # type: ignore[assignment] + + +class CIMultiDict(_CIMixin, MultiDict[_V]): + """Dictionary with the support for duplicate case-insensitive keys.""" + + +class MultiDictProxy(_CSMixin, MultiMapping[_V]): + """Read-only proxy for MultiDict instance.""" + + __slots__ = ("_md",) + + _md: MultiDict[_V] + + def __init__(self, arg: Union[MultiDict[_V], "MultiDictProxy[_V]"]): + if not isinstance(arg, (MultiDict, MultiDictProxy)): + raise TypeError( + f"ctor requires MultiDict or MultiDictProxy instance, not {type(arg)}" + ) + if isinstance(arg, MultiDictProxy): + self._md = arg._md + else: + self._md = arg + + def __reduce__(self) -> NoReturn: + raise TypeError(f"can't pickle {self.__class__.__name__} objects") + + @overload + def getall(self, key: str) -> list[_V]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def getall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Return a list of all values matching the key.""" + if default is not sentinel: + return self._md.getall(key, default) + else: + return self._md.getall(key) + + @overload + def getone(self, key: str) -> _V: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V, _T]: ... + def getone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Get first value matching the key. + + Raises KeyError if the key is not found and no default is provided. + """ + if default is not sentinel: + return self._md.getone(key, default) + else: + return self._md.getone(key) + + # Mapping interface # + + def __getitem__(self, key: str) -> _V: + return self.getone(key) + + @overload + def get(self, key: str, /) -> Union[_V, None]: ... + @overload + def get(self, key: str, /, default: _T) -> Union[_V, _T]: ... + def get(self, key: str, default: Union[_T, None] = None) -> Union[_V, _T, None]: + """Get first value matching the key. + + If the key is not found, returns the default (or None if no default is provided) + """ + return self._md.getone(key, default) + + def __iter__(self) -> Iterator[str]: + return iter(self._md.keys()) + + def __len__(self) -> int: + return len(self._md) + + def keys(self) -> KeysView[str]: + """Return a new view of the dictionary's keys.""" + return self._md.keys() + + def items(self) -> ItemsView[str, _V]: + """Return a new view of the dictionary's items *(key, value) pairs).""" + return self._md.items() + + def values(self) -> _ValuesView[_V]: + """Return a new view of the dictionary's values.""" + return self._md.values() + + def __eq__(self, other: object) -> bool: + return self._md == other + + def __contains__(self, key: object) -> bool: + return key in self._md + + @reprlib.recursive_repr() + def __repr__(self) -> str: + body = ", ".join(f"'{k}': {v!r}" for k, v in self.items()) + return f"<{self.__class__.__name__}({body})>" + + def copy(self) -> MultiDict[_V]: + """Return a copy of itself.""" + return MultiDict(self._md) + + +class CIMultiDictProxy(_CIMixin, MultiDictProxy[_V]): + """Read-only proxy for CIMultiDict instance.""" + + def __init__(self, arg: Union[MultiDict[_V], MultiDictProxy[_V]]): + if not isinstance(arg, (CIMultiDict, CIMultiDictProxy)): + raise TypeError( + "ctor requires CIMultiDict or CIMultiDictProxy instance" + f", not {type(arg)}" + ) + + super().__init__(arg) + + def copy(self) -> CIMultiDict[_V]: + """Return a copy of itself.""" + return CIMultiDict(self._md) + + +def getversion(md: Union[MultiDict[object], MultiDictProxy[object]]) -> int: + if isinstance(md, MultiDictProxy): + md = md._md + elif not isinstance(md, MultiDict): + raise TypeError("Parameter should be multidict or proxy") + return md._version diff --git a/lib/python3.12/site-packages/multidict/py.typed b/lib/python3.12/site-packages/multidict/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..dfe8cc048e7100a97025b954fffa31e4ff859a7d --- /dev/null +++ b/lib/python3.12/site-packages/multidict/py.typed @@ -0,0 +1 @@ +PEP-561 marker. \ No newline at end of file diff --git a/lib/python3.12/site-packages/sentry_sdk/__init__.py b/lib/python3.12/site-packages/sentry_sdk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e149418c381a764bf1489d2a7b31df02d9b34259 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/__init__.py @@ -0,0 +1,67 @@ +from sentry_sdk import profiler +from sentry_sdk import metrics +from sentry_sdk.scope import Scope +from sentry_sdk.transport import Transport, HttpTransport +from sentry_sdk.client import Client + +from sentry_sdk.api import * # noqa +from sentry_sdk.consts import VERSION + +__all__ = [ # noqa + "Hub", + "Scope", + "Client", + "Transport", + "HttpTransport", + "VERSION", + "integrations", + # From sentry_sdk.api + "init", + "add_attachment", + "add_breadcrumb", + "capture_event", + "capture_exception", + "capture_message", + "configure_scope", + "continue_trace", + "flush", + "get_baggage", + "get_client", + "get_global_scope", + "get_isolation_scope", + "get_current_scope", + "get_current_span", + "get_traceparent", + "is_initialized", + "isolation_scope", + "last_event_id", + "new_scope", + "push_scope", + "set_context", + "set_extra", + "set_level", + "set_measurement", + "set_tag", + "set_tags", + "set_user", + "start_span", + "start_transaction", + "trace", + "monitor", + "logger", + "metrics", + "profiler", + "start_session", + "end_session", + "set_transaction_name", + "update_current_span", +] + +# Initialize the debug support after everything is loaded +from sentry_sdk.debug import init_debug_support + +init_debug_support() +del init_debug_support + +# circular imports +from sentry_sdk.hub import Hub diff --git a/lib/python3.12/site-packages/sentry_sdk/_batcher.py b/lib/python3.12/site-packages/sentry_sdk/_batcher.py new file mode 100644 index 0000000000000000000000000000000000000000..4a6ee07e67ba69c92123d16b0aa93cb1103d67d2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_batcher.py @@ -0,0 +1,139 @@ +import os +import random +import threading +from datetime import datetime, timezone +from typing import TYPE_CHECKING, TypeVar, Generic + +from sentry_sdk.utils import format_timestamp, safe_repr, serialize_attribute +from sentry_sdk.envelope import Envelope, Item, PayloadRef + +if TYPE_CHECKING: + from typing import Optional, Callable, Any + +T = TypeVar("T") + + +class Batcher(Generic[T]): + MAX_BEFORE_FLUSH = 100 + MAX_BEFORE_DROP = 1_000 + FLUSH_WAIT_TIME = 5.0 + + TYPE = "" + CONTENT_TYPE = "" + + def __init__( + self, + capture_func: "Callable[[Envelope], None]", + record_lost_func: "Callable[..., None]", + ) -> None: + self._buffer: "list[T]" = [] + self._capture_func = capture_func + self._record_lost_func = record_lost_func + self._running = True + self._lock = threading.Lock() + + self._flush_event: "threading.Event" = threading.Event() + + self._flusher: "Optional[threading.Thread]" = None + self._flusher_pid: "Optional[int]" = None + + def _ensure_thread(self) -> bool: + """For forking processes we might need to restart this thread. + This ensures that our process actually has that thread running. + """ + if not self._running: + return False + + pid = os.getpid() + if self._flusher_pid == pid: + return True + + with self._lock: + # Recheck to make sure another thread didn't get here and start the + # the flusher in the meantime + if self._flusher_pid == pid: + return True + + self._flusher_pid = pid + + self._flusher = threading.Thread(target=self._flush_loop) + self._flusher.daemon = True + + try: + self._flusher.start() + except RuntimeError: + # Unfortunately at this point the interpreter is in a state that no + # longer allows us to spawn a thread and we have to bail. + self._running = False + return False + + return True + + def _flush_loop(self) -> None: + while self._running: + self._flush_event.wait(self.FLUSH_WAIT_TIME + random.random()) + self._flush_event.clear() + self._flush() + + def add(self, item: "T") -> None: + if not self._ensure_thread() or self._flusher is None: + return None + + with self._lock: + if len(self._buffer) >= self.MAX_BEFORE_DROP: + self._record_lost(item) + return None + + self._buffer.append(item) + if len(self._buffer) >= self.MAX_BEFORE_FLUSH: + self._flush_event.set() + + def kill(self) -> None: + if self._flusher is None: + return + + self._running = False + self._flush_event.set() + self._flusher = None + + def flush(self) -> None: + self._flush() + + def _add_to_envelope(self, envelope: "Envelope") -> None: + envelope.add_item( + Item( + type=self.TYPE, + content_type=self.CONTENT_TYPE, + headers={ + "item_count": len(self._buffer), + }, + payload=PayloadRef( + json={ + "items": [ + self._to_transport_format(item) for item in self._buffer + ] + } + ), + ) + ) + + def _flush(self) -> "Optional[Envelope]": + envelope = Envelope( + headers={"sent_at": format_timestamp(datetime.now(timezone.utc))} + ) + with self._lock: + if len(self._buffer) == 0: + return None + + self._add_to_envelope(envelope) + self._buffer.clear() + + self._capture_func(envelope) + return envelope + + def _record_lost(self, item: "T") -> None: + pass + + @staticmethod + def _to_transport_format(item: "T") -> "Any": + pass diff --git a/lib/python3.12/site-packages/sentry_sdk/_compat.py b/lib/python3.12/site-packages/sentry_sdk/_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..dcb590fcfa8d3ae8d110ff3402accf300a2f542b --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_compat.py @@ -0,0 +1,94 @@ +import sys + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import TypeVar + + T = TypeVar("T") + + +PY37 = sys.version_info[0] == 3 and sys.version_info[1] >= 7 +PY38 = sys.version_info[0] == 3 and sys.version_info[1] >= 8 +PY310 = sys.version_info[0] == 3 and sys.version_info[1] >= 10 +PY311 = sys.version_info[0] == 3 and sys.version_info[1] >= 11 + + +def with_metaclass(meta: "Any", *bases: "Any") -> "Any": + class MetaClass(type): + def __new__(metacls: "Any", name: "Any", this_bases: "Any", d: "Any") -> "Any": + return meta(name, bases, d) + + return type.__new__(MetaClass, "temporary_class", (), {}) + + +def check_uwsgi_thread_support() -> bool: + # We check two things here: + # + # 1. uWSGI doesn't run in threaded mode by default -- issue a warning if + # that's the case. + # + # 2. Additionally, if uWSGI is running in preforking mode (default), it needs + # the --py-call-uwsgi-fork-hooks option for the SDK to work properly. This + # is because any background threads spawned before the main process is + # forked are NOT CLEANED UP IN THE CHILDREN BY DEFAULT even if + # --enable-threads is on. One has to explicitly provide + # --py-call-uwsgi-fork-hooks to force uWSGI to run regular cpython + # after-fork hooks that take care of cleaning up stale thread data. + try: + from uwsgi import opt # type: ignore + except ImportError: + return True + + from sentry_sdk.consts import FALSE_VALUES + + def enabled(option: str) -> bool: + value = opt.get(option, False) + if isinstance(value, bool): + return value + + if isinstance(value, bytes): + try: + value = value.decode() + except Exception: + pass + + return value and str(value).lower() not in FALSE_VALUES + + # When `threads` is passed in as a uwsgi option, + # `enable-threads` is implied on. + threads_enabled = "threads" in opt or enabled("enable-threads") + fork_hooks_on = enabled("py-call-uwsgi-fork-hooks") + lazy_mode = enabled("lazy-apps") or enabled("lazy") + + if lazy_mode and not threads_enabled: + from warnings import warn + + warn( + Warning( + "IMPORTANT: " + "We detected the use of uWSGI without thread support. " + "This might lead to unexpected issues. " + 'Please run uWSGI with "--enable-threads" for full support.' + ) + ) + + return False + + elif not lazy_mode and (not threads_enabled or not fork_hooks_on): + from warnings import warn + + warn( + Warning( + "IMPORTANT: " + "We detected the use of uWSGI in preforking mode without " + "thread support. This might lead to crashing workers. " + 'Please run uWSGI with both "--enable-threads" and ' + '"--py-call-uwsgi-fork-hooks" for full support.' + ) + ) + + return False + + return True diff --git a/lib/python3.12/site-packages/sentry_sdk/_init_implementation.py b/lib/python3.12/site-packages/sentry_sdk/_init_implementation.py new file mode 100644 index 0000000000000000000000000000000000000000..c2d77809c77dd9085ff3439e9f641e33a142dd5f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_init_implementation.py @@ -0,0 +1,79 @@ +import warnings + +from typing import TYPE_CHECKING + +import sentry_sdk + +if TYPE_CHECKING: + from typing import Any, ContextManager, Optional + + import sentry_sdk.consts + + +class _InitGuard: + _CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE = ( + "Using the return value of sentry_sdk.init as a context manager " + "and manually calling the __enter__ and __exit__ methods on the " + "return value are deprecated. We are no longer maintaining this " + "functionality, and we will remove it in the next major release." + ) + + def __init__(self, client: "sentry_sdk.Client") -> None: + self._client = client + + def __enter__(self) -> "_InitGuard": + warnings.warn( + self._CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE, + stacklevel=2, + category=DeprecationWarning, + ) + + return self + + def __exit__(self, exc_type: "Any", exc_value: "Any", tb: "Any") -> None: + warnings.warn( + self._CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE, + stacklevel=2, + category=DeprecationWarning, + ) + + c = self._client + if c is not None: + c.close() + + +def _check_python_deprecations() -> None: + # Since we're likely to deprecate Python versions in the future, I'm keeping + # this handy function around. Use this to detect the Python version used and + # to output logger.warning()s if it's deprecated. + pass + + +def _init(*args: "Optional[str]", **kwargs: "Any") -> "ContextManager[Any]": + """Initializes the SDK and optionally integrations. + + This takes the same arguments as the client constructor. + """ + client = sentry_sdk.Client(*args, **kwargs) + sentry_sdk.get_global_scope().set_client(client) + _check_python_deprecations() + rv = _InitGuard(client) + return rv + + +if TYPE_CHECKING: + # Make mypy, PyCharm and other static analyzers think `init` is a type to + # have nicer autocompletion for params. + # + # Use `ClientConstructor` to define the argument types of `init` and + # `ContextManager[Any]` to tell static analyzers about the return type. + + class init(sentry_sdk.consts.ClientConstructor, _InitGuard): # noqa: N801 + pass + +else: + # Alias `init` for actual usage. Go through the lambda indirection to throw + # PyCharm off of the weakly typed signature (it would otherwise discover + # both the weakly typed signature of `_init` and our faked `init` type). + + init = (lambda: _init)() diff --git a/lib/python3.12/site-packages/sentry_sdk/_log_batcher.py b/lib/python3.12/site-packages/sentry_sdk/_log_batcher.py new file mode 100644 index 0000000000000000000000000000000000000000..1c59f7379c7862180adbd3cfc57433cf1670545f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_log_batcher.py @@ -0,0 +1,56 @@ +from typing import TYPE_CHECKING + +from sentry_sdk._batcher import Batcher +from sentry_sdk.utils import serialize_attribute +from sentry_sdk.envelope import Envelope, Item, PayloadRef + +if TYPE_CHECKING: + from typing import Any + from sentry_sdk._types import Log + + +class LogBatcher(Batcher["Log"]): + MAX_BEFORE_FLUSH = 100 + MAX_BEFORE_DROP = 1_000 + FLUSH_WAIT_TIME = 5.0 + + TYPE = "log" + CONTENT_TYPE = "application/vnd.sentry.items.log+json" + + @staticmethod + def _to_transport_format(item: "Log") -> "Any": + if "sentry.severity_number" not in item["attributes"]: + item["attributes"]["sentry.severity_number"] = item["severity_number"] + if "sentry.severity_text" not in item["attributes"]: + item["attributes"]["sentry.severity_text"] = item["severity_text"] + + res = { + "timestamp": int(item["time_unix_nano"]) / 1.0e9, + "trace_id": item.get("trace_id", "00000000-0000-0000-0000-000000000000"), + "span_id": item.get("span_id"), + "level": str(item["severity_text"]), + "body": str(item["body"]), + "attributes": { + k: serialize_attribute(v) for (k, v) in item["attributes"].items() + }, + } + + return res + + def _record_lost(self, item: "Log") -> None: + # Construct log envelope item without sending it to report lost bytes + log_item = Item( + type=self.TYPE, + content_type=self.CONTENT_TYPE, + headers={ + "item_count": 1, + }, + payload=PayloadRef(json={"items": [self._to_transport_format(item)]}), + ) + + self._record_lost_func( + reason="queue_overflow", + data_category="log_item", + item=log_item, + quantity=1, + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/_lru_cache.py b/lib/python3.12/site-packages/sentry_sdk/_lru_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..16c238bcab6f34530d0bb4e88039aad8a9acc9da --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_lru_cache.py @@ -0,0 +1,43 @@ +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +_SENTINEL = object() + + +class LRUCache: + def __init__(self, max_size: int) -> None: + if max_size <= 0: + raise AssertionError(f"invalid max_size: {max_size}") + self.max_size = max_size + self._data: "dict[Any, Any]" = {} + self.hits = self.misses = 0 + self.full = False + + def set(self, key: "Any", value: "Any") -> None: + current = self._data.pop(key, _SENTINEL) + if current is not _SENTINEL: + self._data[key] = value + elif self.full: + self._data.pop(next(iter(self._data))) + self._data[key] = value + else: + self._data[key] = value + self.full = len(self._data) >= self.max_size + + def get(self, key: "Any", default: "Any" = None) -> "Any": + try: + ret = self._data.pop(key) + except KeyError: + self.misses += 1 + ret = default + else: + self.hits += 1 + self._data[key] = ret + + return ret + + def get_all(self) -> "list[tuple[Any, Any]]": + return list(self._data.items()) diff --git a/lib/python3.12/site-packages/sentry_sdk/_metrics_batcher.py b/lib/python3.12/site-packages/sentry_sdk/_metrics_batcher.py new file mode 100644 index 0000000000000000000000000000000000000000..e60a4c86ec94cd6788b8dd7ac40b0ff7e9c0b479 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_metrics_batcher.py @@ -0,0 +1,46 @@ +from typing import TYPE_CHECKING + +from sentry_sdk._batcher import Batcher +from sentry_sdk.utils import serialize_attribute +from sentry_sdk.envelope import Item + +if TYPE_CHECKING: + from typing import Any + from sentry_sdk._types import Metric + + +class MetricsBatcher(Batcher["Metric"]): + MAX_BEFORE_FLUSH = 1000 + MAX_BEFORE_DROP = 10_000 + FLUSH_WAIT_TIME = 5.0 + + TYPE = "trace_metric" + CONTENT_TYPE = "application/vnd.sentry.items.trace-metric+json" + + @staticmethod + def _to_transport_format(item: "Metric") -> "Any": + res = { + "timestamp": item["timestamp"], + "trace_id": item["trace_id"], + "name": item["name"], + "type": item["type"], + "value": item["value"], + "attributes": { + k: serialize_attribute(v) for (k, v) in item["attributes"].items() + }, + } + + if item.get("span_id") is not None: + res["span_id"] = item["span_id"] + + if item.get("unit") is not None: + res["unit"] = item["unit"] + + return res + + def _record_lost(self, item: "Metric") -> None: + self._record_lost_func( + reason="queue_overflow", + data_category="trace_metric", + quantity=1, + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/_queue.py b/lib/python3.12/site-packages/sentry_sdk/_queue.py new file mode 100644 index 0000000000000000000000000000000000000000..9bdb76dddba7528ea8597c58ef74aa86ccd42ef1 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_queue.py @@ -0,0 +1,289 @@ +""" +A fork of Python 3.6's stdlib queue (found in Pythons 'cpython/Lib/queue.py') +with Lock swapped out for RLock to avoid a deadlock while garbage collecting. + +https://github.com/python/cpython/blob/v3.6.12/Lib/queue.py + + +See also +https://codewithoutrules.com/2017/08/16/concurrency-python/ +https://bugs.python.org/issue14976 +https://github.com/sqlalchemy/sqlalchemy/blob/4eb747b61f0c1b1c25bdee3856d7195d10a0c227/lib/sqlalchemy/queue.py#L1 + +We also vendor the code to evade eventlet's broken monkeypatching, see +https://github.com/getsentry/sentry-python/pull/484 + + +Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, +2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation; + +All Rights Reserved + + +PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +-------------------------------------------- + +1. This LICENSE AGREEMENT is between the Python Software Foundation +("PSF"), and the Individual or Organization ("Licensee") accessing and +otherwise using this software ("Python") in source or binary form and +its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, PSF hereby +grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce, +analyze, test, perform and/or display publicly, prepare derivative works, +distribute, and otherwise use Python alone or in any derivative version, +provided, however, that PSF's License Agreement and PSF's notice of copyright, +i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, +2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation; +All Rights Reserved" are retained in Python alone or in any derivative version +prepared by Licensee. + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python. + +4. PSF is making Python available to Licensee on an "AS IS" +basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between PSF and +Licensee. This License Agreement does not grant permission to use PSF +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using Python, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + +""" + +import threading + +from collections import deque +from time import time + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + +__all__ = ["EmptyError", "FullError", "Queue"] + + +class EmptyError(Exception): + "Exception raised by Queue.get(block=0)/get_nowait()." + + pass + + +class FullError(Exception): + "Exception raised by Queue.put(block=0)/put_nowait()." + + pass + + +class Queue: + """Create a queue object with a given maximum size. + + If maxsize is <= 0, the queue size is infinite. + """ + + def __init__(self, maxsize=0): + self.maxsize = maxsize + self._init(maxsize) + + # mutex must be held whenever the queue is mutating. All methods + # that acquire mutex must release it before returning. mutex + # is shared between the three conditions, so acquiring and + # releasing the conditions also acquires and releases mutex. + self.mutex = threading.RLock() + + # Notify not_empty whenever an item is added to the queue; a + # thread waiting to get is notified then. + self.not_empty = threading.Condition(self.mutex) + + # Notify not_full whenever an item is removed from the queue; + # a thread waiting to put is notified then. + self.not_full = threading.Condition(self.mutex) + + # Notify all_tasks_done whenever the number of unfinished tasks + # drops to zero; thread waiting to join() is notified to resume + self.all_tasks_done = threading.Condition(self.mutex) + self.unfinished_tasks = 0 + + def task_done(self): + """Indicate that a formerly enqueued task is complete. + + Used by Queue consumer threads. For each get() used to fetch a task, + a subsequent call to task_done() tells the queue that the processing + on the task is complete. + + If a join() is currently blocking, it will resume when all items + have been processed (meaning that a task_done() call was received + for every item that had been put() into the queue). + + Raises a ValueError if called more times than there were items + placed in the queue. + """ + with self.all_tasks_done: + unfinished = self.unfinished_tasks - 1 + if unfinished <= 0: + if unfinished < 0: + raise ValueError("task_done() called too many times") + self.all_tasks_done.notify_all() + self.unfinished_tasks = unfinished + + def join(self): + """Blocks until all items in the Queue have been gotten and processed. + + The count of unfinished tasks goes up whenever an item is added to the + queue. The count goes down whenever a consumer thread calls task_done() + to indicate the item was retrieved and all work on it is complete. + + When the count of unfinished tasks drops to zero, join() unblocks. + """ + with self.all_tasks_done: + while self.unfinished_tasks: + self.all_tasks_done.wait() + + def qsize(self): + """Return the approximate size of the queue (not reliable!).""" + with self.mutex: + return self._qsize() + + def empty(self): + """Return True if the queue is empty, False otherwise (not reliable!). + + This method is likely to be removed at some point. Use qsize() == 0 + as a direct substitute, but be aware that either approach risks a race + condition where a queue can grow before the result of empty() or + qsize() can be used. + + To create code that needs to wait for all queued tasks to be + completed, the preferred technique is to use the join() method. + """ + with self.mutex: + return not self._qsize() + + def full(self): + """Return True if the queue is full, False otherwise (not reliable!). + + This method is likely to be removed at some point. Use qsize() >= n + as a direct substitute, but be aware that either approach risks a race + condition where a queue can shrink before the result of full() or + qsize() can be used. + """ + with self.mutex: + return 0 < self.maxsize <= self._qsize() + + def put(self, item, block=True, timeout=None): + """Put an item into the queue. + + If optional args 'block' is true and 'timeout' is None (the default), + block if necessary until a free slot is available. If 'timeout' is + a non-negative number, it blocks at most 'timeout' seconds and raises + the FullError exception if no free slot was available within that time. + Otherwise ('block' is false), put an item on the queue if a free slot + is immediately available, else raise the FullError exception ('timeout' + is ignored in that case). + """ + with self.not_full: + if self.maxsize > 0: + if not block: + if self._qsize() >= self.maxsize: + raise FullError() + elif timeout is None: + while self._qsize() >= self.maxsize: + self.not_full.wait() + elif timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + endtime = time() + timeout + while self._qsize() >= self.maxsize: + remaining = endtime - time() + if remaining <= 0.0: + raise FullError() + self.not_full.wait(remaining) + self._put(item) + self.unfinished_tasks += 1 + self.not_empty.notify() + + def get(self, block=True, timeout=None): + """Remove and return an item from the queue. + + If optional args 'block' is true and 'timeout' is None (the default), + block if necessary until an item is available. If 'timeout' is + a non-negative number, it blocks at most 'timeout' seconds and raises + the EmptyError exception if no item was available within that time. + Otherwise ('block' is false), return an item if one is immediately + available, else raise the EmptyError exception ('timeout' is ignored + in that case). + """ + with self.not_empty: + if not block: + if not self._qsize(): + raise EmptyError() + elif timeout is None: + while not self._qsize(): + self.not_empty.wait() + elif timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + endtime = time() + timeout + while not self._qsize(): + remaining = endtime - time() + if remaining <= 0.0: + raise EmptyError() + self.not_empty.wait(remaining) + item = self._get() + self.not_full.notify() + return item + + def put_nowait(self, item): + """Put an item into the queue without blocking. + + Only enqueue the item if a free slot is immediately available. + Otherwise raise the FullError exception. + """ + return self.put(item, block=False) + + def get_nowait(self): + """Remove and return an item from the queue without blocking. + + Only get an item if one is immediately available. Otherwise + raise the EmptyError exception. + """ + return self.get(block=False) + + # Override these methods to implement other queue organizations + # (e.g. stack or priority queue). + # These will only be called with appropriate locks held + + # Initialize the queue representation + def _init(self, maxsize): + self.queue: "Any" = deque() + + def _qsize(self): + return len(self.queue) + + # Put a new item in the queue + def _put(self, item): + self.queue.append(item) + + # Get an item from the queue + def _get(self): + return self.queue.popleft() diff --git a/lib/python3.12/site-packages/sentry_sdk/_span_batcher.py b/lib/python3.12/site-packages/sentry_sdk/_span_batcher.py new file mode 100644 index 0000000000000000000000000000000000000000..38ecc8da51d08979da688fd16d5f4e0f75f4697e --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_span_batcher.py @@ -0,0 +1,126 @@ +import threading +from collections import defaultdict +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from sentry_sdk._batcher import Batcher +from sentry_sdk.consts import SPANSTATUS +from sentry_sdk.envelope import Envelope, Item, PayloadRef +from sentry_sdk.utils import format_timestamp, serialize_attribute, safe_repr + +if TYPE_CHECKING: + from typing import Any, Callable, Optional + from sentry_sdk.traces import StreamedSpan + from sentry_sdk._types import SerializedAttributeValue + + +class SpanBatcher(Batcher["StreamedSpan"]): + # TODO[span-first]: size-based flushes + # TODO[span-first]: adjust flush/drop defaults + MAX_BEFORE_FLUSH = 1000 + MAX_BEFORE_DROP = 5000 + FLUSH_WAIT_TIME = 5.0 + + TYPE = "span" + CONTENT_TYPE = "application/vnd.sentry.items.span.v2+json" + + def __init__( + self, + capture_func: "Callable[[Envelope], None]", + record_lost_func: "Callable[..., None]", + ) -> None: + # Spans from different traces cannot be emitted in the same envelope + # since the envelope contains a shared trace header. That's why we bucket + # by trace_id, so that we can then send the buckets each in its own + # envelope. + # trace_id -> span buffer + self._span_buffer: dict[str, list["StreamedSpan"]] = defaultdict(list) + self._capture_func = capture_func + self._record_lost_func = record_lost_func + self._running = True + self._lock = threading.Lock() + + self._flush_event: "threading.Event" = threading.Event() + + self._flusher: "Optional[threading.Thread]" = None + self._flusher_pid: "Optional[int]" = None + + def get_size(self) -> int: + # caller is responsible for locking before checking this + return sum(len(buffer) for buffer in self._span_buffer.values()) + + def add(self, span: "StreamedSpan") -> None: + if not self._ensure_thread() or self._flusher is None: + return None + + with self._lock: + size = self.get_size() + if size >= self.MAX_BEFORE_DROP: + self._record_lost_func( + reason="queue_overflow", + data_category="span", + quantity=1, + ) + return None + + self._span_buffer[span.trace_id].append(span) + if size + 1 >= self.MAX_BEFORE_FLUSH: + self._flush_event.set() + + @staticmethod + def _to_transport_format(item: "StreamedSpan") -> "Any": + # TODO[span-first] + res: "dict[str, Any]" = { + "name": item.name, + } + + if item._attributes: + res["attributes"] = { + k: serialize_attribute(v) for (k, v) in item._attributes.items() + } + + return res + + def _flush(self) -> None: + with self._lock: + if len(self._span_buffer) == 0: + return None + + envelopes = [] + for trace_id, spans in self._span_buffer.items(): + if spans: + # TODO[span-first] + # dsc = spans[0].dynamic_sampling_context() + dsc = None + + envelope = Envelope( + headers={ + "sent_at": format_timestamp(datetime.now(timezone.utc)), + "trace": dsc, + } + ) + + envelope.add_item( + Item( + type="span", + content_type="application/vnd.sentry.items.span.v2+json", + headers={ + "item_count": len(spans), + }, + payload=PayloadRef( + json={ + "items": [ + self._to_transport_format(span) + for span in spans + ] + } + ), + ) + ) + + envelopes.append(envelope) + + self._span_buffer.clear() + + for envelope in envelopes: + self._capture_func(envelope) diff --git a/lib/python3.12/site-packages/sentry_sdk/_types.py b/lib/python3.12/site-packages/sentry_sdk/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..e5f791fdf07ec926fd448c2fcdc03fe86d7d3f84 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_types.py @@ -0,0 +1,362 @@ +from typing import TYPE_CHECKING, TypeVar, Union + + +# Re-exported for compat, since code out there in the wild might use this variable. +MYPY = TYPE_CHECKING + + +SENSITIVE_DATA_SUBSTITUTE = "[Filtered]" +BLOB_DATA_SUBSTITUTE = "[Blob substitute]" + + +class AnnotatedValue: + """ + Meta information for a data field in the event payload. + This is to tell Relay that we have tampered with the fields value. + See: + https://github.com/getsentry/relay/blob/be12cd49a0f06ea932ed9b9f93a655de5d6ad6d1/relay-general/src/types/meta.rs#L407-L423 + """ + + __slots__ = ("value", "metadata") + + def __init__(self, value: "Optional[Any]", metadata: "Dict[str, Any]") -> None: + self.value = value + self.metadata = metadata + + def __eq__(self, other: "Any") -> bool: + if not isinstance(other, AnnotatedValue): + return False + + return self.value == other.value and self.metadata == other.metadata + + def __str__(self: "AnnotatedValue") -> str: + return str({"value": str(self.value), "metadata": str(self.metadata)}) + + def __len__(self: "AnnotatedValue") -> int: + if self.value is not None: + return len(self.value) + else: + return 0 + + @classmethod + def removed_because_raw_data(cls) -> "AnnotatedValue": + """The value was removed because it could not be parsed. This is done for request body values that are not json nor a form.""" + return AnnotatedValue( + value="", + metadata={ + "rem": [ # Remark + [ + "!raw", # Unparsable raw data + "x", # The fields original value was removed + ] + ] + }, + ) + + @classmethod + def removed_because_over_size_limit(cls, value: "Any" = "") -> "AnnotatedValue": + """ + The actual value was removed because the size of the field exceeded the configured maximum size, + for example specified with the max_request_body_size sdk option. + """ + return AnnotatedValue( + value=value, + metadata={ + "rem": [ # Remark + [ + "!config", # Because of configured maximum size + "x", # The fields original value was removed + ] + ] + }, + ) + + @classmethod + def substituted_because_contains_sensitive_data(cls) -> "AnnotatedValue": + """The actual value was removed because it contained sensitive information.""" + return AnnotatedValue( + value=SENSITIVE_DATA_SUBSTITUTE, + metadata={ + "rem": [ # Remark + [ + "!config", # Because of SDK configuration (in this case the config is the hard coded removal of certain django cookies) + "s", # The fields original value was substituted + ] + ] + }, + ) + + +T = TypeVar("T") +Annotated = Union[AnnotatedValue, T] + + +if TYPE_CHECKING: + from collections.abc import Container, MutableMapping, Sequence + + from datetime import datetime + + from types import TracebackType + from typing import Any + from typing import Callable + from typing import Dict + from typing import Mapping + from typing import NotRequired + from typing import Optional + from typing import Tuple + from typing import Type + from typing_extensions import Literal, TypedDict + + class SDKInfo(TypedDict): + name: str + version: str + packages: "Sequence[Mapping[str, str]]" + + # "critical" is an alias of "fatal" recognized by Relay + LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"] + + DurationUnit = Literal[ + "nanosecond", + "microsecond", + "millisecond", + "second", + "minute", + "hour", + "day", + "week", + ] + + InformationUnit = Literal[ + "bit", + "byte", + "kilobyte", + "kibibyte", + "megabyte", + "mebibyte", + "gigabyte", + "gibibyte", + "terabyte", + "tebibyte", + "petabyte", + "pebibyte", + "exabyte", + "exbibyte", + ] + + FractionUnit = Literal["ratio", "percent"] + MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str] + + MeasurementValue = TypedDict( + "MeasurementValue", + { + "value": float, + "unit": NotRequired[Optional[MeasurementUnit]], + }, + ) + + Event = TypedDict( + "Event", + { + "breadcrumbs": Annotated[ + dict[Literal["values"], list[dict[str, Any]]] + ], # TODO: We can expand on this type + "check_in_id": str, + "contexts": dict[str, dict[str, object]], + "dist": str, + "duration": Optional[float], + "environment": Optional[str], + "errors": list[dict[str, Any]], # TODO: We can expand on this type + "event_id": str, + "exception": dict[ + Literal["values"], list[dict[str, Any]] + ], # TODO: We can expand on this type + "extra": MutableMapping[str, object], + "fingerprint": list[str], + "level": LogLevelStr, + "logentry": Mapping[str, object], + "logger": str, + "measurements": dict[str, MeasurementValue], + "message": str, + "modules": dict[str, str], + "monitor_config": Mapping[str, object], + "monitor_slug": Optional[str], + "platform": Literal["python"], + "profile": object, # Should be sentry_sdk.profiler.Profile, but we can't import that here due to circular imports + "release": Optional[str], + "request": dict[str, object], + "sdk": Mapping[str, object], + "server_name": str, + "spans": Annotated[list[dict[str, object]]], + "stacktrace": dict[ + str, object + ], # We access this key in the code, but I am unsure whether we ever set it + "start_timestamp": datetime, + "status": Optional[str], + "tags": MutableMapping[ + str, str + ], # Tags must be less than 200 characters each + "threads": dict[ + Literal["values"], list[dict[str, Any]] + ], # TODO: We can expand on this type + "timestamp": Optional[datetime], # Must be set before sending the event + "transaction": str, + "transaction_info": Mapping[str, Any], # TODO: We can expand on this type + "type": Literal["check_in", "transaction"], + "user": dict[str, object], + "_dropped_spans": int, + }, + total=False, + ) + + ExcInfo = Union[ + tuple[Type[BaseException], BaseException, Optional[TracebackType]], + tuple[None, None, None], + ] + + # TODO: Make a proper type definition for this (PRs welcome!) + Hint = Dict[str, Any] + + AttributeValue = ( + str + | bool + | float + | int + | list[str] + | list[bool] + | list[float] + | list[int] + | tuple[str, ...] + | tuple[bool, ...] + | tuple[float, ...] + | tuple[int, ...] + ) + Attributes = dict[str, AttributeValue] + + SerializedAttributeValue = TypedDict( + # https://develop.sentry.dev/sdk/telemetry/attributes/#supported-types + "SerializedAttributeValue", + { + "type": Literal[ + "string", + "boolean", + "double", + "integer", + "array", + ], + "value": AttributeValue, + }, + ) + + Log = TypedDict( + "Log", + { + "severity_text": str, + "severity_number": int, + "body": str, + "attributes": Attributes, + "time_unix_nano": int, + "trace_id": Optional[str], + "span_id": Optional[str], + }, + ) + + MetricType = Literal["counter", "gauge", "distribution"] + MetricUnit = Union[DurationUnit, InformationUnit, str] + + Metric = TypedDict( + "Metric", + { + "timestamp": float, + "trace_id": Optional[str], + "span_id": Optional[str], + "name": str, + "type": MetricType, + "value": float, + "unit": Optional[MetricUnit], + "attributes": Attributes, + }, + ) + + MetricProcessor = Callable[[Metric, Hint], Optional[Metric]] + + # TODO: Make a proper type definition for this (PRs welcome!) + Breadcrumb = Dict[str, Any] + + # TODO: Make a proper type definition for this (PRs welcome!) + BreadcrumbHint = Dict[str, Any] + + # TODO: Make a proper type definition for this (PRs welcome!) + SamplingContext = Dict[str, Any] + + EventProcessor = Callable[[Event, Hint], Optional[Event]] + ErrorProcessor = Callable[[Event, ExcInfo], Optional[Event]] + BreadcrumbProcessor = Callable[[Breadcrumb, BreadcrumbHint], Optional[Breadcrumb]] + TransactionProcessor = Callable[[Event, Hint], Optional[Event]] + LogProcessor = Callable[[Log, Hint], Optional[Log]] + + TracesSampler = Callable[[SamplingContext], Union[float, int, bool]] + + # https://github.com/python/mypy/issues/5710 + NotImplementedType = Any + + EventDataCategory = Literal[ + "default", + "error", + "crash", + "transaction", + "security", + "attachment", + "session", + "internal", + "profile", + "profile_chunk", + "monitor", + "span", + "log_item", + "log_byte", + "trace_metric", + ] + SessionStatus = Literal["ok", "exited", "crashed", "abnormal"] + + ContinuousProfilerMode = Literal["thread", "gevent", "unknown"] + ProfilerMode = Union[ContinuousProfilerMode, Literal["sleep"]] + + MonitorConfigScheduleType = Literal["crontab", "interval"] + MonitorConfigScheduleUnit = Literal[ + "year", + "month", + "week", + "day", + "hour", + "minute", + "second", # not supported in Sentry and will result in a warning + ] + + MonitorConfigSchedule = TypedDict( + "MonitorConfigSchedule", + { + "type": MonitorConfigScheduleType, + "value": Union[int, str], + "unit": MonitorConfigScheduleUnit, + }, + total=False, + ) + + MonitorConfig = TypedDict( + "MonitorConfig", + { + "schedule": MonitorConfigSchedule, + "timezone": str, + "checkin_margin": int, + "max_runtime": int, + "failure_issue_threshold": int, + "recovery_threshold": int, + }, + total=False, + ) + + HttpStatusCodeRange = Union[int, Container[int]] + + class TextPart(TypedDict): + type: Literal["text"] + content: str diff --git a/lib/python3.12/site-packages/sentry_sdk/_werkzeug.py b/lib/python3.12/site-packages/sentry_sdk/_werkzeug.py new file mode 100644 index 0000000000000000000000000000000000000000..cdc3026c08acfcdb246a5d01f6e51cccac260038 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/_werkzeug.py @@ -0,0 +1,96 @@ +""" +Copyright (c) 2007 by the Pallets team. + +Some rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +* Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE AND DOCUMENTATION IS PROVIDED BY THE COPYRIGHT HOLDERS AND +CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, +BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT +NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF +USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF +THIS SOFTWARE AND DOCUMENTATION, EVEN IF ADVISED OF THE POSSIBILITY OF +SUCH DAMAGE. +""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Dict + from typing import Iterator + from typing import Tuple + + +# +# `get_headers` comes from `werkzeug.datastructures.EnvironHeaders` +# https://github.com/pallets/werkzeug/blob/0.14.1/werkzeug/datastructures.py#L1361 +# +# We need this function because Django does not give us a "pure" http header +# dict. So we might as well use it for all WSGI integrations. +# +def _get_headers(environ: "Dict[str, str]") -> "Iterator[Tuple[str, str]]": + """ + Returns only proper HTTP headers. + """ + for key, value in environ.items(): + key = str(key) + if key.startswith("HTTP_") and key not in ( + "HTTP_CONTENT_TYPE", + "HTTP_CONTENT_LENGTH", + ): + yield key[5:].replace("_", "-").title(), value + elif key in ("CONTENT_TYPE", "CONTENT_LENGTH"): + yield key.replace("_", "-").title(), value + + +# +# `get_host` comes from `werkzeug.wsgi.get_host` +# https://github.com/pallets/werkzeug/blob/1.0.1/src/werkzeug/wsgi.py#L145 +# +def get_host(environ: "Dict[str, str]", use_x_forwarded_for: bool = False) -> str: + """ + Return the host for the given WSGI environment. + """ + if use_x_forwarded_for and "HTTP_X_FORWARDED_HOST" in environ: + rv = environ["HTTP_X_FORWARDED_HOST"] + if environ["wsgi.url_scheme"] == "http" and rv.endswith(":80"): + rv = rv[:-3] + elif environ["wsgi.url_scheme"] == "https" and rv.endswith(":443"): + rv = rv[:-4] + elif environ.get("HTTP_HOST"): + rv = environ["HTTP_HOST"] + if environ["wsgi.url_scheme"] == "http" and rv.endswith(":80"): + rv = rv[:-3] + elif environ["wsgi.url_scheme"] == "https" and rv.endswith(":443"): + rv = rv[:-4] + elif environ.get("SERVER_NAME"): + rv = environ["SERVER_NAME"] + if (environ["wsgi.url_scheme"], environ["SERVER_PORT"]) not in ( + ("https", "443"), + ("http", "80"), + ): + rv += ":" + environ["SERVER_PORT"] + else: + # In spite of the WSGI spec, SERVER_NAME might not be present. + rv = "unknown" + + return rv diff --git a/lib/python3.12/site-packages/sentry_sdk/api.py b/lib/python3.12/site-packages/sentry_sdk/api.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e2229938341733744eede8f624a41f42cfb907 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/api.py @@ -0,0 +1,525 @@ +import inspect +import warnings +from contextlib import contextmanager + +from sentry_sdk import tracing_utils, Client +from sentry_sdk._init_implementation import init +from sentry_sdk.consts import INSTRUMENTER +from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope +from sentry_sdk.tracing import NoOpSpan, Transaction, trace +from sentry_sdk.crons import monitor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping + + from typing import Any + from typing import Dict + from typing import Generator + from typing import Optional + from typing import overload + from typing import Callable + from typing import TypeVar + from typing import ContextManager + from typing import Union + + from typing_extensions import Unpack + + from sentry_sdk.client import BaseClient + from sentry_sdk._types import ( + Event, + Hint, + Breadcrumb, + BreadcrumbHint, + ExcInfo, + MeasurementUnit, + LogLevelStr, + SamplingContext, + ) + from sentry_sdk.tracing import Span, TransactionKwargs + + T = TypeVar("T") + F = TypeVar("F", bound=Callable[..., Any]) +else: + + def overload(x: "T") -> "T": + return x + + +# When changing this, update __all__ in __init__.py too +__all__ = [ + "init", + "add_attachment", + "add_breadcrumb", + "capture_event", + "capture_exception", + "capture_message", + "configure_scope", + "continue_trace", + "flush", + "get_baggage", + "get_client", + "get_global_scope", + "get_isolation_scope", + "get_current_scope", + "get_current_span", + "get_traceparent", + "is_initialized", + "isolation_scope", + "last_event_id", + "new_scope", + "push_scope", + "set_context", + "set_extra", + "set_level", + "set_measurement", + "set_tag", + "set_tags", + "set_user", + "start_span", + "start_transaction", + "trace", + "monitor", + "start_session", + "end_session", + "set_transaction_name", + "update_current_span", +] + + +def scopemethod(f: "F") -> "F": + f.__doc__ = "%s\n\n%s" % ( + "Alias for :py:meth:`sentry_sdk.Scope.%s`" % f.__name__, + inspect.getdoc(getattr(Scope, f.__name__)), + ) + return f + + +def clientmethod(f: "F") -> "F": + f.__doc__ = "%s\n\n%s" % ( + "Alias for :py:meth:`sentry_sdk.Client.%s`" % f.__name__, + inspect.getdoc(getattr(Client, f.__name__)), + ) + return f + + +@scopemethod +def get_client() -> "BaseClient": + return Scope.get_client() + + +def is_initialized() -> bool: + """ + .. versionadded:: 2.0.0 + + Returns whether Sentry has been initialized or not. + + If a client is available and the client is active + (meaning it is configured to send data) then + Sentry is initialized. + """ + return get_client().is_active() + + +@scopemethod +def get_global_scope() -> "Scope": + return Scope.get_global_scope() + + +@scopemethod +def get_isolation_scope() -> "Scope": + return Scope.get_isolation_scope() + + +@scopemethod +def get_current_scope() -> "Scope": + return Scope.get_current_scope() + + +@scopemethod +def last_event_id() -> "Optional[str]": + """ + See :py:meth:`sentry_sdk.Scope.last_event_id` documentation regarding + this method's limitations. + """ + return Scope.last_event_id() + + +@scopemethod +def capture_event( + event: "Event", + hint: "Optional[Hint]" = None, + scope: "Optional[Any]" = None, + **scope_kwargs: "Any", +) -> "Optional[str]": + return get_current_scope().capture_event(event, hint, scope=scope, **scope_kwargs) + + +@scopemethod +def capture_message( + message: str, + level: "Optional[LogLevelStr]" = None, + scope: "Optional[Any]" = None, + **scope_kwargs: "Any", +) -> "Optional[str]": + return get_current_scope().capture_message( + message, level, scope=scope, **scope_kwargs + ) + + +@scopemethod +def capture_exception( + error: "Optional[Union[BaseException, ExcInfo]]" = None, + scope: "Optional[Any]" = None, + **scope_kwargs: "Any", +) -> "Optional[str]": + return get_current_scope().capture_exception(error, scope=scope, **scope_kwargs) + + +@scopemethod +def add_attachment( + bytes: "Union[None, bytes, Callable[[], bytes]]" = None, + filename: "Optional[str]" = None, + path: "Optional[str]" = None, + content_type: "Optional[str]" = None, + add_to_transactions: bool = False, +) -> None: + return get_isolation_scope().add_attachment( + bytes, filename, path, content_type, add_to_transactions + ) + + +@scopemethod +def add_breadcrumb( + crumb: "Optional[Breadcrumb]" = None, + hint: "Optional[BreadcrumbHint]" = None, + **kwargs: "Any", +) -> None: + return get_isolation_scope().add_breadcrumb(crumb, hint, **kwargs) + + +@overload +def configure_scope() -> "ContextManager[Scope]": + pass + + +@overload +def configure_scope( # noqa: F811 + callback: "Callable[[Scope], None]", +) -> None: + pass + + +def configure_scope( # noqa: F811 + callback: "Optional[Callable[[Scope], None]]" = None, +) -> "Optional[ContextManager[Scope]]": + """ + Reconfigures the scope. + + :param callback: If provided, call the callback with the current scope. + + :returns: If no callback is provided, returns a context manager that returns the scope. + """ + warnings.warn( + "sentry_sdk.configure_scope is deprecated and will be removed in the next major version. " + "Please consult our migration guide to learn how to migrate to the new API: " + "https://docs.sentry.io/platforms/python/migration/1.x-to-2.x#scope-configuring", + DeprecationWarning, + stacklevel=2, + ) + + scope = get_isolation_scope() + scope.generate_propagation_context() + + if callback is not None: + # TODO: used to return None when client is None. Check if this changes behavior. + callback(scope) + + return None + + @contextmanager + def inner() -> "Generator[Scope, None, None]": + yield scope + + return inner() + + +@overload +def push_scope() -> "ContextManager[Scope]": + pass + + +@overload +def push_scope( # noqa: F811 + callback: "Callable[[Scope], None]", +) -> None: + pass + + +def push_scope( # noqa: F811 + callback: "Optional[Callable[[Scope], None]]" = None, +) -> "Optional[ContextManager[Scope]]": + """ + Pushes a new layer on the scope stack. + + :param callback: If provided, this method pushes a scope, calls + `callback`, and pops the scope again. + + :returns: If no `callback` is provided, a context manager that should + be used to pop the scope again. + """ + warnings.warn( + "sentry_sdk.push_scope is deprecated and will be removed in the next major version. " + "Please consult our migration guide to learn how to migrate to the new API: " + "https://docs.sentry.io/platforms/python/migration/1.x-to-2.x#scope-pushing", + DeprecationWarning, + stacklevel=2, + ) + + if callback is not None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + with push_scope() as scope: + callback(scope) + return None + + return _ScopeManager() + + +@scopemethod +def set_tag(key: str, value: "Any") -> None: + return get_isolation_scope().set_tag(key, value) + + +@scopemethod +def set_tags(tags: "Mapping[str, object]") -> None: + return get_isolation_scope().set_tags(tags) + + +@scopemethod +def set_context(key: str, value: "Dict[str, Any]") -> None: + return get_isolation_scope().set_context(key, value) + + +@scopemethod +def set_extra(key: str, value: "Any") -> None: + return get_isolation_scope().set_extra(key, value) + + +@scopemethod +def set_user(value: "Optional[Dict[str, Any]]") -> None: + return get_isolation_scope().set_user(value) + + +@scopemethod +def set_level(value: "LogLevelStr") -> None: + return get_isolation_scope().set_level(value) + + +@clientmethod +def flush( + timeout: "Optional[float]" = None, + callback: "Optional[Callable[[int, float], None]]" = None, +) -> None: + return get_client().flush(timeout=timeout, callback=callback) + + +@scopemethod +def start_span( + **kwargs: "Any", +) -> "Span": + return get_current_scope().start_span(**kwargs) + + +@scopemethod +def start_transaction( + transaction: "Optional[Transaction]" = None, + instrumenter: str = INSTRUMENTER.SENTRY, + custom_sampling_context: "Optional[SamplingContext]" = None, + **kwargs: "Unpack[TransactionKwargs]", +) -> "Union[Transaction, NoOpSpan]": + """ + Start and return a transaction on the current scope. + + Start an existing transaction if given, otherwise create and start a new + transaction with kwargs. + + This is the entry point to manual tracing instrumentation. + + A tree structure can be built by adding child spans to the transaction, + and child spans to other spans. To start a new child span within the + transaction or any span, call the respective `.start_child()` method. + + Every child span must be finished before the transaction is finished, + otherwise the unfinished spans are discarded. + + When used as context managers, spans and transactions are automatically + finished at the end of the `with` block. If not using context managers, + call the `.finish()` method. + + When the transaction is finished, it will be sent to Sentry with all its + finished child spans. + + :param transaction: The transaction to start. If omitted, we create and + start a new transaction. + :param instrumenter: This parameter is meant for internal use only. It + will be removed in the next major version. + :param custom_sampling_context: The transaction's custom sampling context. + :param kwargs: Optional keyword arguments to be passed to the Transaction + constructor. See :py:class:`sentry_sdk.tracing.Transaction` for + available arguments. + """ + return get_current_scope().start_transaction( + transaction, instrumenter, custom_sampling_context, **kwargs + ) + + +def set_measurement(name: str, value: float, unit: "MeasurementUnit" = "") -> None: + """ + .. deprecated:: 2.28.0 + This function is deprecated and will be removed in the next major release. + """ + transaction = get_current_scope().transaction + if transaction is not None: + transaction.set_measurement(name, value, unit) + + +def get_current_span(scope: "Optional[Scope]" = None) -> "Optional[Span]": + """ + Returns the currently active span if there is one running, otherwise `None` + """ + return tracing_utils.get_current_span(scope) + + +def get_traceparent() -> "Optional[str]": + """ + Returns the traceparent either from the active span or from the scope. + """ + return get_current_scope().get_traceparent() + + +def get_baggage() -> "Optional[str]": + """ + Returns Baggage either from the active span or from the scope. + """ + baggage = get_current_scope().get_baggage() + if baggage is not None: + return baggage.serialize() + + return None + + +def continue_trace( + environ_or_headers: "Dict[str, Any]", + op: "Optional[str]" = None, + name: "Optional[str]" = None, + source: "Optional[str]" = None, + origin: str = "manual", +) -> "Transaction": + """ + Sets the propagation context from environment or headers and returns a transaction. + """ + return get_isolation_scope().continue_trace( + environ_or_headers, op, name, source, origin + ) + + +@scopemethod +def start_session( + session_mode: str = "application", +) -> None: + return get_isolation_scope().start_session(session_mode=session_mode) + + +@scopemethod +def end_session() -> None: + return get_isolation_scope().end_session() + + +@scopemethod +def set_transaction_name(name: str, source: "Optional[str]" = None) -> None: + return get_current_scope().set_transaction_name(name, source) + + +def update_current_span( + op: "Optional[str]" = None, + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Union[str, int, float, bool]]]" = None, + data: "Optional[dict[str, Any]]" = None, +) -> None: + """ + Update the current active span with the provided parameters. + + This function allows you to modify properties of the currently active span. + If no span is currently active, this function will do nothing. + + :param op: The operation name for the span. This is a high-level description + of what the span represents (e.g., "http.client", "db.query"). + You can use predefined constants from :py:class:`sentry_sdk.consts.OP` + or provide your own string. If not provided, the span's operation will + remain unchanged. + :type op: str or None + + :param name: The human-readable name/description for the span. This provides + more specific details about what the span represents (e.g., "GET /api/users", + "SELECT * FROM users"). If not provided, the span's name will remain unchanged. + :type name: str or None + + :param data: A dictionary of key-value pairs to add as data to the span. This + data will be merged with any existing span data. If not provided, + no data will be added. + + .. deprecated:: 2.35.0 + Use ``attributes`` instead. The ``data`` parameter will be removed + in a future version. + :type data: dict[str, Union[str, int, float, bool]] or None + + :param attributes: A dictionary of key-value pairs to add as attributes to the span. + Attribute values must be strings, integers, floats, or booleans. These + attributes will be merged with any existing span data. If not provided, + no attributes will be added. + :type attributes: dict[str, Union[str, int, float, bool]] or None + + :returns: None + + .. versionadded:: 2.35.0 + + Example:: + + import sentry_sdk + from sentry_sdk.consts import OP + + sentry_sdk.update_current_span( + op=OP.FUNCTION, + name="process_user_data", + attributes={"user_id": 123, "batch_size": 50} + ) + """ + current_span = get_current_span() + + if current_span is None: + return + + if op is not None: + current_span.op = op + + if name is not None: + # internally it is still description + current_span.description = name + + if data is not None and attributes is not None: + raise ValueError( + "Cannot provide both `data` and `attributes`. Please use only `attributes`." + ) + + if data is not None: + warnings.warn( + "The `data` parameter is deprecated. Please use `attributes` instead.", + DeprecationWarning, + stacklevel=2, + ) + attributes = data + + if attributes is not None: + current_span.update_data(attributes) diff --git a/lib/python3.12/site-packages/sentry_sdk/attachments.py b/lib/python3.12/site-packages/sentry_sdk/attachments.py new file mode 100644 index 0000000000000000000000000000000000000000..8ad85f433512c2a89683693bfece57133ffe229f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/attachments.py @@ -0,0 +1,72 @@ +import os +import mimetypes + +from sentry_sdk.envelope import Item, PayloadRef + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional, Union, Callable + + +class Attachment: + """Additional files/data to send along with an event. + + This class stores attachments that can be sent along with an event. Attachments are files or other data, e.g. + config or log files, that are relevant to an event. Attachments are set on the ``Scope``, and are sent along with + all non-transaction events (or all events including transactions if ``add_to_transactions`` is ``True``) that are + captured within the ``Scope``. + + To add an attachment to a ``Scope``, use :py:meth:`sentry_sdk.Scope.add_attachment`. The parameters for + ``add_attachment`` are the same as the parameters for this class's constructor. + + :param bytes: Raw bytes of the attachment, or a function that returns the raw bytes. Must be provided unless + ``path`` is provided. + :param filename: The filename of the attachment. Must be provided unless ``path`` is provided. + :param path: Path to a file to attach. Must be provided unless ``bytes`` is provided. + :param content_type: The content type of the attachment. If not provided, it will be guessed from the ``filename`` + parameter, if available, or the ``path`` parameter if ``filename`` is ``None``. + :param add_to_transactions: Whether to add this attachment to transactions. Defaults to ``False``. + """ + + def __init__( + self, + bytes: "Union[None, bytes, Callable[[], bytes]]" = None, + filename: "Optional[str]" = None, + path: "Optional[str]" = None, + content_type: "Optional[str]" = None, + add_to_transactions: bool = False, + ) -> None: + if bytes is None and path is None: + raise TypeError("path or raw bytes required for attachment") + if filename is None and path is not None: + filename = os.path.basename(path) + if filename is None: + raise TypeError("filename is required for attachment") + if content_type is None: + content_type = mimetypes.guess_type(filename)[0] + self.bytes = bytes + self.filename = filename + self.path = path + self.content_type = content_type + self.add_to_transactions = add_to_transactions + + def to_envelope_item(self) -> "Item": + """Returns an envelope item for this attachment.""" + payload: "Union[None, PayloadRef, bytes]" = None + if self.bytes is not None: + if callable(self.bytes): + payload = self.bytes() + else: + payload = self.bytes + else: + payload = PayloadRef(path=self.path) + return Item( + payload=payload, + type="attachment", + content_type=self.content_type, + filename=self.filename, + ) + + def __repr__(self) -> str: + return "" % (self.filename,) diff --git a/lib/python3.12/site-packages/sentry_sdk/client.py b/lib/python3.12/site-packages/sentry_sdk/client.py new file mode 100644 index 0000000000000000000000000000000000000000..f540a49f352187b9e94f9ef4fe86651ffee7a88a --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/client.py @@ -0,0 +1,1080 @@ +import os +import uuid +import random +import socket +from collections.abc import Mapping +from datetime import datetime, timezone +from importlib import import_module +from typing import TYPE_CHECKING, List, Dict, cast, overload +import warnings + +import sentry_sdk +from sentry_sdk._compat import PY37, check_uwsgi_thread_support +from sentry_sdk._metrics_batcher import MetricsBatcher +from sentry_sdk._span_batcher import SpanBatcher +from sentry_sdk.utils import ( + AnnotatedValue, + ContextVar, + capture_internal_exceptions, + current_stacktrace, + env_to_bool, + format_timestamp, + get_sdk_name, + get_type_name, + get_default_release, + handle_in_app, + is_gevent, + logger, + get_before_send_log, + get_before_send_metric, + has_logs_enabled, + has_metrics_enabled, +) +from sentry_sdk.serializer import serialize +from sentry_sdk.tracing import trace +from sentry_sdk.tracing_utils import has_span_streaming_enabled +from sentry_sdk.transport import BaseHttpTransport, make_transport +from sentry_sdk.consts import ( + SPANDATA, + DEFAULT_MAX_VALUE_LENGTH, + DEFAULT_OPTIONS, + INSTRUMENTER, + VERSION, + ClientConstructor, +) +from sentry_sdk.integrations import _DEFAULT_INTEGRATIONS, setup_integrations +from sentry_sdk.integrations.dedupe import DedupeIntegration +from sentry_sdk.sessions import SessionFlusher +from sentry_sdk.envelope import Envelope +from sentry_sdk.profiler.continuous_profiler import setup_continuous_profiler +from sentry_sdk.profiler.transaction_profiler import ( + has_profiling_enabled, + Profile, + setup_profiler, +) +from sentry_sdk.scrubber import EventScrubber +from sentry_sdk.monitor import Monitor + +if TYPE_CHECKING: + from typing import Any + from typing import Callable + from typing import Optional + from typing import Sequence + from typing import Type + from typing import Union + from typing import TypeVar + + from sentry_sdk._types import Event, Hint, SDKInfo, Log, Metric, EventDataCategory + from sentry_sdk.integrations import Integration + from sentry_sdk.scope import Scope + from sentry_sdk.session import Session + from sentry_sdk.spotlight import SpotlightClient + from sentry_sdk.traces import StreamedSpan + from sentry_sdk.transport import Transport, Item + from sentry_sdk._log_batcher import LogBatcher + from sentry_sdk._metrics_batcher import MetricsBatcher + from sentry_sdk.utils import Dsn + + I = TypeVar("I", bound=Integration) # noqa: E741 + +_client_init_debug = ContextVar("client_init_debug") + + +SDK_INFO: "SDKInfo" = { + "name": "sentry.python", # SDK name will be overridden after integrations have been loaded with sentry_sdk.integrations.setup_integrations() + "version": VERSION, + "packages": [{"name": "pypi:sentry-sdk", "version": VERSION}], +} + + +def _get_options(*args: "Optional[str]", **kwargs: "Any") -> "Dict[str, Any]": + if args and (isinstance(args[0], (bytes, str)) or args[0] is None): + dsn: "Optional[str]" = args[0] + args = args[1:] + else: + dsn = None + + if len(args) > 1: + raise TypeError("Only single positional argument is expected") + + rv = dict(DEFAULT_OPTIONS) + options = dict(*args, **kwargs) + if dsn is not None and options.get("dsn") is None: + options["dsn"] = dsn + + for key, value in options.items(): + if key not in rv: + raise TypeError("Unknown option %r" % (key,)) + + rv[key] = value + + if rv["dsn"] is None: + rv["dsn"] = os.environ.get("SENTRY_DSN") + + if rv["release"] is None: + rv["release"] = get_default_release() + + if rv["environment"] is None: + rv["environment"] = os.environ.get("SENTRY_ENVIRONMENT") or "production" + + if rv["debug"] is None: + rv["debug"] = env_to_bool(os.environ.get("SENTRY_DEBUG"), strict=True) or False + + if rv["server_name"] is None and hasattr(socket, "gethostname"): + rv["server_name"] = socket.gethostname() + + if rv["instrumenter"] is None: + rv["instrumenter"] = INSTRUMENTER.SENTRY + + if rv["project_root"] is None: + try: + project_root = os.getcwd() + except Exception: + project_root = None + + rv["project_root"] = project_root + + if rv["enable_tracing"] is True and rv["traces_sample_rate"] is None: + rv["traces_sample_rate"] = 1.0 + + if rv["event_scrubber"] is None: + rv["event_scrubber"] = EventScrubber( + send_default_pii=( + False if rv["send_default_pii"] is None else rv["send_default_pii"] + ) + ) + + if rv["socket_options"] and not isinstance(rv["socket_options"], list): + logger.warning( + "Ignoring socket_options because of unexpected format. See urllib3.HTTPConnection.socket_options for the expected format." + ) + rv["socket_options"] = None + + if rv["keep_alive"] is None: + rv["keep_alive"] = ( + env_to_bool(os.environ.get("SENTRY_KEEP_ALIVE"), strict=True) or False + ) + + if rv["enable_tracing"] is not None: + warnings.warn( + "The `enable_tracing` parameter is deprecated. Please use `traces_sample_rate` instead.", + DeprecationWarning, + stacklevel=2, + ) + + return rv + + +try: + # Python 3.6+ + module_not_found_error = ModuleNotFoundError +except Exception: + # Older Python versions + module_not_found_error = ImportError # type: ignore + + +class BaseClient: + """ + .. versionadded:: 2.0.0 + + The basic definition of a client that is used for sending data to Sentry. + """ + + spotlight: "Optional[SpotlightClient]" = None + + def __init__(self, options: "Optional[Dict[str, Any]]" = None) -> None: + self.options: "Dict[str, Any]" = ( + options if options is not None else DEFAULT_OPTIONS + ) + + self.transport: "Optional[Transport]" = None + self.monitor: "Optional[Monitor]" = None + self.log_batcher: "Optional[LogBatcher]" = None + self.metrics_batcher: "Optional[MetricsBatcher]" = None + self.span_batcher: "Optional[SpanBatcher]" = None + self.integrations: "dict[str, Integration]" = {} + + def __getstate__(self, *args: "Any", **kwargs: "Any") -> "Any": + return {"options": {}} + + def __setstate__(self, *args: "Any", **kwargs: "Any") -> None: + pass + + @property + def dsn(self) -> "Optional[str]": + return None + + @property + def parsed_dsn(self) -> "Optional[Dsn]": + return None + + def should_send_default_pii(self) -> bool: + return False + + def is_active(self) -> bool: + """ + .. versionadded:: 2.0.0 + + Returns whether the client is active (able to send data to Sentry) + """ + return False + + def capture_event(self, *args: "Any", **kwargs: "Any") -> "Optional[str]": + return None + + def _capture_log(self, log: "Log", scope: "Scope") -> None: + pass + + def _capture_metric(self, metric: "Metric", scope: "Scope") -> None: + pass + + def _capture_span(self, span: "StreamedSpan", scope: "Scope") -> None: + pass + + def capture_session(self, *args: "Any", **kwargs: "Any") -> None: + return None + + if TYPE_CHECKING: + + @overload + def get_integration(self, name_or_class: str) -> "Optional[Integration]": ... + + @overload + def get_integration(self, name_or_class: "type[I]") -> "Optional[I]": ... + + def get_integration( + self, name_or_class: "Union[str, type[Integration]]" + ) -> "Optional[Integration]": + return None + + def close(self, *args: "Any", **kwargs: "Any") -> None: + return None + + def flush(self, *args: "Any", **kwargs: "Any") -> None: + return None + + def __enter__(self) -> "BaseClient": + return self + + def __exit__(self, exc_type: "Any", exc_value: "Any", tb: "Any") -> None: + return None + + +class NonRecordingClient(BaseClient): + """ + .. versionadded:: 2.0.0 + + A client that does not send any events to Sentry. This is used as a fallback when the Sentry SDK is not yet initialized. + """ + + pass + + +class _Client(BaseClient): + """ + The client is internally responsible for capturing the events and + forwarding them to sentry through the configured transport. It takes + the client options as keyword arguments and optionally the DSN as first + argument. + + Alias of :py:class:`sentry_sdk.Client`. (Was created for better intelisense support) + """ + + def __init__(self, *args: "Any", **kwargs: "Any") -> None: + super(_Client, self).__init__(options=get_options(*args, **kwargs)) + self._init_impl() + + def __getstate__(self) -> "Any": + return {"options": self.options} + + def __setstate__(self, state: "Any") -> None: + self.options = state["options"] + self._init_impl() + + def _setup_instrumentation( + self, functions_to_trace: "Sequence[Dict[str, str]]" + ) -> None: + """ + Instruments the functions given in the list `functions_to_trace` with the `@sentry_sdk.tracing.trace` decorator. + """ + for function in functions_to_trace: + class_name = None + function_qualname = function["qualified_name"] + module_name, function_name = function_qualname.rsplit(".", 1) + + try: + # Try to import module and function + # ex: "mymodule.submodule.funcname" + + module_obj = import_module(module_name) + function_obj = getattr(module_obj, function_name) + setattr(module_obj, function_name, trace(function_obj)) + logger.debug("Enabled tracing for %s", function_qualname) + except module_not_found_error: + try: + # Try to import a class + # ex: "mymodule.submodule.MyClassName.member_function" + + module_name, class_name = module_name.rsplit(".", 1) + module_obj = import_module(module_name) + class_obj = getattr(module_obj, class_name) + function_obj = getattr(class_obj, function_name) + function_type = type(class_obj.__dict__[function_name]) + traced_function = trace(function_obj) + + if function_type in (staticmethod, classmethod): + traced_function = staticmethod(traced_function) + + setattr(class_obj, function_name, traced_function) + setattr(module_obj, class_name, class_obj) + logger.debug("Enabled tracing for %s", function_qualname) + + except Exception as e: + logger.warning( + "Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.", + function_qualname, + e, + ) + + except Exception as e: + logger.warning( + "Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.", + function_qualname, + e, + ) + + def _init_impl(self) -> None: + old_debug = _client_init_debug.get(False) + + def _capture_envelope(envelope: "Envelope") -> None: + if self.spotlight is not None: + self.spotlight.capture_envelope(envelope) + if self.transport is not None: + self.transport.capture_envelope(envelope) + + def _record_lost_event( + reason: str, + data_category: "EventDataCategory", + item: "Optional[Item]" = None, + quantity: int = 1, + ) -> None: + if self.transport is not None: + self.transport.record_lost_event( + reason=reason, + data_category=data_category, + item=item, + quantity=quantity, + ) + + try: + _client_init_debug.set(self.options["debug"]) + self.transport = make_transport(self.options) + + self.monitor = None + if self.transport: + if self.options["enable_backpressure_handling"]: + self.monitor = Monitor(self.transport) + + # Setup Spotlight before creating batchers so _capture_envelope can use it. + # setup_spotlight handles all config/env var resolution per the SDK spec. + from sentry_sdk.spotlight import setup_spotlight + + self.spotlight = setup_spotlight(self.options) + if self.spotlight is not None and not self.options["dsn"]: + sample_all = lambda *_args, **_kwargs: 1.0 + self.options["send_default_pii"] = True + self.options["error_sampler"] = sample_all + self.options["traces_sampler"] = sample_all + self.options["profiles_sampler"] = sample_all + + self.session_flusher = SessionFlusher(capture_func=_capture_envelope) + + self.log_batcher = None + + if has_logs_enabled(self.options): + from sentry_sdk._log_batcher import LogBatcher + + self.log_batcher = LogBatcher( + capture_func=_capture_envelope, + record_lost_func=_record_lost_event, + ) + + self.metrics_batcher = None + if has_metrics_enabled(self.options): + self.metrics_batcher = MetricsBatcher( + capture_func=_capture_envelope, + record_lost_func=_record_lost_event, + ) + + self.span_batcher = None + if has_span_streaming_enabled(self.options): + self.span_batcher = SpanBatcher( + capture_func=_capture_envelope, + record_lost_func=_record_lost_event, + ) + + max_request_body_size = ("always", "never", "small", "medium") + if self.options["max_request_body_size"] not in max_request_body_size: + raise ValueError( + "Invalid value for max_request_body_size. Must be one of {}".format( + max_request_body_size + ) + ) + + if self.options["_experiments"].get("otel_powered_performance", False): + logger.debug( + "[OTel] Enabling experimental OTel-powered performance monitoring." + ) + self.options["instrumenter"] = INSTRUMENTER.OTEL + if ( + "sentry_sdk.integrations.opentelemetry.integration.OpenTelemetryIntegration" + not in _DEFAULT_INTEGRATIONS + ): + _DEFAULT_INTEGRATIONS.append( + "sentry_sdk.integrations.opentelemetry.integration.OpenTelemetryIntegration", + ) + + self.integrations = setup_integrations( + self.options["integrations"], + with_defaults=self.options["default_integrations"], + with_auto_enabling_integrations=self.options[ + "auto_enabling_integrations" + ], + disabled_integrations=self.options["disabled_integrations"], + options=self.options, + ) + + sdk_name = get_sdk_name(list(self.integrations.keys())) + SDK_INFO["name"] = sdk_name + logger.debug("Setting SDK name to '%s'", sdk_name) + + if has_profiling_enabled(self.options): + try: + setup_profiler(self.options) + except Exception as e: + logger.debug("Can not set up profiler. (%s)", e) + else: + try: + setup_continuous_profiler( + self.options, + sdk_info=SDK_INFO, + capture_func=_capture_envelope, + ) + except Exception as e: + logger.debug("Can not set up continuous profiler. (%s)", e) + + finally: + _client_init_debug.set(old_debug) + + self._setup_instrumentation(self.options.get("functions_to_trace", [])) + + if ( + self.monitor + or self.log_batcher + or self.metrics_batcher + or self.span_batcher + or has_profiling_enabled(self.options) + or isinstance(self.transport, BaseHttpTransport) + ): + # If we have anything on that could spawn a background thread, we + # need to check if it's safe to use them. + check_uwsgi_thread_support() + + def is_active(self) -> bool: + """ + .. versionadded:: 2.0.0 + + Returns whether the client is active (able to send data to Sentry) + """ + return True + + def should_send_default_pii(self) -> bool: + """ + .. versionadded:: 2.0.0 + + Returns whether the client should send default PII (Personally Identifiable Information) data to Sentry. + """ + return self.options.get("send_default_pii") or False + + @property + def dsn(self) -> "Optional[str]": + """Returns the configured DSN as string.""" + return self.options["dsn"] + + @property + def parsed_dsn(self) -> "Optional[Dsn]": + """Returns the configured parsed DSN object.""" + return self.transport.parsed_dsn if self.transport else None + + def _prepare_event( + self, + event: "Event", + hint: "Hint", + scope: "Optional[Scope]", + ) -> "Optional[Event]": + previous_total_spans: "Optional[int]" = None + previous_total_breadcrumbs: "Optional[int]" = None + + if event.get("timestamp") is None: + event["timestamp"] = datetime.now(timezone.utc) + + is_transaction = event.get("type") == "transaction" + + if scope is not None: + spans_before = len(cast(List[Dict[str, object]], event.get("spans", []))) + event_ = scope.apply_to_event(event, hint, self.options) + + # one of the event/error processors returned None + if event_ is None: + if self.transport: + self.transport.record_lost_event( + "event_processor", + data_category=("transaction" if is_transaction else "error"), + ) + if is_transaction: + self.transport.record_lost_event( + "event_processor", + data_category="span", + quantity=spans_before + 1, # +1 for the transaction itself + ) + return None + + event = event_ + spans_delta = spans_before - len( + cast(List[Dict[str, object]], event.get("spans", [])) + ) + span_recorder_dropped_spans: int = event.pop("_dropped_spans", 0) + + if is_transaction and self.transport is not None: + if spans_delta > 0: + self.transport.record_lost_event( + "event_processor", data_category="span", quantity=spans_delta + ) + if span_recorder_dropped_spans > 0: + self.transport.record_lost_event( + "buffer_overflow", + data_category="span", + quantity=span_recorder_dropped_spans, + ) + + dropped_spans: int = span_recorder_dropped_spans + spans_delta + if dropped_spans > 0: + previous_total_spans = spans_before + dropped_spans + if scope._n_breadcrumbs_truncated > 0: + breadcrumbs = event.get("breadcrumbs", {}) + values = ( + breadcrumbs.get("values", []) + if not isinstance(breadcrumbs, AnnotatedValue) + else [] + ) + previous_total_breadcrumbs = ( + len(values) + scope._n_breadcrumbs_truncated + ) + + if ( + not is_transaction + and self.options["attach_stacktrace"] + and "exception" not in event + and "stacktrace" not in event + and "threads" not in event + ): + with capture_internal_exceptions(): + event["threads"] = { + "values": [ + { + "stacktrace": current_stacktrace( + include_local_variables=self.options.get( + "include_local_variables", True + ), + max_value_length=self.options.get( + "max_value_length", DEFAULT_MAX_VALUE_LENGTH + ), + ), + "crashed": False, + "current": True, + } + ] + } + + for key in "release", "environment", "server_name", "dist": + if event.get(key) is None and self.options[key] is not None: + event[key] = str(self.options[key]).strip() + if event.get("sdk") is None: + sdk_info = dict(SDK_INFO) + sdk_info["integrations"] = sorted(self.integrations.keys()) + event["sdk"] = sdk_info + + if event.get("platform") is None: + event["platform"] = "python" + + event = handle_in_app( + event, + self.options["in_app_exclude"], + self.options["in_app_include"], + self.options["project_root"], + ) + + if event is not None: + event_scrubber = self.options["event_scrubber"] + if event_scrubber: + event_scrubber.scrub_event(event) + + if scope is not None and scope._gen_ai_original_message_count: + spans: "List[Dict[str, Any]] | AnnotatedValue" = event.get("spans", []) + if isinstance(spans, list): + for span in spans: + span_id = span.get("span_id", None) + span_data = span.get("data", {}) + if ( + span_id + and span_id in scope._gen_ai_original_message_count + and SPANDATA.GEN_AI_REQUEST_MESSAGES in span_data + ): + span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES] = AnnotatedValue( + span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES], + {"len": scope._gen_ai_original_message_count[span_id]}, + ) + if previous_total_spans is not None: + event["spans"] = AnnotatedValue( + event.get("spans", []), {"len": previous_total_spans} + ) + if previous_total_breadcrumbs is not None: + event["breadcrumbs"] = AnnotatedValue( + event.get("breadcrumbs", {"values": []}), + {"len": previous_total_breadcrumbs}, + ) + + # Postprocess the event here so that annotated types do + # generally not surface in before_send + if event is not None: + event = cast( + "Event", + serialize( + cast("Dict[str, Any]", event), + max_request_body_size=self.options.get("max_request_body_size"), + max_value_length=self.options.get("max_value_length"), + custom_repr=self.options.get("custom_repr"), + ), + ) + + before_send = self.options["before_send"] + if ( + before_send is not None + and event is not None + and event.get("type") != "transaction" + ): + new_event = None + with capture_internal_exceptions(): + new_event = before_send(event, hint or {}) + if new_event is None: + logger.info("before send dropped event") + if self.transport: + self.transport.record_lost_event( + "before_send", data_category="error" + ) + + # If this is an exception, reset the DedupeIntegration. It still + # remembers the dropped exception as the last exception, meaning + # that if the same exception happens again and is not dropped + # in before_send, it'd get dropped by DedupeIntegration. + if event.get("exception"): + DedupeIntegration.reset_last_seen() + + event = new_event + + before_send_transaction = self.options["before_send_transaction"] + if ( + before_send_transaction is not None + and event is not None + and event.get("type") == "transaction" + ): + new_event = None + spans_before = len(cast(List[Dict[str, object]], event.get("spans", []))) + with capture_internal_exceptions(): + new_event = before_send_transaction(event, hint or {}) + if new_event is None: + logger.info("before send transaction dropped event") + if self.transport: + self.transport.record_lost_event( + reason="before_send", data_category="transaction" + ) + self.transport.record_lost_event( + reason="before_send", + data_category="span", + quantity=spans_before + 1, # +1 for the transaction itself + ) + else: + spans_delta = spans_before - len(new_event.get("spans", [])) + if spans_delta > 0 and self.transport is not None: + self.transport.record_lost_event( + reason="before_send", data_category="span", quantity=spans_delta + ) + + event = new_event + + return event + + def _is_ignored_error(self, event: "Event", hint: "Hint") -> bool: + exc_info = hint.get("exc_info") + if exc_info is None: + return False + + error = exc_info[0] + error_type_name = get_type_name(exc_info[0]) + error_full_name = "%s.%s" % (exc_info[0].__module__, error_type_name) + + for ignored_error in self.options["ignore_errors"]: + # String types are matched against the type name in the + # exception only + if isinstance(ignored_error, str): + if ignored_error == error_full_name or ignored_error == error_type_name: + return True + else: + if issubclass(error, ignored_error): + return True + + return False + + def _should_capture( + self, + event: "Event", + hint: "Hint", + scope: "Optional[Scope]" = None, + ) -> bool: + # Transactions are sampled independent of error events. + is_transaction = event.get("type") == "transaction" + if is_transaction: + return True + + ignoring_prevents_recursion = scope is not None and not scope._should_capture + if ignoring_prevents_recursion: + return False + + ignored_by_config_option = self._is_ignored_error(event, hint) + if ignored_by_config_option: + return False + + return True + + def _should_sample_error( + self, + event: "Event", + hint: "Hint", + ) -> bool: + error_sampler = self.options.get("error_sampler", None) + + if callable(error_sampler): + with capture_internal_exceptions(): + sample_rate = error_sampler(event, hint) + else: + sample_rate = self.options["sample_rate"] + + try: + not_in_sample_rate = sample_rate < 1.0 and random.random() >= sample_rate + except NameError: + logger.warning( + "The provided error_sampler raised an error. Defaulting to sampling the event." + ) + + # If the error_sampler raised an error, we should sample the event, since the default behavior + # (when no sample_rate or error_sampler is provided) is to sample all events. + not_in_sample_rate = False + except TypeError: + parameter, verb = ( + ("error_sampler", "returned") + if callable(error_sampler) + else ("sample_rate", "contains") + ) + logger.warning( + "The provided %s %s an invalid value of %s. The value should be a float or a bool. Defaulting to sampling the event." + % (parameter, verb, repr(sample_rate)) + ) + + # If the sample_rate has an invalid value, we should sample the event, since the default behavior + # (when no sample_rate or error_sampler is provided) is to sample all events. + not_in_sample_rate = False + + if not_in_sample_rate: + # because we will not sample this event, record a "lost event". + if self.transport: + self.transport.record_lost_event("sample_rate", data_category="error") + + return False + + return True + + def _update_session_from_event( + self, + session: "Session", + event: "Event", + ) -> None: + crashed = False + errored = False + user_agent = None + + exceptions = (event.get("exception") or {}).get("values") + if exceptions: + errored = True + for error in exceptions: + if isinstance(error, AnnotatedValue): + error = error.value or {} + mechanism = error.get("mechanism") + if isinstance(mechanism, Mapping) and mechanism.get("handled") is False: + crashed = True + break + + user = event.get("user") + + if session.user_agent is None: + headers = (event.get("request") or {}).get("headers") + headers_dict = headers if isinstance(headers, dict) else {} + for k, v in headers_dict.items(): + if k.lower() == "user-agent": + user_agent = v + break + + session.update( + status="crashed" if crashed else None, + user=user, + user_agent=user_agent, + errors=session.errors + (errored or crashed), + ) + + def capture_event( + self, + event: "Event", + hint: "Optional[Hint]" = None, + scope: "Optional[Scope]" = None, + ) -> "Optional[str]": + """Captures an event. + + :param event: A ready-made event that can be directly sent to Sentry. + + :param hint: Contains metadata about the event that can be read from `before_send`, such as the original exception object or a HTTP request object. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + + :returns: An event ID. May be `None` if there is no DSN set or of if the SDK decided to discard the event for other reasons. In such situations setting `debug=True` on `init()` may help. + """ + hint: "Hint" = dict(hint or ()) + + if not self._should_capture(event, hint, scope): + return None + + profile = event.pop("profile", None) + + event_id = event.get("event_id") + if event_id is None: + event["event_id"] = event_id = uuid.uuid4().hex + event_opt = self._prepare_event(event, hint, scope) + if event_opt is None: + return None + + # whenever we capture an event we also check if the session needs + # to be updated based on that information. + session = scope._session if scope else None + if session: + self._update_session_from_event(session, event) + + is_transaction = event_opt.get("type") == "transaction" + is_checkin = event_opt.get("type") == "check_in" + + if ( + not is_transaction + and not is_checkin + and not self._should_sample_error(event, hint) + ): + return None + + attachments = hint.get("attachments") + + trace_context = event_opt.get("contexts", {}).get("trace") or {} + dynamic_sampling_context = trace_context.pop("dynamic_sampling_context", {}) + + headers: "dict[str, object]" = { + "event_id": event_opt["event_id"], + "sent_at": format_timestamp(datetime.now(timezone.utc)), + } + + if dynamic_sampling_context: + headers["trace"] = dynamic_sampling_context + + envelope = Envelope(headers=headers) + + if is_transaction: + if isinstance(profile, Profile): + envelope.add_profile(profile.to_json(event_opt, self.options)) + envelope.add_transaction(event_opt) + elif is_checkin: + envelope.add_checkin(event_opt) + else: + envelope.add_event(event_opt) + + for attachment in attachments or (): + envelope.add_item(attachment.to_envelope_item()) + + return_value = None + if self.spotlight: + self.spotlight.capture_envelope(envelope) + return_value = event_id + + if self.transport is not None: + self.transport.capture_envelope(envelope) + return_value = event_id + + return return_value + + def _capture_telemetry( + self, + telemetry: "Optional[Union[Log, Metric, StreamedSpan]]", + ty: str, + scope: "Scope", + ) -> None: + # Capture attributes-based telemetry (logs, metrics, spansV2) + if telemetry is None: + return + + scope.apply_to_telemetry(telemetry) + + before_send = None + if ty == "log": + before_send = get_before_send_log(self.options) + elif ty == "metric": + before_send = get_before_send_metric(self.options) # type: ignore + + if before_send is not None: + telemetry = before_send(telemetry, {}) # type: ignore + + if telemetry is None: + return + + batcher = None + if ty == "log": + batcher = self.log_batcher + elif ty == "metric": + batcher = self.metrics_batcher # type: ignore + elif ty == "span": + batcher = self.span_batcher # type: ignore + + if batcher is not None: + batcher.add(telemetry) # type: ignore + + def _capture_log(self, log: "Optional[Log]", scope: "Scope") -> None: + self._capture_telemetry(log, "log", scope) + + def _capture_metric(self, metric: "Optional[Metric]", scope: "Scope") -> None: + self._capture_telemetry(metric, "metric", scope) + + def _capture_span(self, span: "Optional[StreamedSpan]", scope: "Scope") -> None: + self._capture_telemetry(span, "span", scope) + + def capture_session( + self, + session: "Session", + ) -> None: + if not session.release: + logger.info("Discarded session update because of missing release") + else: + self.session_flusher.add_session(session) + + if TYPE_CHECKING: + + @overload + def get_integration(self, name_or_class: str) -> "Optional[Integration]": ... + + @overload + def get_integration(self, name_or_class: "type[I]") -> "Optional[I]": ... + + def get_integration( + self, + name_or_class: "Union[str, Type[Integration]]", + ) -> "Optional[Integration]": + """Returns the integration for this client by name or class. + If the client does not have that integration then `None` is returned. + """ + if isinstance(name_or_class, str): + integration_name = name_or_class + elif name_or_class.identifier is not None: + integration_name = name_or_class.identifier + else: + raise ValueError("Integration has no name") + + return self.integrations.get(integration_name) + + def close( + self, + timeout: "Optional[float]" = None, + callback: "Optional[Callable[[int, float], None]]" = None, + ) -> None: + """ + Close the client and shut down the transport. Arguments have the same + semantics as :py:meth:`Client.flush`. + """ + if self.transport is not None: + self.flush(timeout=timeout, callback=callback) + self.session_flusher.kill() + if self.log_batcher is not None: + self.log_batcher.kill() + if self.metrics_batcher is not None: + self.metrics_batcher.kill() + if self.span_batcher is not None: + self.span_batcher.kill() + if self.monitor: + self.monitor.kill() + self.transport.kill() + self.transport = None + + def flush( + self, + timeout: "Optional[float]" = None, + callback: "Optional[Callable[[int, float], None]]" = None, + ) -> None: + """ + Wait for the current events to be sent. + + :param timeout: Wait for at most `timeout` seconds. If no `timeout` is provided, the `shutdown_timeout` option value is used. + + :param callback: Is invoked with the number of pending events and the configured timeout. + """ + if self.transport is not None: + if timeout is None: + timeout = self.options["shutdown_timeout"] + self.session_flusher.flush() + if self.log_batcher is not None: + self.log_batcher.flush() + if self.metrics_batcher is not None: + self.metrics_batcher.flush() + if self.span_batcher is not None: + self.span_batcher.flush() + self.transport.flush(timeout=timeout, callback=callback) + + def __enter__(self) -> "_Client": + return self + + def __exit__(self, exc_type: "Any", exc_value: "Any", tb: "Any") -> None: + self.close() + + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + # Make mypy, PyCharm and other static analyzers think `get_options` is a + # type to have nicer autocompletion for params. + # + # Use `ClientConstructor` to define the argument types of `init` and + # `Dict[str, Any]` to tell static analyzers about the return type. + + class get_options(ClientConstructor, Dict[str, Any]): # noqa: N801 + pass + + class Client(ClientConstructor, _Client): + pass + +else: + # Alias `get_options` for actual usage. Go through the lambda indirection + # to throw PyCharm off of the weakly typed signature (it would otherwise + # discover both the weakly typed signature of `_init` and our faked `init` + # type). + + get_options = (lambda: _get_options)() + Client = (lambda: _Client)() diff --git a/lib/python3.12/site-packages/sentry_sdk/consts.py b/lib/python3.12/site-packages/sentry_sdk/consts.py new file mode 100644 index 0000000000000000000000000000000000000000..6eaa1edccf5f3fbb683f35d46fe10501a6cd1fe0 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/consts.py @@ -0,0 +1,1493 @@ +import itertools +from enum import Enum +from typing import TYPE_CHECKING + +# up top to prevent circular import due to integration import +# This is more or less an arbitrary large-ish value for now, so that we allow +# pretty long strings (like LLM prompts), but still have *some* upper limit +# until we verify that removing the trimming completely is safe. +DEFAULT_MAX_VALUE_LENGTH = 100_000 + +DEFAULT_MAX_STACK_FRAMES = 100 +DEFAULT_ADD_FULL_STACK = False + + +# Also needs to be at the top to prevent circular import +class EndpointType(Enum): + """ + The type of an endpoint. This is an enum, rather than a constant, for historical reasons + (the old /store endpoint). The enum also preserve future compatibility, in case we ever + have a new endpoint. + """ + + ENVELOPE = "envelope" + OTLP_TRACES = "integration/otlp/v1/traces" + + +class CompressionAlgo(Enum): + GZIP = "gzip" + BROTLI = "br" + + +if TYPE_CHECKING: + from typing import ( + AbstractSet, + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + ) + + from typing_extensions import Literal, TypedDict + + import sentry_sdk + from sentry_sdk._types import ( + BreadcrumbProcessor, + ContinuousProfilerMode, + Event, + EventProcessor, + Hint, + Log, + MeasurementUnit, + Metric, + ProfilerMode, + TracesSampler, + TransactionProcessor, + ) + + # Experiments are feature flags to enable and disable certain unstable SDK + # functionality. Changing them from the defaults (`None`) in production + # code is highly discouraged. They are not subject to any stability + # guarantees such as the ones from semantic versioning. + Experiments = TypedDict( + "Experiments", + { + "max_spans": Optional[int], + "max_flags": Optional[int], + "record_sql_params": Optional[bool], + "continuous_profiling_auto_start": Optional[bool], + "continuous_profiling_mode": Optional[ContinuousProfilerMode], + "otel_powered_performance": Optional[bool], + "transport_zlib_compression_level": Optional[int], + "transport_compression_level": Optional[int], + "transport_compression_algo": Optional[CompressionAlgo], + "transport_num_pools": Optional[int], + "transport_http2": Optional[bool], + "enable_logs": Optional[bool], + "before_send_log": Optional[Callable[[Log, Hint], Optional[Log]]], + "enable_metrics": Optional[bool], + "before_send_metric": Optional[Callable[[Metric, Hint], Optional[Metric]]], + "trace_lifecycle": Optional[Literal["static", "stream"]], + }, + total=False, + ) + +DEFAULT_QUEUE_SIZE = 100 +DEFAULT_MAX_BREADCRUMBS = 100 +MATCH_ALL = r".*" + +FALSE_VALUES = [ + "false", + "no", + "off", + "n", + "0", +] + + +class SPANTEMPLATE(str, Enum): + DEFAULT = "default" + AI_AGENT = "ai_agent" + AI_TOOL = "ai_tool" + AI_CHAT = "ai_chat" + + def __str__(self) -> str: + return self.value + + +class INSTRUMENTER: + SENTRY = "sentry" + OTEL = "otel" + + +class SPANNAME: + DB_COMMIT = "COMMIT" + DB_ROLLBACK = "ROLLBACK" + + +class SPANDATA: + """ + Additional information describing the type of the span. + See: https://develop.sentry.dev/sdk/performance/span-data-conventions/ + """ + + AI_CITATIONS = "ai.citations" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + References or sources cited by the AI model in its response. + Example: ["Smith et al. 2020", "Jones 2019"] + """ + + AI_DOCUMENTS = "ai.documents" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Documents or content chunks used as context for the AI model. + Example: ["doc1.txt", "doc2.pdf"] + """ + + AI_FINISH_REASON = "ai.finish_reason" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_FINISH_REASONS instead. + + The reason why the model stopped generating. + Example: "length" + """ + + AI_FREQUENCY_PENALTY = "ai.frequency_penalty" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_FREQUENCY_PENALTY instead. + + Used to reduce repetitiveness of generated tokens. + Example: 0.5 + """ + + AI_FUNCTION_CALL = "ai.function_call" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_TOOL_CALLS instead. + + For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls + """ + + AI_GENERATION_ID = "ai.generation_id" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_ID instead. + + Unique identifier for the completion. + Example: "gen_123abc" + """ + + AI_INPUT_MESSAGES = "ai.input_messages" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_MESSAGES instead. + + The input messages to an LLM call. + Example: [{"role": "user", "message": "hello"}] + """ + + AI_LOGIT_BIAS = "ai.logit_bias" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + For an AI model call, the logit bias + """ + + AI_METADATA = "ai.metadata" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Extra metadata passed to an AI pipeline step. + Example: {"executed_function": "add_integers"} + """ + + AI_MODEL_ID = "ai.model_id" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_MODEL or GEN_AI_RESPONSE_MODEL instead. + + The unique descriptor of the model being executed. + Example: gpt-4 + """ + + AI_PIPELINE_NAME = "ai.pipeline.name" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_PIPELINE_NAME instead. + + Name of the AI pipeline or chain being executed. + Example: "qa-pipeline" + """ + + AI_PREAMBLE = "ai.preamble" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + For an AI model call, the preamble parameter. + Preambles are a part of the prompt used to adjust the model's overall behavior and conversation style. + Example: "You are now a clown." + """ + + AI_PRESENCE_PENALTY = "ai.presence_penalty" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_PRESENCE_PENALTY instead. + + Used to reduce repetitiveness of generated tokens. + Example: 0.5 + """ + + AI_RAW_PROMPTING = "ai.raw_prompting" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Minimize pre-processing done to the prompt sent to the LLM. + Example: true + """ + + AI_RESPONSE_FORMAT = "ai.response_format" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + For an AI model call, the format of the response + """ + + AI_RESPONSES = "ai.responses" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_TEXT instead. + + The responses to an AI model call. Always as a list. + Example: ["hello", "world"] + """ + + AI_SEARCH_QUERIES = "ai.search_queries" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Queries used to search for relevant context or documents. + Example: ["climate change effects", "renewable energy"] + """ + + AI_SEARCH_REQUIRED = "ai.is_search_required" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Boolean indicating if the model needs to perform a search. + Example: true + """ + + AI_SEARCH_RESULTS = "ai.search_results" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Results returned from search queries for context. + Example: ["Result 1", "Result 2"] + """ + + AI_SEED = "ai.seed" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_SEED instead. + + The seed, ideally models given the same seed and same other parameters will produce the exact same output. + Example: 123.45 + """ + + AI_STREAMING = "ai.streaming" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_STREAMING instead. + + Whether or not the AI model call's response was streamed back asynchronously + Example: true + """ + + AI_TAGS = "ai.tags" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Tags that describe an AI pipeline step. + Example: {"executed_function": "add_integers"} + """ + + AI_TEMPERATURE = "ai.temperature" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_TEMPERATURE instead. + + For an AI model call, the temperature parameter. Temperature essentially means how random the output will be. + Example: 0.5 + """ + + AI_TEXTS = "ai.texts" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Raw text inputs provided to the model. + Example: ["What is machine learning?"] + """ + + AI_TOP_K = "ai.top_k" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_TOP_K instead. + + For an AI model call, the top_k parameter. Top_k essentially controls how random the output will be. + Example: 35 + """ + + AI_TOP_P = "ai.top_p" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_TOP_P instead. + + For an AI model call, the top_p parameter. Top_p essentially controls how random the output will be. + Example: 0.5 + """ + + AI_TOOL_CALLS = "ai.tool_calls" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_RESPONSE_TOOL_CALLS instead. + + For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls + """ + + AI_TOOLS = "ai.tools" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_REQUEST_AVAILABLE_TOOLS instead. + + For an AI model call, the functions that are available + """ + + AI_WARNINGS = "ai.warnings" + """ + .. deprecated:: + This attribute is deprecated. Use GEN_AI_* attributes instead. + + Warning messages generated during model execution. + Example: ["Token limit exceeded"] + """ + + CACHE_HIT = "cache.hit" + """ + A boolean indicating whether the requested data was found in the cache. + Example: true + """ + + CACHE_ITEM_SIZE = "cache.item_size" + """ + The size of the requested data in bytes. + Example: 58 + """ + + CACHE_KEY = "cache.key" + """ + The key of the requested data. + Example: template.cache.some_item.867da7e2af8e6b2f3aa7213a4080edb3 + """ + + CODE_FILEPATH = "code.filepath" + """ + The source code file name that identifies the code unit as uniquely as possible (preferably an absolute file path). + Example: "/app/myapplication/http/handler/server.py" + """ + + CODE_FUNCTION = "code.function" + """ + The method or function name, or equivalent (usually rightmost part of the code unit's name). + Example: "server_request" + """ + + CODE_LINENO = "code.lineno" + """ + The line number in `code.filepath` best representing the operation. It SHOULD point within the code unit named in `code.function`. + Example: 42 + """ + + CODE_NAMESPACE = "code.namespace" + """ + The "namespace" within which `code.function` is defined. Usually the qualified class or module name, such that `code.namespace` + some separator + `code.function` form a unique identifier for the code unit. + Example: "http.handler" + """ + + DB_MONGODB_COLLECTION = "db.mongodb.collection" + """ + The MongoDB collection being accessed within the database. + See: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/database/mongodb.md#attributes + Example: public.users; customers + """ + + DB_NAME = "db.name" + """ + The name of the database being accessed. For commands that switch the database, this should be set to the target database (even if the command fails). + Example: myDatabase + """ + + DB_OPERATION = "db.operation" + """ + The name of the operation being executed, e.g. the MongoDB command name such as findAndModify, or the SQL keyword. + See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md + Example: findAndModify, HMSET, SELECT + """ + + DB_SYSTEM = "db.system" + """ + An identifier for the database management system (DBMS) product being used. + See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md + Example: postgresql + """ + + DB_USER = "db.user" + """ + The name of the database user used for connecting to the database. + See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md + Example: my_user + """ + + GEN_AI_AGENT_NAME = "gen_ai.agent.name" + """ + The name of the agent being used. + Example: "ResearchAssistant" + """ + + GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" + """ + The unique identifier for the conversation/thread with the AI model. + Example: "conv_abc123" + """ + + GEN_AI_CHOICE = "gen_ai.choice" + """ + The model's response message. + Example: "The weather in Paris is rainy and overcast, with temperatures around 57°F" + """ + + GEN_AI_EMBEDDINGS_INPUT = "gen_ai.embeddings.input" + """ + The input to the embeddings operation. + Example: "Hello!" + """ + + GEN_AI_OPERATION_NAME = "gen_ai.operation.name" + """ + The name of the operation being performed. + Example: "chat" + """ + + GEN_AI_PIPELINE_NAME = "gen_ai.pipeline.name" + """ + Name of the AI pipeline or chain being executed. + Example: "qa-pipeline" + """ + + GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons" + """ + The reason why the model stopped generating. + Example: "COMPLETE" + """ + + GEN_AI_RESPONSE_ID = "gen_ai.response.id" + """ + Unique identifier for the completion. + Example: "gen_123abc" + """ + + GEN_AI_RESPONSE_MODEL = "gen_ai.response.model" + """ + Exact model identifier used to generate the response + Example: gpt-4o-mini-2024-07-18 + """ + + GEN_AI_RESPONSE_STREAMING = "gen_ai.response.streaming" + """ + Whether or not the AI model call's response was streamed back asynchronously + Example: true + """ + + GEN_AI_RESPONSE_TEXT = "gen_ai.response.text" + """ + The model's response text messages. + Example: ["The weather in Paris is rainy and overcast, with temperatures around 57°F", "The weather in London is sunny and warm, with temperatures around 65°F"] + """ + + GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN = "gen_ai.response.time_to_first_token" + """ + The time it took to receive the first token from the model. + Example: 0.1 + """ + + GEN_AI_RESPONSE_TOOL_CALLS = "gen_ai.response.tool_calls" + """ + The tool calls in the model's response. + Example: [{"name": "get_weather", "arguments": {"location": "Paris"}}] + """ + + GEN_AI_REQUEST_AVAILABLE_TOOLS = "gen_ai.request.available_tools" + """ + The available tools for the model. + Example: [{"name": "get_weather", "description": "Get the weather for a given location"}, {"name": "get_news", "description": "Get the news for a given topic"}] + """ + + GEN_AI_REQUEST_FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty" + """ + The frequency penalty parameter used to reduce repetitiveness of generated tokens. + Example: 0.1 + """ + + GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens" + """ + The maximum number of tokens to generate in the response. + Example: 2048 + """ + + GEN_AI_SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions" + """ + The system instructions passed to the model. + Example: [{"type": "text", "text": "You are a helpful assistant."},{"type": "text", "text": "Be concise and clear."}] + """ + + GEN_AI_REQUEST_MESSAGES = "gen_ai.request.messages" + """ + The messages passed to the model. The "content" can be a string or an array of objects. + Example: [{role: "system", "content: "Generate a random number."}, {"role": "user", "content": [{"text": "Generate a random number between 0 and 10.", "type": "text"}]}] + """ + + GEN_AI_REQUEST_MODEL = "gen_ai.request.model" + """ + The model identifier being used for the request. + Example: "gpt-4-turbo" + """ + + GEN_AI_REQUEST_PRESENCE_PENALTY = "gen_ai.request.presence_penalty" + """ + The presence penalty parameter used to reduce repetitiveness of generated tokens. + Example: 0.1 + """ + + GEN_AI_REQUEST_SEED = "gen_ai.request.seed" + """ + The seed, ideally models given the same seed and same other parameters will produce the exact same output. + Example: "1234567890" + """ + + GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature" + """ + The temperature parameter used to control randomness in the output. + Example: 0.7 + """ + + GEN_AI_REQUEST_TOP_K = "gen_ai.request.top_k" + """ + Limits the model to only consider the K most likely next tokens, where K is an integer (e.g., top_k=20 means only the 20 highest probability tokens are considered). + Example: 35 + """ + + GEN_AI_REQUEST_TOP_P = "gen_ai.request.top_p" + """ + The top_p parameter used to control diversity via nucleus sampling. + Example: 1.0 + """ + + GEN_AI_SYSTEM = "gen_ai.system" + """ + The name of the AI system being used. + Example: "openai" + """ + + GEN_AI_TOOL_DESCRIPTION = "gen_ai.tool.description" + """ + The description of the tool being used. + Example: "Searches the web for current information about a topic" + """ + + GEN_AI_TOOL_INPUT = "gen_ai.tool.input" + """ + The input of the tool being used. + Example: {"location": "Paris"} + """ + + GEN_AI_TOOL_NAME = "gen_ai.tool.name" + """ + The name of the tool being used. + Example: "web_search" + """ + + GEN_AI_TOOL_OUTPUT = "gen_ai.tool.output" + """ + The output of the tool being used. + Example: "rainy, 57°F" + """ + + GEN_AI_TOOL_TYPE = "gen_ai.tool.type" + """ + The type of tool being used. + Example: "function" + """ + + GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens" + """ + The number of tokens in the input. + Example: 150 + """ + + GEN_AI_USAGE_INPUT_TOKENS_CACHED = "gen_ai.usage.input_tokens.cached" + """ + The number of cached tokens in the input. + Example: 50 + """ + + GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE = "gen_ai.usage.input_tokens.cache_write" + """ + The number of tokens written to the cache when processing the AI input (prompt). + Example: 100 + """ + + GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens" + """ + The number of tokens in the output. + Example: 250 + """ + + GEN_AI_USAGE_OUTPUT_TOKENS_REASONING = "gen_ai.usage.output_tokens.reasoning" + """ + The number of tokens used for reasoning in the output. + Example: 75 + """ + + GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens" + """ + The total number of tokens used (input + output). + Example: 400 + """ + + GEN_AI_USER_MESSAGE = "gen_ai.user.message" + """ + The user message passed to the model. + Example: "What's the weather in Paris?" + """ + + HTTP_FRAGMENT = "http.fragment" + """ + The Fragments present in the URL. + Example: #foo=bar + """ + + HTTP_METHOD = "http.method" + """ + The HTTP method used. + Example: GET + """ + + HTTP_QUERY = "http.query" + """ + The Query string present in the URL. + Example: ?foo=bar&bar=baz + """ + + HTTP_STATUS_CODE = "http.response.status_code" + """ + The HTTP status code as an integer. + Example: 418 + """ + + MESSAGING_DESTINATION_NAME = "messaging.destination.name" + """ + The destination name where the message is being consumed from, + e.g. the queue name or topic. + """ + + MESSAGING_MESSAGE_ID = "messaging.message.id" + """ + The message's identifier. + """ + + MESSAGING_MESSAGE_RECEIVE_LATENCY = "messaging.message.receive.latency" + """ + The latency between when the task was enqueued and when it was started to be processed. + """ + + MESSAGING_MESSAGE_RETRY_COUNT = "messaging.message.retry.count" + """ + Number of retries/attempts to process a message. + """ + + MESSAGING_SYSTEM = "messaging.system" + """ + The messaging system's name, e.g. `kafka`, `aws_sqs` + """ + + NETWORK_PEER_ADDRESS = "network.peer.address" + """ + Peer address of the network connection - IP address or Unix domain socket name. + Example: 10.1.2.80, /tmp/my.sock, localhost + """ + + NETWORK_PEER_PORT = "network.peer.port" + """ + Peer port number of the network connection. + Example: 6379 + """ + + NETWORK_TRANSPORT = "network.transport" + """ + The transport protocol used for the network connection. + Example: "tcp", "udp", "unix" + """ + + PROFILER_ID = "profiler_id" + """ + Label identifying the profiler id that the span occurred in. This should be a string. + Example: "5249fbada8d5416482c2f6e47e337372" + """ + + SERVER_ADDRESS = "server.address" + """ + Name of the database host. + Example: example.com + """ + + SERVER_PORT = "server.port" + """ + Logical server port number + Example: 80; 8080; 443 + """ + + SERVER_SOCKET_ADDRESS = "server.socket.address" + """ + Physical server IP address or Unix socket address. + Example: 10.5.3.2 + """ + + SERVER_SOCKET_PORT = "server.socket.port" + """ + Physical server port. + Recommended: If different than server.port. + Example: 16456 + """ + + THREAD_ID = "thread.id" + """ + Identifier of a thread from where the span originated. This should be a string. + Example: "7972576320" + """ + + THREAD_NAME = "thread.name" + """ + Label identifying a thread from where the span originated. This should be a string. + Example: "MainThread" + """ + + MCP_TOOL_NAME = "mcp.tool.name" + """ + The name of the MCP tool being called. + Example: "get_weather" + """ + + MCP_PROMPT_NAME = "mcp.prompt.name" + """ + The name of the MCP prompt being retrieved. + Example: "code_review" + """ + + MCP_RESOURCE_URI = "mcp.resource.uri" + """ + The URI of the MCP resource being accessed. + Example: "file:///path/to/resource" + """ + + MCP_METHOD_NAME = "mcp.method.name" + """ + The MCP protocol method name being called. + Example: "tools/call", "prompts/get", "resources/read" + """ + + MCP_REQUEST_ID = "mcp.request.id" + """ + The unique identifier for the MCP request. + Example: "req_123abc" + """ + + MCP_TOOL_RESULT_CONTENT = "mcp.tool.result.content" + """ + The result/output content from an MCP tool execution. + Example: "The weather is sunny" + """ + + MCP_TOOL_RESULT_CONTENT_COUNT = "mcp.tool.result.content_count" + """ + The number of items/keys in the MCP tool result. + Example: 5 + """ + + MCP_TOOL_RESULT_IS_ERROR = "mcp.tool.result.is_error" + """ + Whether the MCP tool execution resulted in an error. + Example: True + """ + + MCP_PROMPT_RESULT_MESSAGE_CONTENT = "mcp.prompt.result.message_content" + """ + The message content from an MCP prompt retrieval. + Example: "Review the following code..." + """ + + MCP_PROMPT_RESULT_MESSAGE_ROLE = "mcp.prompt.result.message_role" + """ + The role of the message in an MCP prompt retrieval (only set for single-message prompts). + Example: "user", "assistant", "system" + """ + + MCP_PROMPT_RESULT_MESSAGE_COUNT = "mcp.prompt.result.message_count" + """ + The number of messages in an MCP prompt result. + Example: 1, 3 + """ + + MCP_RESOURCE_PROTOCOL = "mcp.resource.protocol" + """ + The protocol/scheme of the MCP resource URI. + Example: "file", "http", "https" + """ + + MCP_TRANSPORT = "mcp.transport" + """ + The transport method used for MCP communication. + Example: "http", "sse", "stdio" + """ + + MCP_SESSION_ID = "mcp.session.id" + """ + The session identifier for the MCP connection. + Example: "a1b2c3d4e5f6" + """ + + +class SPANSTATUS: + """ + The status of a Sentry span. + + See: https://develop.sentry.dev/sdk/event-payloads/contexts/#trace-context + """ + + ABORTED = "aborted" + ALREADY_EXISTS = "already_exists" + CANCELLED = "cancelled" + DATA_LOSS = "data_loss" + DEADLINE_EXCEEDED = "deadline_exceeded" + FAILED_PRECONDITION = "failed_precondition" + INTERNAL_ERROR = "internal_error" + INVALID_ARGUMENT = "invalid_argument" + NOT_FOUND = "not_found" + OK = "ok" + OUT_OF_RANGE = "out_of_range" + PERMISSION_DENIED = "permission_denied" + RESOURCE_EXHAUSTED = "resource_exhausted" + UNAUTHENTICATED = "unauthenticated" + UNAVAILABLE = "unavailable" + UNIMPLEMENTED = "unimplemented" + UNKNOWN_ERROR = "unknown_error" + + +class OP: + ANTHROPIC_MESSAGES_CREATE = "ai.messages.create.anthropic" + CACHE_GET = "cache.get" + CACHE_PUT = "cache.put" + COHERE_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.cohere" + COHERE_EMBEDDINGS_CREATE = "ai.embeddings.create.cohere" + DB = "db" + DB_REDIS = "db.redis" + EVENT_DJANGO = "event.django" + FUNCTION = "function" + FUNCTION_AWS = "function.aws" + FUNCTION_GCP = "function.gcp" + GEN_AI_CHAT = "gen_ai.chat" + GEN_AI_CREATE_AGENT = "gen_ai.create_agent" + GEN_AI_EMBEDDINGS = "gen_ai.embeddings" + GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool" + GEN_AI_GENERATE_TEXT = "gen_ai.generate_text" + GEN_AI_HANDOFF = "gen_ai.handoff" + GEN_AI_PIPELINE = "gen_ai.pipeline" + GEN_AI_INVOKE_AGENT = "gen_ai.invoke_agent" + GEN_AI_RESPONSES = "gen_ai.responses" + GRAPHQL_EXECUTE = "graphql.execute" + GRAPHQL_MUTATION = "graphql.mutation" + GRAPHQL_PARSE = "graphql.parse" + GRAPHQL_RESOLVE = "graphql.resolve" + GRAPHQL_SUBSCRIPTION = "graphql.subscription" + GRAPHQL_QUERY = "graphql.query" + GRAPHQL_VALIDATE = "graphql.validate" + GRPC_CLIENT = "grpc.client" + GRPC_SERVER = "grpc.server" + HTTP_CLIENT = "http.client" + HTTP_CLIENT_STREAM = "http.client.stream" + HTTP_SERVER = "http.server" + MIDDLEWARE_DJANGO = "middleware.django" + MIDDLEWARE_LITESTAR = "middleware.litestar" + MIDDLEWARE_LITESTAR_RECEIVE = "middleware.litestar.receive" + MIDDLEWARE_LITESTAR_SEND = "middleware.litestar.send" + MIDDLEWARE_STARLETTE = "middleware.starlette" + MIDDLEWARE_STARLETTE_RECEIVE = "middleware.starlette.receive" + MIDDLEWARE_STARLETTE_SEND = "middleware.starlette.send" + MIDDLEWARE_STARLITE = "middleware.starlite" + MIDDLEWARE_STARLITE_RECEIVE = "middleware.starlite.receive" + MIDDLEWARE_STARLITE_SEND = "middleware.starlite.send" + HUGGINGFACE_HUB_CHAT_COMPLETIONS_CREATE = ( + "ai.chat_completions.create.huggingface_hub" + ) + QUEUE_PROCESS = "queue.process" + QUEUE_PUBLISH = "queue.publish" + QUEUE_SUBMIT_ARQ = "queue.submit.arq" + QUEUE_TASK_ARQ = "queue.task.arq" + QUEUE_SUBMIT_CELERY = "queue.submit.celery" + QUEUE_TASK_CELERY = "queue.task.celery" + QUEUE_TASK_RQ = "queue.task.rq" + QUEUE_SUBMIT_HUEY = "queue.submit.huey" + QUEUE_TASK_HUEY = "queue.task.huey" + QUEUE_SUBMIT_RAY = "queue.submit.ray" + QUEUE_TASK_RAY = "queue.task.ray" + QUEUE_TASK_DRAMATIQ = "queue.task.dramatiq" + QUEUE_SUBMIT_DJANGO = "queue.submit.django" + SUBPROCESS = "subprocess" + SUBPROCESS_WAIT = "subprocess.wait" + SUBPROCESS_COMMUNICATE = "subprocess.communicate" + TEMPLATE_RENDER = "template.render" + VIEW_RENDER = "view.render" + VIEW_RESPONSE_RENDER = "view.response.render" + WEBSOCKET_SERVER = "websocket.server" + SOCKET_CONNECTION = "socket.connection" + SOCKET_DNS = "socket.dns" + MCP_SERVER = "mcp.server" + + +# This type exists to trick mypy and PyCharm into thinking `init` and `Client` +# take these arguments (even though they take opaque **kwargs) +class ClientConstructor: + def __init__( + self, + dsn: "Optional[str]" = None, + *, + max_breadcrumbs: int = DEFAULT_MAX_BREADCRUMBS, + release: "Optional[str]" = None, + environment: "Optional[str]" = None, + server_name: "Optional[str]" = None, + shutdown_timeout: float = 2, + integrations: "Sequence[sentry_sdk.integrations.Integration]" = [], # noqa: B006 + in_app_include: "List[str]" = [], # noqa: B006 + in_app_exclude: "List[str]" = [], # noqa: B006 + default_integrations: bool = True, + dist: "Optional[str]" = None, + transport: "Optional[Union[sentry_sdk.transport.Transport, Type[sentry_sdk.transport.Transport], Callable[[Event], None]]]" = None, + transport_queue_size: int = DEFAULT_QUEUE_SIZE, + sample_rate: float = 1.0, + send_default_pii: "Optional[bool]" = None, + http_proxy: "Optional[str]" = None, + https_proxy: "Optional[str]" = None, + ignore_errors: "Sequence[Union[type, str]]" = [], # noqa: B006 + max_request_body_size: str = "medium", + socket_options: "Optional[List[Tuple[int, int, int | bytes]]]" = None, + keep_alive: "Optional[bool]" = None, + before_send: "Optional[EventProcessor]" = None, + before_breadcrumb: "Optional[BreadcrumbProcessor]" = None, + debug: "Optional[bool]" = None, + attach_stacktrace: bool = False, + ca_certs: "Optional[str]" = None, + propagate_traces: bool = True, + traces_sample_rate: "Optional[float]" = None, + traces_sampler: "Optional[TracesSampler]" = None, + profiles_sample_rate: "Optional[float]" = None, + profiles_sampler: "Optional[TracesSampler]" = None, + profiler_mode: "Optional[ProfilerMode]" = None, + profile_lifecycle: 'Literal["manual", "trace"]' = "manual", + profile_session_sample_rate: "Optional[float]" = None, + auto_enabling_integrations: bool = True, + disabled_integrations: "Optional[Sequence[sentry_sdk.integrations.Integration]]" = None, + auto_session_tracking: bool = True, + send_client_reports: bool = True, + _experiments: "Experiments" = {}, # noqa: B006 + proxy_headers: "Optional[Dict[str, str]]" = None, + instrumenter: "Optional[str]" = INSTRUMENTER.SENTRY, + before_send_transaction: "Optional[TransactionProcessor]" = None, + project_root: "Optional[str]" = None, + enable_tracing: "Optional[bool]" = None, + include_local_variables: "Optional[bool]" = True, + include_source_context: "Optional[bool]" = True, + trace_propagation_targets: "Optional[Sequence[str]]" = [ # noqa: B006 + MATCH_ALL + ], + functions_to_trace: "Sequence[Dict[str, str]]" = [], # noqa: B006 + event_scrubber: "Optional[sentry_sdk.scrubber.EventScrubber]" = None, + max_value_length: int = DEFAULT_MAX_VALUE_LENGTH, + enable_backpressure_handling: bool = True, + error_sampler: "Optional[Callable[[Event, Hint], Union[float, bool]]]" = None, + enable_db_query_source: bool = True, + db_query_source_threshold_ms: int = 100, + enable_http_request_source: bool = True, + http_request_source_threshold_ms: int = 100, + spotlight: "Optional[Union[bool, str]]" = None, + cert_file: "Optional[str]" = None, + key_file: "Optional[str]" = None, + custom_repr: "Optional[Callable[..., Optional[str]]]" = None, + add_full_stack: bool = DEFAULT_ADD_FULL_STACK, + max_stack_frames: "Optional[int]" = DEFAULT_MAX_STACK_FRAMES, + enable_logs: bool = False, + before_send_log: "Optional[Callable[[Log, Hint], Optional[Log]]]" = None, + trace_ignore_status_codes: "AbstractSet[int]" = frozenset(), + enable_metrics: bool = True, + before_send_metric: "Optional[Callable[[Metric, Hint], Optional[Metric]]]" = None, + org_id: "Optional[str]" = None, + strict_trace_continuation: bool = False, + ) -> None: + """Initialize the Sentry SDK with the given parameters. All parameters described here can be used in a call to `sentry_sdk.init()`. + + :param dsn: The DSN tells the SDK where to send the events. + + If this option is not set, the SDK will just not send any data. + + The `dsn` config option takes precedence over the environment variable. + + Learn more about `DSN utilization `_. + + :param debug: Turns debug mode on or off. + + When `True`, the SDK will attempt to print out debugging information. This can be useful if something goes + wrong with event sending. + + The default is always `False`. It's generally not recommended to turn it on in production because of the + increase in log output. + + The `debug` config option takes precedence over the environment variable. + + :param release: Sets the release. + + If not set, the SDK will try to automatically configure a release out of the box but it's a better idea to + manually set it to guarantee that the release is in sync with your deploy integrations. + + Release names are strings, but some formats are detected by Sentry and might be rendered differently. + + See `the releases documentation `_ to learn how the SDK tries to + automatically configure a release. + + The `release` config option takes precedence over the environment variable. + + Learn more about how to send release data so Sentry can tell you about regressions between releases and + identify the potential source in `the product documentation `_. + + :param environment: Sets the environment. This string is freeform and set to `production` by default. + + A release can be associated with more than one environment to separate them in the UI (think `staging` vs + `production` or similar). + + The `environment` config option takes precedence over the environment variable. + + :param dist: The distribution of the application. + + Distributions are used to disambiguate build or deployment variants of the same release of an application. + + The dist can be for example a build number. + + :param sample_rate: Configures the sample rate for error events, in the range of `0.0` to `1.0`. + + The default is `1.0`, which means that 100% of error events will be sent. If set to `0.1`, only 10% of + error events will be sent. + + Events are picked randomly. + + :param error_sampler: Dynamically configures the sample rate for error events on a per-event basis. + + This configuration option accepts a function, which takes two parameters (the `event` and the `hint`), and + which returns a boolean (indicating whether the event should be sent to Sentry) or a floating-point number + between `0.0` and `1.0`, inclusive. + + The number indicates the probability the event is sent to Sentry; the SDK will randomly decide whether to + send the event with the given probability. + + If this configuration option is specified, the `sample_rate` option is ignored. + + :param ignore_errors: A list of exception class names that shouldn't be sent to Sentry. + + Errors that are an instance of these exceptions or a subclass of them, will be filtered out before they're + sent to Sentry. + + By default, all errors are sent. + + :param max_breadcrumbs: This variable controls the total amount of breadcrumbs that should be captured. + + This defaults to `100`, but you can set this to any number. + + However, you should be aware that Sentry has a `maximum payload size `_ + and any events exceeding that payload size will be dropped. + + :param attach_stacktrace: When enabled, stack traces are automatically attached to all messages logged. + + Stack traces are always attached to exceptions; however, when this option is set, stack traces are also + sent with messages. + + This option means that stack traces appear next to all log messages. + + Grouping in Sentry is different for events with stack traces and without. As a result, you will get new + groups as you enable or disable this flag for certain events. + + :param send_default_pii: If this flag is enabled, `certain personally identifiable information (PII) + `_ is added by active integrations. + + If you enable this option, be sure to manually remove what you don't want to send using our features for + managing `Sensitive Data `_. + + :param event_scrubber: Scrubs the event payload for sensitive information such as cookies, sessions, and + passwords from a `denylist`. + + It can additionally be used to scrub from another `pii_denylist` if `send_default_pii` is disabled. + + See how to `configure the scrubber here `_. + + :param include_source_context: When enabled, source context will be included in events sent to Sentry. + + This source context includes the five lines of code above and below the line of code where an error + happened. + + :param include_local_variables: When enabled, the SDK will capture a snapshot of local variables to send with + the event to help with debugging. + + :param add_full_stack: When capturing errors, Sentry stack traces typically only include frames that start the + moment an error occurs. + + But if the `add_full_stack` option is enabled (set to `True`), all frames from the start of execution will + be included in the stack trace sent to Sentry. + + :param max_stack_frames: This option limits the number of stack frames that will be captured when + `add_full_stack` is enabled. + + :param server_name: This option can be used to supply a server name. + + When provided, the name of the server is sent along and persisted in the event. + + For many integrations, the server name actually corresponds to the device hostname, even in situations + where the machine is not actually a server. + + :param project_root: The full path to the root directory of your application. + + The `project_root` is used to mark frames in a stack trace either as being in your application or outside + of the application. + + :param in_app_include: A list of string prefixes of module names that belong to the app. + + This option takes precedence over `in_app_exclude`. + + Sentry differentiates stack frames that are directly related to your application ("in application") from + stack frames that come from other packages such as the standard library, frameworks, or other dependencies. + + The application package is automatically marked as `inApp`. + + The difference is visible in [sentry.io](https://sentry.io), where only the "in application" frames are + displayed by default. + + :param in_app_exclude: A list of string prefixes of module names that do not belong to the app, but rather to + third-party packages. + + Modules considered not part of the app will be hidden from stack traces by default. + + This option can be overridden using `in_app_include`. + + :param max_request_body_size: This parameter controls whether integrations should capture HTTP request bodies. + It can be set to one of the following values: + + - `never`: Request bodies are never sent. + - `small`: Only small request bodies will be captured. The cutoff for small depends on the SDK (typically + 4KB). + - `medium`: Medium and small requests will be captured (typically 10KB). + - `always`: The SDK will always capture the request body as long as Sentry can make sense of it. + + Please note that the Sentry server [limits HTTP request body size](https://develop.sentry.dev/sdk/ + expected-features/data-handling/#variable-size). The server always enforces its size limit, regardless of + how you configure this option. + + :param max_value_length: The number of characters after which the values containing text in the event payload + will be truncated. + + WARNING: If the value you set for this is exceptionally large, the event may exceed 1 MiB and will be + dropped by Sentry. + + :param ca_certs: A path to an alternative CA bundle file in PEM-format. + + :param send_client_reports: Set this boolean to `False` to disable sending of client reports. + + Client reports allow the client to send status reports about itself to Sentry, such as information about + events that were dropped before being sent. + + :param integrations: List of integrations to enable in addition to `auto-enabling integrations (overview) + `_. + + This setting can be used to override the default config options for a specific auto-enabling integration + or to add an integration that is not auto-enabled. + + :param disabled_integrations: List of integrations that will be disabled. + + This setting can be used to explicitly turn off specific `auto-enabling integrations (list) + `_ or + `default `_ integrations. + + :param auto_enabling_integrations: Configures whether `auto-enabling integrations (configuration) + `_ should be enabled. + + When set to `False`, no auto-enabling integrations will be enabled by default, even if the corresponding + framework/library is detected. + + :param default_integrations: Configures whether `default integrations + `_ should be enabled. + + Setting `default_integrations` to `False` disables all default integrations **as well as all auto-enabling + integrations**, unless they are specifically added in the `integrations` option, described above. + + :param before_send: This function is called with an SDK-specific message or error event object, and can return + a modified event object, or `null` to skip reporting the event. + + This can be used, for instance, for manual PII stripping before sending. + + By the time `before_send` is executed, all scope data has already been applied to the event. Further + modification of the scope won't have any effect. + + :param before_send_transaction: This function is called with an SDK-specific transaction event object, and can + return a modified transaction event object, or `null` to skip reporting the event. + + One way this might be used is for manual PII stripping before sending. + + :param before_breadcrumb: This function is called with an SDK-specific breadcrumb object before the breadcrumb + is added to the scope. + + When nothing is returned from the function, the breadcrumb is dropped. + + To pass the breadcrumb through, return the first argument, which contains the breadcrumb object. + + The callback typically gets a second argument (called a "hint") which contains the original object from + which the breadcrumb was created to further customize what the breadcrumb should look like. + + :param transport: Switches out the transport used to send events. + + How this works depends on the SDK. It can, for instance, be used to capture events for unit-testing or to + send it through some more complex setup that requires proxy authentication. + + :param transport_queue_size: The maximum number of events that will be queued before the transport is forced to + flush. + + :param http_proxy: When set, a proxy can be configured that should be used for outbound requests. + + This is also used for HTTPS requests unless a separate `https_proxy` is configured. However, not all SDKs + support a separate HTTPS proxy. + + SDKs will attempt to default to the system-wide configured proxy, if possible. For instance, on Unix + systems, the `http_proxy` environment variable will be picked up. + + :param https_proxy: Configures a separate proxy for outgoing HTTPS requests. + + This value might not be supported by all SDKs. When not supported the `http-proxy` value is also used for + HTTPS requests at all times. + + :param proxy_headers: A dict containing additional proxy headers (usually for authentication) to be forwarded + to `urllib3`'s `ProxyManager `_. + + :param shutdown_timeout: Controls how many seconds to wait before shutting down. + + Sentry SDKs send events from a background queue. This queue is given a certain amount to drain pending + events. The default is SDK specific but typically around two seconds. + + Setting this value too low may cause problems for sending events from command line applications. + + Setting the value too high will cause the application to block for a long time for users experiencing + network connectivity problems. + + :param keep_alive: Determines whether to keep the connection alive between requests. + + This can be useful in environments where you encounter frequent network issues such as connection resets. + + :param cert_file: Path to the client certificate to use. + + If set, supersedes the `CLIENT_CERT_FILE` environment variable. + + :param key_file: Path to the key file to use. + + If set, supersedes the `CLIENT_KEY_FILE` environment variable. + + :param socket_options: An optional list of socket options to use. + + These provide fine-grained, low-level control over the way the SDK connects to Sentry. + + If provided, the options will override the default `urllib3` `socket options + `_. + + :param traces_sample_rate: A number between `0` and `1`, controlling the percentage chance a given transaction + will be sent to Sentry. + + (`0` represents 0% while `1` represents 100%.) Applies equally to all transactions created in the app. + + Either this or `traces_sampler` must be defined to enable tracing. + + If `traces_sample_rate` is `0`, this means that no new traces will be created. However, if you have + another service (for example a JS frontend) that makes requests to your service that include trace + information, those traces will be continued and thus transactions will be sent to Sentry. + + If you want to disable all tracing you need to set `traces_sample_rate=None`. In this case, no new traces + will be started and no incoming traces will be continued. + + :param traces_sampler: A function responsible for determining the percentage chance a given transaction will be + sent to Sentry. + + It will automatically be passed information about the transaction and the context in which it's being + created, and must return a number between `0` (0% chance of being sent) and `1` (100% chance of being + sent). + + Can also be used for filtering transactions, by returning `0` for those that are unwanted. + + Either this or `traces_sample_rate` must be defined to enable tracing. + + :param trace_propagation_targets: An optional property that controls which downstream services receive tracing + data, in the form of a `sentry-trace` and a `baggage` header attached to any outgoing HTTP requests. + + The option may contain a list of strings or regex against which the URLs of outgoing requests are matched. + + If one of the entries in the list matches the URL of an outgoing request, trace data will be attached to + that request. + + String entries do not have to be full matches, meaning the URL of a request is matched when it _contains_ + a string provided through the option. + + If `trace_propagation_targets` is not provided, trace data is attached to every outgoing request from the + instrumented client. + + :param functions_to_trace: An optional list of functions that should be set up for tracing. + + For each function in the list, a span will be created when the function is executed. + + Functions in the list are represented as strings containing the fully qualified name of the function. + + This is a convenient option, making it possible to have one central place for configuring what functions + to trace, instead of having custom instrumentation scattered all over your code base. + + To learn more, see the `Custom Instrumentation `_ documentation. + + :param enable_backpressure_handling: When enabled, a new monitor thread will be spawned to perform health + checks on the SDK. + + If the system is unhealthy, the SDK will keep halving the `traces_sample_rate` set by you in 10 second + intervals until recovery. + + This down sampling helps ensure that the system stays stable and reduces SDK overhead under high load. + + This option is enabled by default. + + :param enable_db_query_source: When enabled, the source location will be added to database queries. + + :param db_query_source_threshold_ms: The threshold in milliseconds for adding the source location to database + queries. + + The query location will be added to the query for queries slower than the specified threshold. + + :param enable_http_request_source: When enabled, the source location will be added to outgoing HTTP requests. + + :param http_request_source_threshold_ms: The threshold in milliseconds for adding the source location to an + outgoing HTTP request. + + The request location will be added to the request for requests slower than the specified threshold. + + :param custom_repr: A custom `repr `_ function to run + while serializing an object. + + Use this to control how your custom objects and classes are visible in Sentry. + + Return a string for that repr value to be used or `None` to continue serializing how Sentry would have + done it anyway. + + :param profiles_sample_rate: A number between `0` and `1`, controlling the percentage chance a given sampled + transaction will be profiled. + + (`0` represents 0% while `1` represents 100%.) Applies equally to all transactions created in the app. + + This is relative to the tracing sample rate - e.g. `0.5` means 50% of sampled transactions will be + profiled. + + :param profiles_sampler: + + :param profiler_mode: + + :param profile_lifecycle: + + :param profile_session_sample_rate: + + :param enable_tracing: + + :param propagate_traces: + + :param auto_session_tracking: + + :param spotlight: + + :param instrumenter: + + :param enable_logs: Set `enable_logs` to True to enable the SDK to emit + Sentry logs. Defaults to False. + + :param before_send_log: An optional function to modify or filter out logs + before they're sent to Sentry. Any modifications to the log in this + function will be retained. If the function returns None, the log will + not be sent to Sentry. + + :param trace_ignore_status_codes: An optional property that disables tracing for + HTTP requests with certain status codes. + + Requests are not traced if the status code is contained in the provided set. + + If `trace_ignore_status_codes` is not provided, requests with any status code + may be traced. + + :param strict_trace_continuation: If set to `True`, the SDK will only continue a trace if the `org_id` of the incoming trace found in the + `baggage` header matches the `org_id` of the current Sentry client and only if BOTH are present. + + If set to `False`, consistency of `org_id` will only be enforced if both are present. If either are missing, the trace will be continued. + + The client's organization ID is extracted from the DSN or can be set with the `org_id` option. + If the organization IDs do not match, the SDK will start a new trace instead of continuing the incoming one. + This is useful to prevent traces of unknown third-party services from being continued in your application. + + :param org_id: An optional organization ID. The SDK will try to extract if from the DSN in most cases + but you can provide it explicitly for self-hosted and Relay setups. This value is used for + trace propagation and for features like `strict_trace_continuation`. + + :param _experiments: + """ + pass + + +def _get_default_options() -> "dict[str, Any]": + import inspect + + a = inspect.getfullargspec(ClientConstructor.__init__) + defaults = a.defaults or () + kwonlydefaults = a.kwonlydefaults or {} + + return dict( + itertools.chain( + zip(a.args[-len(defaults) :], defaults), + kwonlydefaults.items(), + ) + ) + + +DEFAULT_OPTIONS = _get_default_options() +del _get_default_options + + +VERSION = "2.52.0" diff --git a/lib/python3.12/site-packages/sentry_sdk/debug.py b/lib/python3.12/site-packages/sentry_sdk/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..513ba8813f9176da0a2f3a87c735d998b637410e --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/debug.py @@ -0,0 +1,37 @@ +import sys +import logging +import warnings + +from sentry_sdk import get_client +from sentry_sdk.client import _client_init_debug +from sentry_sdk.utils import logger +from logging import LogRecord + + +class _DebugFilter(logging.Filter): + def filter(self, record: "LogRecord") -> bool: + if _client_init_debug.get(False): + return True + + return get_client().options["debug"] + + +def init_debug_support() -> None: + if not logger.handlers: + configure_logger() + + +def configure_logger() -> None: + _handler = logging.StreamHandler(sys.stderr) + _handler.setFormatter(logging.Formatter(" [sentry] %(levelname)s: %(message)s")) + logger.addHandler(_handler) + logger.setLevel(logging.DEBUG) + logger.addFilter(_DebugFilter()) + + +def configure_debug_hub() -> None: + warnings.warn( + "configure_debug_hub is deprecated. Please remove calls to it, as it is a no-op.", + DeprecationWarning, + stacklevel=2, + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/envelope.py b/lib/python3.12/site-packages/sentry_sdk/envelope.py new file mode 100644 index 0000000000000000000000000000000000000000..307fb26fd62c97bc53c5b07d01451ae8aa0e741a --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/envelope.py @@ -0,0 +1,336 @@ +import io +import json +import mimetypes + +from sentry_sdk.session import Session +from sentry_sdk.utils import json_dumps, capture_internal_exceptions + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Optional + from typing import Union + from typing import Dict + from typing import List + from typing import Iterator + + from sentry_sdk._types import Event, EventDataCategory + + +def parse_json(data: "Union[bytes, str]") -> "Any": + # on some python 3 versions this needs to be bytes + if isinstance(data, bytes): + data = data.decode("utf-8", "replace") + return json.loads(data) + + +class Envelope: + """ + Represents a Sentry Envelope. The calling code is responsible for adhering to the constraints + documented in the Sentry docs: https://develop.sentry.dev/sdk/envelopes/#data-model. In particular, + each envelope may have at most one Item with type "event" or "transaction" (but not both). + """ + + def __init__( + self, + headers: "Optional[Dict[str, Any]]" = None, + items: "Optional[List[Item]]" = None, + ) -> None: + if headers is not None: + headers = dict(headers) + self.headers = headers or {} + if items is None: + items = [] + else: + items = list(items) + self.items = items + + @property + def description(self) -> str: + return "envelope with %s items (%s)" % ( + len(self.items), + ", ".join(x.data_category for x in self.items), + ) + + def add_event( + self, + event: "Event", + ) -> None: + self.add_item(Item(payload=PayloadRef(json=event), type="event")) + + def add_transaction( + self, + transaction: "Event", + ) -> None: + self.add_item(Item(payload=PayloadRef(json=transaction), type="transaction")) + + def add_profile( + self, + profile: "Any", + ) -> None: + self.add_item(Item(payload=PayloadRef(json=profile), type="profile")) + + def add_profile_chunk( + self, + profile_chunk: "Any", + ) -> None: + self.add_item( + Item( + payload=PayloadRef(json=profile_chunk), + type="profile_chunk", + headers={"platform": profile_chunk.get("platform", "python")}, + ) + ) + + def add_checkin( + self, + checkin: "Any", + ) -> None: + self.add_item(Item(payload=PayloadRef(json=checkin), type="check_in")) + + def add_session( + self, + session: "Union[Session, Any]", + ) -> None: + if isinstance(session, Session): + session = session.to_json() + self.add_item(Item(payload=PayloadRef(json=session), type="session")) + + def add_sessions( + self, + sessions: "Any", + ) -> None: + self.add_item(Item(payload=PayloadRef(json=sessions), type="sessions")) + + def add_item( + self, + item: "Item", + ) -> None: + self.items.append(item) + + def get_event(self) -> "Optional[Event]": + for items in self.items: + event = items.get_event() + if event is not None: + return event + return None + + def get_transaction_event(self) -> "Optional[Event]": + for item in self.items: + event = item.get_transaction_event() + if event is not None: + return event + return None + + def __iter__(self) -> "Iterator[Item]": + return iter(self.items) + + def serialize_into( + self, + f: "Any", + ) -> None: + f.write(json_dumps(self.headers)) + f.write(b"\n") + for item in self.items: + item.serialize_into(f) + + def serialize(self) -> bytes: + out = io.BytesIO() + self.serialize_into(out) + return out.getvalue() + + @classmethod + def deserialize_from( + cls, + f: "Any", + ) -> "Envelope": + headers = parse_json(f.readline()) + items = [] + while 1: + item = Item.deserialize_from(f) + if item is None: + break + items.append(item) + return cls(headers=headers, items=items) + + @classmethod + def deserialize( + cls, + bytes: bytes, + ) -> "Envelope": + return cls.deserialize_from(io.BytesIO(bytes)) + + def __repr__(self) -> str: + return "" % (self.headers, self.items) + + +class PayloadRef: + def __init__( + self, + bytes: "Optional[bytes]" = None, + path: "Optional[Union[bytes, str]]" = None, + json: "Optional[Any]" = None, + ) -> None: + self.json = json + self.bytes = bytes + self.path = path + + def get_bytes(self) -> bytes: + if self.bytes is None: + if self.path is not None: + with capture_internal_exceptions(): + with open(self.path, "rb") as f: + self.bytes = f.read() + elif self.json is not None: + self.bytes = json_dumps(self.json) + return self.bytes or b"" + + @property + def inferred_content_type(self) -> str: + if self.json is not None: + return "application/json" + elif self.path is not None: + path = self.path + if isinstance(path, bytes): + path = path.decode("utf-8", "replace") + ty = mimetypes.guess_type(path)[0] + if ty: + return ty + return "application/octet-stream" + + def __repr__(self) -> str: + return "" % (self.inferred_content_type,) + + +class Item: + def __init__( + self, + payload: "Union[bytes, str, PayloadRef]", + headers: "Optional[Dict[str, Any]]" = None, + type: "Optional[str]" = None, + content_type: "Optional[str]" = None, + filename: "Optional[str]" = None, + ): + if headers is not None: + headers = dict(headers) + elif headers is None: + headers = {} + self.headers = headers + if isinstance(payload, bytes): + payload = PayloadRef(bytes=payload) + elif isinstance(payload, str): + payload = PayloadRef(bytes=payload.encode("utf-8")) + else: + payload = payload + + if filename is not None: + headers["filename"] = filename + if type is not None: + headers["type"] = type + if content_type is not None: + headers["content_type"] = content_type + elif "content_type" not in headers: + headers["content_type"] = payload.inferred_content_type + + self.payload = payload + + def __repr__(self) -> str: + return "" % ( + self.headers, + self.payload, + self.data_category, + ) + + @property + def type(self) -> "Optional[str]": + return self.headers.get("type") + + @property + def data_category(self) -> "EventDataCategory": + ty = self.headers.get("type") + if ty == "session" or ty == "sessions": + return "session" + elif ty == "attachment": + return "attachment" + elif ty == "transaction": + return "transaction" + elif ty == "event": + return "error" + elif ty == "log": + return "log_item" + elif ty == "trace_metric": + return "trace_metric" + elif ty == "client_report": + return "internal" + elif ty == "profile": + return "profile" + elif ty == "profile_chunk": + return "profile_chunk" + elif ty == "check_in": + return "monitor" + else: + return "default" + + def get_bytes(self) -> bytes: + return self.payload.get_bytes() + + def get_event(self) -> "Optional[Event]": + """ + Returns an error event if there is one. + """ + if self.type == "event" and self.payload.json is not None: + return self.payload.json + return None + + def get_transaction_event(self) -> "Optional[Event]": + if self.type == "transaction" and self.payload.json is not None: + return self.payload.json + return None + + def serialize_into( + self, + f: "Any", + ) -> None: + headers = dict(self.headers) + bytes = self.get_bytes() + headers["length"] = len(bytes) + f.write(json_dumps(headers)) + f.write(b"\n") + f.write(bytes) + f.write(b"\n") + + def serialize(self) -> bytes: + out = io.BytesIO() + self.serialize_into(out) + return out.getvalue() + + @classmethod + def deserialize_from( + cls, + f: "Any", + ) -> "Optional[Item]": + line = f.readline().rstrip() + if not line: + return None + headers = parse_json(line) + length = headers.get("length") + if length is not None: + payload = f.read(length) + f.readline() + else: + # if no length was specified we need to read up to the end of line + # and remove it (if it is present, i.e. not the very last char in an eof terminated envelope) + payload = f.readline().rstrip(b"\n") + if headers.get("type") in ("event", "transaction"): + rv = cls(headers=headers, payload=PayloadRef(json=parse_json(payload))) + else: + rv = cls(headers=headers, payload=payload) + return rv + + @classmethod + def deserialize( + cls, + bytes: bytes, + ) -> "Optional[Item]": + return cls.deserialize_from(io.BytesIO(bytes)) diff --git a/lib/python3.12/site-packages/sentry_sdk/feature_flags.py b/lib/python3.12/site-packages/sentry_sdk/feature_flags.py new file mode 100644 index 0000000000000000000000000000000000000000..c9f3f303f98111e32a539feb7f609c053406b745 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/feature_flags.py @@ -0,0 +1,65 @@ +import copy +import sentry_sdk +from sentry_sdk._lru_cache import LRUCache +from threading import Lock + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from typing import TypedDict + + FlagData = TypedDict("FlagData", {"flag": str, "result": bool}) + + +DEFAULT_FLAG_CAPACITY = 100 + + +class FlagBuffer: + def __init__(self, capacity: int) -> None: + self.capacity = capacity + self.lock = Lock() + + # Buffer is private. The name is mangled to discourage use. If you use this attribute + # directly you're on your own! + self.__buffer = LRUCache(capacity) + + def clear(self) -> None: + self.__buffer = LRUCache(self.capacity) + + def __deepcopy__(self, memo: "dict[int, Any]") -> "FlagBuffer": + with self.lock: + buffer = FlagBuffer(self.capacity) + buffer.__buffer = copy.deepcopy(self.__buffer, memo) + return buffer + + def get(self) -> "list[FlagData]": + with self.lock: + return [ + {"flag": key, "result": value} for key, value in self.__buffer.get_all() + ] + + def set(self, flag: str, result: bool) -> None: + if isinstance(result, FlagBuffer): + # If someone were to insert `self` into `self` this would create a circular dependency + # on the lock. This is of course a deadlock. However, this is far outside the expected + # usage of this class. We guard against it here for completeness and to document this + # expected failure mode. + raise ValueError( + "FlagBuffer instances can not be inserted into the dictionary." + ) + + with self.lock: + self.__buffer.set(flag, result) + + +def add_feature_flag(flag: str, result: bool) -> None: + """ + Records a flag and its value to be sent on subsequent error events. + We recommend you do this on flag evaluations. Flags are buffered per Sentry scope. + """ + flags = sentry_sdk.get_isolation_scope().flags + flags.set(flag, result) + + span = sentry_sdk.get_current_span() + if span: + span.set_flag(f"flag.evaluation.{flag}", result) diff --git a/lib/python3.12/site-packages/sentry_sdk/hub.py b/lib/python3.12/site-packages/sentry_sdk/hub.py new file mode 100644 index 0000000000000000000000000000000000000000..0e5d7df9f954e1b63f95bcfd245a75df200edea5 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/hub.py @@ -0,0 +1,741 @@ +import warnings +from contextlib import contextmanager +from typing import TYPE_CHECKING + +from sentry_sdk import ( + get_client, + get_current_scope, + get_global_scope, + get_isolation_scope, +) +from sentry_sdk._compat import with_metaclass +from sentry_sdk.client import Client +from sentry_sdk.consts import INSTRUMENTER +from sentry_sdk.scope import _ScopeManager +from sentry_sdk.tracing import ( + NoOpSpan, + Span, + Transaction, +) +from sentry_sdk.utils import ( + ContextVar, + logger, +) + +if TYPE_CHECKING: + from typing import ( + Any, + Callable, + ContextManager, + Dict, + Generator, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, + overload, + ) + + from typing_extensions import Unpack + + from sentry_sdk._types import ( + Breadcrumb, + BreadcrumbHint, + Event, + ExcInfo, + Hint, + LogLevelStr, + SamplingContext, + ) + from sentry_sdk.client import BaseClient + from sentry_sdk.integrations import Integration + from sentry_sdk.scope import Scope + from sentry_sdk.tracing import TransactionKwargs + + T = TypeVar("T") + +else: + + def overload(x: "T") -> "T": + return x + + +class SentryHubDeprecationWarning(DeprecationWarning): + """ + A custom deprecation warning to inform users that the Hub is deprecated. + """ + + _MESSAGE = ( + "`sentry_sdk.Hub` is deprecated and will be removed in a future major release. " + "Please consult our 1.x to 2.x migration guide for details on how to migrate " + "`Hub` usage to the new API: " + "https://docs.sentry.io/platforms/python/migration/1.x-to-2.x" + ) + + def __init__(self, *_: object) -> None: + super().__init__(self._MESSAGE) + + +@contextmanager +def _suppress_hub_deprecation_warning() -> "Generator[None, None, None]": + """Utility function to suppress deprecation warnings for the Hub.""" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=SentryHubDeprecationWarning) + yield + + +_local = ContextVar("sentry_current_hub") + + +class HubMeta(type): + @property + def current(cls) -> "Hub": + """Returns the current instance of the hub.""" + warnings.warn(SentryHubDeprecationWarning(), stacklevel=2) + rv = _local.get(None) + if rv is None: + with _suppress_hub_deprecation_warning(): + # This will raise a deprecation warning; suppress it since we already warned above. + rv = Hub(GLOBAL_HUB) + _local.set(rv) + return rv + + @property + def main(cls) -> "Hub": + """Returns the main instance of the hub.""" + warnings.warn(SentryHubDeprecationWarning(), stacklevel=2) + return GLOBAL_HUB + + +class Hub(with_metaclass(HubMeta)): # type: ignore + """ + .. deprecated:: 2.0.0 + The Hub is deprecated. Its functionality will be merged into :py:class:`sentry_sdk.scope.Scope`. + + The hub wraps the concurrency management of the SDK. Each thread has + its own hub but the hub might transfer with the flow of execution if + context vars are available. + + If the hub is used with a with statement it's temporarily activated. + """ + + _stack: "List[Tuple[Optional[Client], Scope]]" = None # type: ignore[assignment] + _scope: "Optional[Scope]" = None + + # Mypy doesn't pick up on the metaclass. + + if TYPE_CHECKING: + current: "Hub" = None # type: ignore[assignment] + main: "Optional[Hub]" = None + + def __init__( + self, + client_or_hub: "Optional[Union[Hub, Client]]" = None, + scope: "Optional[Any]" = None, + ) -> None: + warnings.warn(SentryHubDeprecationWarning(), stacklevel=2) + + current_scope = None + + if isinstance(client_or_hub, Hub): + client = get_client() + if scope is None: + # hub cloning is going on, we use a fork of the current/isolation scope for context manager + scope = get_isolation_scope().fork() + current_scope = get_current_scope().fork() + else: + client = client_or_hub # type: ignore + get_global_scope().set_client(client) + + if scope is None: # so there is no Hub cloning going on + # just the current isolation scope is used for context manager + scope = get_isolation_scope() + current_scope = get_current_scope() + + if current_scope is None: + # just the current current scope is used for context manager + current_scope = get_current_scope() + + self._stack = [(client, scope)] # type: ignore + self._last_event_id: "Optional[str]" = None + self._old_hubs: "List[Hub]" = [] + + self._old_current_scopes: "List[Scope]" = [] + self._old_isolation_scopes: "List[Scope]" = [] + self._current_scope: "Scope" = current_scope + self._scope: "Scope" = scope + + def __enter__(self) -> "Hub": + self._old_hubs.append(Hub.current) + _local.set(self) + + current_scope = get_current_scope() + self._old_current_scopes.append(current_scope) + scope._current_scope.set(self._current_scope) + + isolation_scope = get_isolation_scope() + self._old_isolation_scopes.append(isolation_scope) + scope._isolation_scope.set(self._scope) + + return self + + def __exit__( + self, + exc_type: "Optional[type]", + exc_value: "Optional[BaseException]", + tb: "Optional[Any]", + ) -> None: + old = self._old_hubs.pop() + _local.set(old) + + old_current_scope = self._old_current_scopes.pop() + scope._current_scope.set(old_current_scope) + + old_isolation_scope = self._old_isolation_scopes.pop() + scope._isolation_scope.set(old_isolation_scope) + + def run( + self, + callback: "Callable[[], T]", + ) -> "T": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + + Runs a callback in the context of the hub. Alternatively the + with statement can be used on the hub directly. + """ + with self: + return callback() + + def get_integration( + self, + name_or_class: "Union[str, Type[Integration]]", + ) -> "Any": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.client._Client.get_integration` instead. + + Returns the integration for this hub by name or class. If there + is no client bound or the client does not have that integration + then `None` is returned. + + If the return value is not `None` the hub is guaranteed to have a + client attached. + """ + return get_client().get_integration(name_or_class) + + @property + def client(self) -> "Optional[BaseClient]": + """ + .. deprecated:: 2.0.0 + This property is deprecated and will be removed in a future release. + Please use :py:func:`sentry_sdk.api.get_client` instead. + + Returns the current client on the hub. + """ + client = get_client() + + if not client.is_active(): + return None + + return client + + @property + def scope(self) -> "Scope": + """ + .. deprecated:: 2.0.0 + This property is deprecated and will be removed in a future release. + Returns the current scope on the hub. + """ + return get_isolation_scope() + + def last_event_id(self) -> "Optional[str]": + """ + Returns the last event ID. + + .. deprecated:: 1.40.5 + This function is deprecated and will be removed in a future release. The functions `capture_event`, `capture_message`, and `capture_exception` return the event ID directly. + """ + logger.warning( + "Deprecated: last_event_id is deprecated. This will be removed in the future. The functions `capture_event`, `capture_message`, and `capture_exception` return the event ID directly." + ) + return self._last_event_id + + def bind_client( + self, + new: "Optional[BaseClient]", + ) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.set_client` instead. + + Binds a new client to the hub. + """ + get_global_scope().set_client(new) + + def capture_event( + self, + event: "Event", + hint: "Optional[Hint]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.capture_event` instead. + + Captures an event. + + Alias of :py:meth:`sentry_sdk.Scope.capture_event`. + + :param event: A ready-made event that can be directly sent to Sentry. + + :param hint: Contains metadata about the event that can be read from `before_send`, such as the original exception object or a HTTP request object. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + """ + last_event_id = get_current_scope().capture_event( + event, hint, scope=scope, **scope_kwargs + ) + + is_transaction = event.get("type") == "transaction" + if last_event_id is not None and not is_transaction: + self._last_event_id = last_event_id + + return last_event_id + + def capture_message( + self, + message: str, + level: "Optional[LogLevelStr]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.capture_message` instead. + + Captures a message. + + Alias of :py:meth:`sentry_sdk.Scope.capture_message`. + + :param message: The string to send as the message to Sentry. + + :param level: If no level is provided, the default level is `info`. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`). + """ + last_event_id = get_current_scope().capture_message( + message, level=level, scope=scope, **scope_kwargs + ) + + if last_event_id is not None: + self._last_event_id = last_event_id + + return last_event_id + + def capture_exception( + self, + error: "Optional[Union[BaseException, ExcInfo]]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.capture_exception` instead. + + Captures an exception. + + Alias of :py:meth:`sentry_sdk.Scope.capture_exception`. + + :param error: An exception to capture. If `None`, `sys.exc_info()` will be used. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`). + """ + last_event_id = get_current_scope().capture_exception( + error, scope=scope, **scope_kwargs + ) + + if last_event_id is not None: + self._last_event_id = last_event_id + + return last_event_id + + def add_breadcrumb( + self, + crumb: "Optional[Breadcrumb]" = None, + hint: "Optional[BreadcrumbHint]" = None, + **kwargs: "Any", + ) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.add_breadcrumb` instead. + + Adds a breadcrumb. + + :param crumb: Dictionary with the data as the sentry v7/v8 protocol expects. + + :param hint: An optional value that can be used by `before_breadcrumb` + to customize the breadcrumbs that are emitted. + """ + get_isolation_scope().add_breadcrumb(crumb, hint, **kwargs) + + def start_span( + self, instrumenter: str = INSTRUMENTER.SENTRY, **kwargs: "Any" + ) -> "Span": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.start_span` instead. + + Start a span whose parent is the currently active span or transaction, if any. + + The return value is a :py:class:`sentry_sdk.tracing.Span` instance, + typically used as a context manager to start and stop timing in a `with` + block. + + Only spans contained in a transaction are sent to Sentry. Most + integrations start a transaction at the appropriate time, for example + for every incoming HTTP request. Use + :py:meth:`sentry_sdk.start_transaction` to start a new transaction when + one is not already in progress. + + For supported `**kwargs` see :py:class:`sentry_sdk.tracing.Span`. + """ + scope = get_current_scope() + return scope.start_span(instrumenter=instrumenter, **kwargs) + + def start_transaction( + self, + transaction: "Optional[Transaction]" = None, + instrumenter: str = INSTRUMENTER.SENTRY, + custom_sampling_context: "Optional[SamplingContext]" = None, + **kwargs: "Unpack[TransactionKwargs]", + ) -> "Union[Transaction, NoOpSpan]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.start_transaction` instead. + + Start and return a transaction. + + Start an existing transaction if given, otherwise create and start a new + transaction with kwargs. + + This is the entry point to manual tracing instrumentation. + + A tree structure can be built by adding child spans to the transaction, + and child spans to other spans. To start a new child span within the + transaction or any span, call the respective `.start_child()` method. + + Every child span must be finished before the transaction is finished, + otherwise the unfinished spans are discarded. + + When used as context managers, spans and transactions are automatically + finished at the end of the `with` block. If not using context managers, + call the `.finish()` method. + + When the transaction is finished, it will be sent to Sentry with all its + finished child spans. + + For supported `**kwargs` see :py:class:`sentry_sdk.tracing.Transaction`. + """ + scope = get_current_scope() + + # For backwards compatibility, we allow passing the scope as the hub. + # We need a major release to make this nice. (if someone searches the code: deprecated) + # Type checking disabled for this line because deprecated keys are not allowed in the type signature. + kwargs["hub"] = scope # type: ignore + + return scope.start_transaction( + transaction, instrumenter, custom_sampling_context, **kwargs + ) + + def continue_trace( + self, + environ_or_headers: "Dict[str, Any]", + op: "Optional[str]" = None, + name: "Optional[str]" = None, + source: "Optional[str]" = None, + ) -> "Transaction": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.continue_trace` instead. + + Sets the propagation context from environment or headers and returns a transaction. + """ + return get_isolation_scope().continue_trace( + environ_or_headers=environ_or_headers, op=op, name=name, source=source + ) + + @overload + def push_scope( + self, + callback: "Optional[None]" = None, + ) -> "ContextManager[Scope]": + pass + + @overload + def push_scope( # noqa: F811 + self, + callback: "Callable[[Scope], None]", + ) -> None: + pass + + def push_scope( # noqa + self, + callback: "Optional[Callable[[Scope], None]]" = None, + continue_trace: bool = True, + ) -> "Optional[ContextManager[Scope]]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + + Pushes a new layer on the scope stack. + + :param callback: If provided, this method pushes a scope, calls + `callback`, and pops the scope again. + + :returns: If no `callback` is provided, a context manager that should + be used to pop the scope again. + """ + if callback is not None: + with self.push_scope() as scope: + callback(scope) + return None + + return _ScopeManager(self) + + def pop_scope_unsafe(self) -> "Tuple[Optional[Client], Scope]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + + Pops a scope layer from the stack. + + Try to use the context manager :py:meth:`push_scope` instead. + """ + rv = self._stack.pop() + assert self._stack, "stack must have at least one layer" + return rv + + @overload + def configure_scope( + self, + callback: "Optional[None]" = None, + ) -> "ContextManager[Scope]": + pass + + @overload + def configure_scope( # noqa: F811 + self, + callback: "Callable[[Scope], None]", + ) -> None: + pass + + def configure_scope( # noqa + self, + callback: "Optional[Callable[[Scope], None]]" = None, + continue_trace: bool = True, + ) -> "Optional[ContextManager[Scope]]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + + Reconfigures the scope. + + :param callback: If provided, call the callback with the current scope. + + :returns: If no callback is provided, returns a context manager that returns the scope. + """ + scope = get_isolation_scope() + + if continue_trace: + scope.generate_propagation_context() + + if callback is not None: + # TODO: used to return None when client is None. Check if this changes behavior. + callback(scope) + + return None + + @contextmanager + def inner() -> "Generator[Scope, None, None]": + yield scope + + return inner() + + def start_session( + self, + session_mode: str = "application", + ) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.start_session` instead. + + Starts a new session. + """ + get_isolation_scope().start_session( + session_mode=session_mode, + ) + + def end_session(self) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.end_session` instead. + + Ends the current session if there is one. + """ + get_isolation_scope().end_session() + + def stop_auto_session_tracking(self) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.stop_auto_session_tracking` instead. + + Stops automatic session tracking. + + This temporarily session tracking for the current scope when called. + To resume session tracking call `resume_auto_session_tracking`. + """ + get_isolation_scope().stop_auto_session_tracking() + + def resume_auto_session_tracking(self) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.resume_auto_session_tracking` instead. + + Resumes automatic session tracking for the current scope if + disabled earlier. This requires that generally automatic session + tracking is enabled. + """ + get_isolation_scope().resume_auto_session_tracking() + + def flush( + self, + timeout: "Optional[float]" = None, + callback: "Optional[Callable[[int, float], None]]" = None, + ) -> None: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.client._Client.flush` instead. + + Alias for :py:meth:`sentry_sdk.client._Client.flush` + """ + return get_client().flush(timeout=timeout, callback=callback) + + def get_traceparent(self) -> "Optional[str]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.get_traceparent` instead. + + Returns the traceparent either from the active span or from the scope. + """ + current_scope = get_current_scope() + traceparent = current_scope.get_traceparent() + + if traceparent is None: + isolation_scope = get_isolation_scope() + traceparent = isolation_scope.get_traceparent() + + return traceparent + + def get_baggage(self) -> "Optional[str]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.get_baggage` instead. + + Returns Baggage either from the active span or from the scope. + """ + current_scope = get_current_scope() + baggage = current_scope.get_baggage() + + if baggage is None: + isolation_scope = get_isolation_scope() + baggage = isolation_scope.get_baggage() + + if baggage is not None: + return baggage.serialize() + + return None + + def iter_trace_propagation_headers( + self, span: "Optional[Span]" = None + ) -> "Generator[Tuple[str, str], None, None]": + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.iter_trace_propagation_headers` instead. + + Return HTTP headers which allow propagation of trace data. Data taken + from the span representing the request, if available, or the current + span on the scope if not. + """ + return get_current_scope().iter_trace_propagation_headers( + span=span, + ) + + def trace_propagation_meta(self, span: "Optional[Span]" = None) -> str: + """ + .. deprecated:: 2.0.0 + This function is deprecated and will be removed in a future release. + Please use :py:meth:`sentry_sdk.Scope.trace_propagation_meta` instead. + + Return meta tags which should be injected into HTML templates + to allow propagation of trace information. + """ + if span is not None: + logger.warning( + "The parameter `span` in trace_propagation_meta() is deprecated and will be removed in the future." + ) + + return get_current_scope().trace_propagation_meta( + span=span, + ) + + +with _suppress_hub_deprecation_warning(): + # Suppress deprecation warning for the Hub here, since we still always + # import this module. + GLOBAL_HUB = Hub() +_local.set(GLOBAL_HUB) + + +# Circular imports +from sentry_sdk import scope diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/__init__.py b/lib/python3.12/site-packages/sentry_sdk/integrations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd12a6011f4cd21898ec836af9fb0fa1eadbea97 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/__init__.py @@ -0,0 +1,348 @@ +from abc import ABC, abstractmethod +from threading import Lock +from typing import TYPE_CHECKING + +from sentry_sdk.utils import logger + +if TYPE_CHECKING: + from collections.abc import Sequence + from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Type, Union + + +_DEFAULT_FAILED_REQUEST_STATUS_CODES = frozenset(range(500, 600)) + + +_installer_lock = Lock() + +# Set of all integration identifiers we have attempted to install +_processed_integrations: "Set[str]" = set() + +# Set of all integration identifiers we have actually installed +_installed_integrations: "Set[str]" = set() + + +def _generate_default_integrations_iterator( + integrations: "List[str]", + auto_enabling_integrations: "List[str]", +) -> "Callable[[bool], Iterator[Type[Integration]]]": + def iter_default_integrations( + with_auto_enabling_integrations: bool, + ) -> "Iterator[Type[Integration]]": + """Returns an iterator of the default integration classes:""" + from importlib import import_module + + if with_auto_enabling_integrations: + all_import_strings = integrations + auto_enabling_integrations + else: + all_import_strings = integrations + + for import_string in all_import_strings: + try: + module, cls = import_string.rsplit(".", 1) + yield getattr(import_module(module), cls) + except (DidNotEnable, SyntaxError) as e: + logger.debug( + "Did not import default integration %s: %s", import_string, e + ) + + if isinstance(iter_default_integrations.__doc__, str): + for import_string in integrations: + iter_default_integrations.__doc__ += "\n- `{}`".format(import_string) + + return iter_default_integrations + + +_DEFAULT_INTEGRATIONS = [ + # stdlib/base runtime integrations + "sentry_sdk.integrations.argv.ArgvIntegration", + "sentry_sdk.integrations.atexit.AtexitIntegration", + "sentry_sdk.integrations.dedupe.DedupeIntegration", + "sentry_sdk.integrations.excepthook.ExcepthookIntegration", + "sentry_sdk.integrations.logging.LoggingIntegration", + "sentry_sdk.integrations.modules.ModulesIntegration", + "sentry_sdk.integrations.stdlib.StdlibIntegration", + "sentry_sdk.integrations.threading.ThreadingIntegration", +] + +_AUTO_ENABLING_INTEGRATIONS = [ + "sentry_sdk.integrations.aiohttp.AioHttpIntegration", + "sentry_sdk.integrations.anthropic.AnthropicIntegration", + "sentry_sdk.integrations.ariadne.AriadneIntegration", + "sentry_sdk.integrations.arq.ArqIntegration", + "sentry_sdk.integrations.asyncpg.AsyncPGIntegration", + "sentry_sdk.integrations.boto3.Boto3Integration", + "sentry_sdk.integrations.bottle.BottleIntegration", + "sentry_sdk.integrations.celery.CeleryIntegration", + "sentry_sdk.integrations.chalice.ChaliceIntegration", + "sentry_sdk.integrations.clickhouse_driver.ClickhouseDriverIntegration", + "sentry_sdk.integrations.cohere.CohereIntegration", + "sentry_sdk.integrations.django.DjangoIntegration", + "sentry_sdk.integrations.falcon.FalconIntegration", + "sentry_sdk.integrations.fastapi.FastApiIntegration", + "sentry_sdk.integrations.flask.FlaskIntegration", + "sentry_sdk.integrations.gql.GQLIntegration", + "sentry_sdk.integrations.google_genai.GoogleGenAIIntegration", + "sentry_sdk.integrations.graphene.GrapheneIntegration", + "sentry_sdk.integrations.httpx.HttpxIntegration", + "sentry_sdk.integrations.huey.HueyIntegration", + "sentry_sdk.integrations.huggingface_hub.HuggingfaceHubIntegration", + "sentry_sdk.integrations.langchain.LangchainIntegration", + "sentry_sdk.integrations.langgraph.LanggraphIntegration", + "sentry_sdk.integrations.litestar.LitestarIntegration", + "sentry_sdk.integrations.loguru.LoguruIntegration", + "sentry_sdk.integrations.mcp.MCPIntegration", + "sentry_sdk.integrations.openai.OpenAIIntegration", + "sentry_sdk.integrations.openai_agents.OpenAIAgentsIntegration", + "sentry_sdk.integrations.pydantic_ai.PydanticAIIntegration", + "sentry_sdk.integrations.pymongo.PyMongoIntegration", + "sentry_sdk.integrations.pyramid.PyramidIntegration", + "sentry_sdk.integrations.quart.QuartIntegration", + "sentry_sdk.integrations.redis.RedisIntegration", + "sentry_sdk.integrations.rq.RqIntegration", + "sentry_sdk.integrations.sanic.SanicIntegration", + "sentry_sdk.integrations.sqlalchemy.SqlalchemyIntegration", + "sentry_sdk.integrations.starlette.StarletteIntegration", + "sentry_sdk.integrations.starlite.StarliteIntegration", + "sentry_sdk.integrations.strawberry.StrawberryIntegration", + "sentry_sdk.integrations.tornado.TornadoIntegration", +] + +iter_default_integrations = _generate_default_integrations_iterator( + integrations=_DEFAULT_INTEGRATIONS, + auto_enabling_integrations=_AUTO_ENABLING_INTEGRATIONS, +) + +del _generate_default_integrations_iterator + + +_MIN_VERSIONS = { + "aiohttp": (3, 4), + "anthropic": (0, 16), + "ariadne": (0, 20), + "arq": (0, 23), + "asyncpg": (0, 23), + "beam": (2, 12), + "boto3": (1, 12), # botocore + "bottle": (0, 12), + "celery": (4, 4, 7), + "chalice": (1, 16, 0), + "clickhouse_driver": (0, 2, 0), + "cohere": (5, 4, 0), + "django": (1, 8), + "dramatiq": (1, 9), + "falcon": (1, 4), + "fastapi": (0, 79, 0), + "flask": (1, 1, 4), + "gql": (3, 4, 1), + "graphene": (3, 3), + "google_genai": (1, 29, 0), # google-genai + "grpc": (1, 32, 0), # grpcio + "httpx": (0, 16, 0), + "huggingface_hub": (0, 24, 7), + "langchain": (0, 1, 0), + "langgraph": (0, 6, 6), + "launchdarkly": (9, 8, 0), + "litellm": (1, 77, 5), + "loguru": (0, 7, 0), + "mcp": (1, 15, 0), + "openai": (1, 0, 0), + "openai_agents": (0, 0, 19), + "openfeature": (0, 7, 1), + "pydantic_ai": (1, 0, 0), + "pymongo": (3, 5, 0), + "quart": (0, 16, 0), + "ray": (2, 7, 0), + "requests": (2, 0, 0), + "rq": (0, 6), + "sanic": (0, 8), + "sqlalchemy": (1, 2), + "starlette": (0, 16), + "starlite": (1, 48), + "statsig": (0, 55, 3), + "strawberry": (0, 209, 5), + "tornado": (6, 0), + "typer": (0, 15), + "unleash": (6, 0, 1), +} + + +_INTEGRATION_DEACTIVATES = { + "langchain": {"openai", "anthropic", "google_genai"}, + "openai_agents": {"openai"}, + "pydantic_ai": {"openai", "anthropic"}, +} + + +def setup_integrations( + integrations: "Sequence[Integration]", + with_defaults: bool = True, + with_auto_enabling_integrations: bool = False, + disabled_integrations: "Optional[Sequence[Union[type[Integration], Integration]]]" = None, + options: "Optional[Dict[str, Any]]" = None, +) -> "Dict[str, Integration]": + """ + Given a list of integration instances, this installs them all. + + When `with_defaults` is set to `True` all default integrations are added + unless they were already provided before. + + `disabled_integrations` takes precedence over `with_defaults` and + `with_auto_enabling_integrations`. + + Some integrations are designed to automatically deactivate other integrations + in order to avoid conflicts and prevent duplicate telemetry from being collected. + For example, enabling the `langchain` integration will auto-deactivate both the + `openai` and `anthropic` integrations. + + Users can override this behavior by: + - Explicitly providing an integration in the `integrations=[]` list, or + - Disabling the higher-level integration via the `disabled_integrations` option. + """ + integrations = dict( + (integration.identifier, integration) for integration in integrations or () + ) + + logger.debug("Setting up integrations (with default = %s)", with_defaults) + + user_provided_integrations = set(integrations.keys()) + + # Integrations that will not be enabled + disabled_integrations = [ + integration if isinstance(integration, type) else type(integration) + for integration in disabled_integrations or [] + ] + + # Integrations that are not explicitly set up by the user. + used_as_default_integration = set() + + if with_defaults: + for integration_cls in iter_default_integrations( + with_auto_enabling_integrations + ): + if integration_cls.identifier not in integrations: + instance = integration_cls() + integrations[instance.identifier] = instance + used_as_default_integration.add(instance.identifier) + + disabled_integration_identifiers = { + integration.identifier for integration in disabled_integrations + } + + for integration, targets_to_deactivate in _INTEGRATION_DEACTIVATES.items(): + if ( + integration in integrations + and integration not in disabled_integration_identifiers + ): + for target in targets_to_deactivate: + if target not in user_provided_integrations: + for cls in iter_default_integrations(True): + if cls.identifier == target: + if cls not in disabled_integrations: + disabled_integrations.append(cls) + logger.debug( + "Auto-deactivating %s integration because %s integration is active", + target, + integration, + ) + + for identifier, integration in integrations.items(): + with _installer_lock: + if identifier not in _processed_integrations: + if type(integration) in disabled_integrations: + logger.debug("Ignoring integration %s", identifier) + else: + logger.debug( + "Setting up previously not enabled integration %s", identifier + ) + try: + type(integration).setup_once() + integration.setup_once_with_options(options) + except DidNotEnable as e: + if identifier not in used_as_default_integration: + raise + + logger.debug( + "Did not enable default integration %s: %s", identifier, e + ) + else: + _installed_integrations.add(identifier) + + _processed_integrations.add(identifier) + + integrations = { + identifier: integration + for identifier, integration in integrations.items() + if identifier in _installed_integrations + } + + for identifier in integrations: + logger.debug("Enabling integration %s", identifier) + + return integrations + + +def _check_minimum_version( + integration: "type[Integration]", + version: "Optional[tuple[int, ...]]", + package: "Optional[str]" = None, +) -> None: + package = package or integration.identifier + + if version is None: + raise DidNotEnable(f"Unparsable {package} version.") + + min_version = _MIN_VERSIONS.get(integration.identifier) + if min_version is None: + return + + if version < min_version: + raise DidNotEnable( + f"Integration only supports {package} {'.'.join(map(str, min_version))} or newer." + ) + + +class DidNotEnable(Exception): # noqa: N818 + """ + The integration could not be enabled due to a trivial user error like + `flask` not being installed for the `FlaskIntegration`. + + This exception is silently swallowed for default integrations, but reraised + for explicitly enabled integrations. + """ + + +class Integration(ABC): + """Baseclass for all integrations. + + To accept options for an integration, implement your own constructor that + saves those options on `self`. + """ + + install = None + """Legacy method, do not implement.""" + + identifier: "str" = None # type: ignore[assignment] + """String unique ID of integration type""" + + @staticmethod + @abstractmethod + def setup_once() -> None: + """ + Initialize the integration. + + This function is only called once, ever. Configuration is not available + at this point, so the only thing to do here is to hook into exception + handlers, and perhaps do monkeypatches. + + Inside those hooks `Integration.current` can be used to access the + instance again. + """ + pass + + def setup_once_with_options( + self, options: "Optional[Dict[str, Any]]" = None + ) -> None: + """ + Called after setup_once in rare cases on the instance and with options since we don't have those available above. + """ + pass diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/_asgi_common.py b/lib/python3.12/site-packages/sentry_sdk/integrations/_asgi_common.py new file mode 100644 index 0000000000000000000000000000000000000000..a8022c6bb10903c9a4b56fe5d459884b3140fcd6 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/_asgi_common.py @@ -0,0 +1,107 @@ +import urllib + +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.integrations._wsgi_common import _filter_headers + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Optional + from typing import Union + from typing_extensions import Literal + + from sentry_sdk.utils import AnnotatedValue + + +def _get_headers(asgi_scope: "Any") -> "Dict[str, str]": + """ + Extract headers from the ASGI scope, in the format that the Sentry protocol expects. + """ + headers: "Dict[str, str]" = {} + for raw_key, raw_value in asgi_scope["headers"]: + key = raw_key.decode("latin-1") + value = raw_value.decode("latin-1") + if key in headers: + headers[key] = headers[key] + ", " + value + else: + headers[key] = value + + return headers + + +def _get_url( + asgi_scope: "Dict[str, Any]", + default_scheme: "Literal['ws', 'http']", + host: "Optional[Union[AnnotatedValue, str]]", +) -> str: + """ + Extract URL from the ASGI scope, without also including the querystring. + """ + scheme = asgi_scope.get("scheme", default_scheme) + + server = asgi_scope.get("server", None) + path = asgi_scope.get("root_path", "") + asgi_scope.get("path", "") + + if host: + return "%s://%s%s" % (scheme, host, path) + + if server is not None: + host, port = server + default_port = {"http": 80, "https": 443, "ws": 80, "wss": 443}.get(scheme) + if port != default_port: + return "%s://%s:%s%s" % (scheme, host, port, path) + return "%s://%s%s" % (scheme, host, path) + return path + + +def _get_query(asgi_scope: "Any") -> "Any": + """ + Extract querystring from the ASGI scope, in the format that the Sentry protocol expects. + """ + qs = asgi_scope.get("query_string") + if not qs: + return None + return urllib.parse.unquote(qs.decode("latin-1")) + + +def _get_ip(asgi_scope: "Any") -> str: + """ + Extract IP Address from the ASGI scope based on request headers with fallback to scope client. + """ + headers = _get_headers(asgi_scope) + try: + return headers["x-forwarded-for"].split(",")[0].strip() + except (KeyError, IndexError): + pass + + try: + return headers["x-real-ip"] + except KeyError: + pass + + return asgi_scope.get("client")[0] + + +def _get_request_data(asgi_scope: "Any") -> "Dict[str, Any]": + """ + Returns data related to the HTTP request from the ASGI scope. + """ + request_data: "Dict[str, Any]" = {} + ty = asgi_scope["type"] + if ty in ("http", "websocket"): + request_data["method"] = asgi_scope.get("method") + + request_data["headers"] = headers = _filter_headers(_get_headers(asgi_scope)) + request_data["query_string"] = _get_query(asgi_scope) + + request_data["url"] = _get_url( + asgi_scope, "http" if ty == "http" else "ws", headers.get("host") + ) + + client = asgi_scope.get("client") + if client and should_send_default_pii(): + request_data["env"] = {"REMOTE_ADDR": _get_ip(asgi_scope)} + + return request_data diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/_wsgi_common.py b/lib/python3.12/site-packages/sentry_sdk/integrations/_wsgi_common.py new file mode 100644 index 0000000000000000000000000000000000000000..688e965be43d29b143f6fd0a77eb6f6eed4aab9c --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/_wsgi_common.py @@ -0,0 +1,258 @@ +from contextlib import contextmanager +import json +from copy import deepcopy + +import sentry_sdk +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import AnnotatedValue, logger + +try: + from django.http.request import RawPostDataException +except ImportError: + RawPostDataException = None + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Iterator + from typing import Mapping + from typing import MutableMapping + from typing import Optional + from typing import Union + from sentry_sdk._types import Event, HttpStatusCodeRange + + +SENSITIVE_ENV_KEYS = ( + "REMOTE_ADDR", + "HTTP_X_FORWARDED_FOR", + "HTTP_SET_COOKIE", + "HTTP_COOKIE", + "HTTP_AUTHORIZATION", + "HTTP_X_API_KEY", + "HTTP_X_FORWARDED_FOR", + "HTTP_X_REAL_IP", +) + +SENSITIVE_HEADERS = tuple( + x[len("HTTP_") :] for x in SENSITIVE_ENV_KEYS if x.startswith("HTTP_") +) + +DEFAULT_HTTP_METHODS_TO_CAPTURE = ( + "CONNECT", + "DELETE", + "GET", + # "HEAD", # do not capture HEAD requests by default + # "OPTIONS", # do not capture OPTIONS requests by default + "PATCH", + "POST", + "PUT", + "TRACE", +) + + +# This noop context manager can be replaced with "from contextlib import nullcontext" when we drop Python 3.6 support +@contextmanager +def nullcontext() -> "Iterator[None]": + yield + + +def request_body_within_bounds( + client: "Optional[sentry_sdk.client.BaseClient]", content_length: int +) -> bool: + if client is None: + return False + + bodies = client.options["max_request_body_size"] + return not ( + bodies == "never" + or (bodies == "small" and content_length > 10**3) + or (bodies == "medium" and content_length > 10**4) + ) + + +class RequestExtractor: + """ + Base class for request extraction. + """ + + # It does not make sense to make this class an ABC because it is not used + # for typing, only so that child classes can inherit common methods from + # it. Only some child classes implement all methods that raise + # NotImplementedError in this class. + + def __init__(self, request: "Any") -> None: + self.request = request + + def extract_into_event(self, event: "Event") -> None: + client = sentry_sdk.get_client() + if not client.is_active(): + return + + data: "Optional[Union[AnnotatedValue, Dict[str, Any]]]" = None + + content_length = self.content_length() + request_info = event.get("request", {}) + + if should_send_default_pii(): + request_info["cookies"] = dict(self.cookies()) + + if not request_body_within_bounds(client, content_length): + data = AnnotatedValue.removed_because_over_size_limit() + else: + # First read the raw body data + # It is important to read this first because if it is Django + # it will cache the body and then we can read the cached version + # again in parsed_body() (or json() or wherever). + raw_data = None + try: + raw_data = self.raw_data() + except (RawPostDataException, ValueError): + # If DjangoRestFramework is used it already read the body for us + # so reading it here will fail. We can ignore this. + pass + + parsed_body = self.parsed_body() + if parsed_body is not None: + data = parsed_body + elif raw_data: + data = AnnotatedValue.removed_because_raw_data() + else: + data = None + + if data is not None: + request_info["data"] = data + + event["request"] = deepcopy(request_info) + + def content_length(self) -> int: + try: + return int(self.env().get("CONTENT_LENGTH", 0)) + except ValueError: + return 0 + + def cookies(self) -> "MutableMapping[str, Any]": + raise NotImplementedError() + + def raw_data(self) -> "Optional[Union[str, bytes]]": + raise NotImplementedError() + + def form(self) -> "Optional[Dict[str, Any]]": + raise NotImplementedError() + + def parsed_body(self) -> "Optional[Dict[str, Any]]": + try: + form = self.form() + except Exception: + form = None + try: + files = self.files() + except Exception: + files = None + + if form or files: + data = {} + if form: + data = dict(form.items()) + if files: + for key in files.keys(): + data[key] = AnnotatedValue.removed_because_raw_data() + + return data + + return self.json() + + def is_json(self) -> bool: + return _is_json_content_type(self.env().get("CONTENT_TYPE")) + + def json(self) -> "Optional[Any]": + try: + if not self.is_json(): + return None + + try: + raw_data = self.raw_data() + except (RawPostDataException, ValueError): + # The body might have already been read, in which case this will + # fail + raw_data = None + + if raw_data is None: + return None + + if isinstance(raw_data, str): + return json.loads(raw_data) + else: + return json.loads(raw_data.decode("utf-8")) + except ValueError: + pass + + return None + + def files(self) -> "Optional[Dict[str, Any]]": + raise NotImplementedError() + + def size_of_file(self, file: "Any") -> int: + raise NotImplementedError() + + def env(self) -> "Dict[str, Any]": + raise NotImplementedError() + + +def _is_json_content_type(ct: "Optional[str]") -> bool: + mt = (ct or "").split(";", 1)[0] + return ( + mt == "application/json" + or (mt.startswith("application/")) + and mt.endswith("+json") + ) + + +def _filter_headers( + headers: "Mapping[str, str]", +) -> "Mapping[str, Union[AnnotatedValue, str]]": + if should_send_default_pii(): + return headers + + return { + k: ( + v + if k.upper().replace("-", "_") not in SENSITIVE_HEADERS + else AnnotatedValue.removed_because_over_size_limit() + ) + for k, v in headers.items() + } + + +def _in_http_status_code_range( + code: object, code_ranges: "list[HttpStatusCodeRange]" +) -> bool: + for target in code_ranges: + if isinstance(target, int): + if code == target: + return True + continue + + try: + if code in target: + return True + except TypeError: + logger.warning( + "failed_request_status_codes has to be a list of integers or containers" + ) + + return False + + +class HttpCodeRangeContainer: + """ + Wrapper to make it possible to use list[HttpStatusCodeRange] as a Container[int]. + Used for backwards compatibility with the old `failed_request_status_codes` option. + """ + + def __init__(self, code_ranges: "list[HttpStatusCodeRange]") -> None: + self._code_ranges = code_ranges + + def __contains__(self, item: object) -> bool: + return _in_http_status_code_range(item, self._code_ranges) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/aiohttp.py b/lib/python3.12/site-packages/sentry_sdk/integrations/aiohttp.py new file mode 100644 index 0000000000000000000000000000000000000000..46ee5f67b6e52b20d280370d01e18a6446c39559 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/aiohttp.py @@ -0,0 +1,362 @@ +import sys +import weakref +from functools import wraps + +import sentry_sdk +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA +from sentry_sdk.integrations import ( + _DEFAULT_FAILED_REQUEST_STATUS_CODES, + _check_minimum_version, + Integration, + DidNotEnable, +) +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.sessions import track_session +from sentry_sdk.integrations._wsgi_common import ( + _filter_headers, + request_body_within_bounds, +) +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SOURCE_FOR_STYLE, + TransactionSource, +) +from sentry_sdk.tracing_utils import should_propagate_trace, add_http_request_source +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + logger, + parse_url, + parse_version, + reraise, + transaction_from_function, + HAS_REAL_CONTEXTVARS, + CONTEXTVARS_ERROR_MESSAGE, + SENSITIVE_DATA_SUBSTITUTE, + AnnotatedValue, +) + +try: + import asyncio + + from aiohttp import __version__ as AIOHTTP_VERSION + from aiohttp import ClientSession, TraceConfig + from aiohttp.web import Application, HTTPException, UrlDispatcher +except ImportError: + raise DidNotEnable("AIOHTTP not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aiohttp.web_request import Request + from aiohttp.web_urldispatcher import UrlMappingMatchInfo + from aiohttp import TraceRequestStartParams, TraceRequestEndParams + + from collections.abc import Set + from types import SimpleNamespace + from typing import Any + from typing import Optional + from typing import Tuple + from typing import Union + + from sentry_sdk.utils import ExcInfo + from sentry_sdk._types import Event, EventProcessor + + +TRANSACTION_STYLE_VALUES = ("handler_name", "method_and_path_pattern") + + +class AioHttpIntegration(Integration): + identifier = "aiohttp" + origin = f"auto.http.{identifier}" + + def __init__( + self, + transaction_style: str = "handler_name", + *, + failed_request_status_codes: "Set[int]" = _DEFAULT_FAILED_REQUEST_STATUS_CODES, + ) -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + self._failed_request_status_codes = failed_request_status_codes + + @staticmethod + def setup_once() -> None: + version = parse_version(AIOHTTP_VERSION) + _check_minimum_version(AioHttpIntegration, version) + + if not HAS_REAL_CONTEXTVARS: + # We better have contextvars or we're going to leak state between + # requests. + raise DidNotEnable( + "The aiohttp integration for Sentry requires Python 3.7+ " + " or aiocontextvars package." + CONTEXTVARS_ERROR_MESSAGE + ) + + ignore_logger("aiohttp.server") + + old_handle = Application._handle + + async def sentry_app_handle( + self: "Any", request: "Request", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(AioHttpIntegration) + if integration is None: + return await old_handle(self, request, *args, **kwargs) + + weak_request = weakref.ref(request) + + with sentry_sdk.isolation_scope() as scope: + with track_session(scope, session_mode="request"): + # Scope data will not leak between requests because aiohttp + # create a task to wrap each request. + scope.generate_propagation_context() + scope.clear_breadcrumbs() + scope.add_event_processor(_make_request_processor(weak_request)) + + headers = dict(request.headers) + transaction = continue_trace( + headers, + op=OP.HTTP_SERVER, + # If this transaction name makes it to the UI, AIOHTTP's + # URL resolver did not find a route or died trying. + name="generic AIOHTTP request", + source=TransactionSource.ROUTE, + origin=AioHttpIntegration.origin, + ) + with sentry_sdk.start_transaction( + transaction, + custom_sampling_context={"aiohttp_request": request}, + ): + try: + response = await old_handle(self, request) + except HTTPException as e: + transaction.set_http_status(e.status_code) + + if ( + e.status_code + in integration._failed_request_status_codes + ): + _capture_exception() + + raise + except (asyncio.CancelledError, ConnectionResetError): + transaction.set_status(SPANSTATUS.CANCELLED) + raise + except Exception: + # This will probably map to a 500 but seems like we + # have no way to tell. Do not set span status. + reraise(*_capture_exception()) + + try: + # A valid response handler will return a valid response with a status. But, if the handler + # returns an invalid response (e.g. None), the line below will raise an AttributeError. + # Even though this is likely invalid, we need to handle this case to ensure we don't break + # the application. + response_status = response.status + except AttributeError: + pass + else: + transaction.set_http_status(response_status) + + return response + + Application._handle = sentry_app_handle + + old_urldispatcher_resolve = UrlDispatcher.resolve + + @wraps(old_urldispatcher_resolve) + async def sentry_urldispatcher_resolve( + self: "UrlDispatcher", request: "Request" + ) -> "UrlMappingMatchInfo": + rv = await old_urldispatcher_resolve(self, request) + + integration = sentry_sdk.get_client().get_integration(AioHttpIntegration) + if integration is None: + return rv + + name = None + + try: + if integration.transaction_style == "handler_name": + name = transaction_from_function(rv.handler) + elif integration.transaction_style == "method_and_path_pattern": + route_info = rv.get_info() + pattern = route_info.get("path") or route_info.get("formatter") + name = "{} {}".format(request.method, pattern) + except Exception: + pass + + if name is not None: + sentry_sdk.get_current_scope().set_transaction_name( + name, + source=SOURCE_FOR_STYLE[integration.transaction_style], + ) + + return rv + + UrlDispatcher.resolve = sentry_urldispatcher_resolve + + old_client_session_init = ClientSession.__init__ + + @ensure_integration_enabled(AioHttpIntegration, old_client_session_init) + def init(*args: "Any", **kwargs: "Any") -> None: + client_trace_configs = list(kwargs.get("trace_configs") or ()) + trace_config = create_trace_config() + client_trace_configs.append(trace_config) + + kwargs["trace_configs"] = client_trace_configs + return old_client_session_init(*args, **kwargs) + + ClientSession.__init__ = init + + +def create_trace_config() -> "TraceConfig": + async def on_request_start( + session: "ClientSession", + trace_config_ctx: "SimpleNamespace", + params: "TraceRequestStartParams", + ) -> None: + if sentry_sdk.get_client().get_integration(AioHttpIntegration) is None: + return + + method = params.method.upper() + + parsed_url = None + with capture_internal_exceptions(): + parsed_url = parse_url(str(params.url), sanitize=False) + + span = sentry_sdk.start_span( + op=OP.HTTP_CLIENT, + name="%s %s" + % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE), + origin=AioHttpIntegration.origin, + ) + span.set_data(SPANDATA.HTTP_METHOD, method) + if parsed_url is not None: + span.set_data("url", parsed_url.url) + span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query) + span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment) + + client = sentry_sdk.get_client() + + if should_propagate_trace(client, str(params.url)): + for ( + key, + value, + ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers( + span=span + ): + logger.debug( + "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format( + key=key, value=value, url=params.url + ) + ) + if key == BAGGAGE_HEADER_NAME and params.headers.get( + BAGGAGE_HEADER_NAME + ): + # do not overwrite any existing baggage, just append to it + params.headers[key] += "," + value + else: + params.headers[key] = value + + trace_config_ctx.span = span + + async def on_request_end( + session: "ClientSession", + trace_config_ctx: "SimpleNamespace", + params: "TraceRequestEndParams", + ) -> None: + if trace_config_ctx.span is None: + return + + span = trace_config_ctx.span + span.set_http_status(int(params.response.status)) + span.set_data("reason", params.response.reason) + span.finish() + + with capture_internal_exceptions(): + add_http_request_source(span) + + trace_config = TraceConfig() + + trace_config.on_request_start.append(on_request_start) + trace_config.on_request_end.append(on_request_end) + + return trace_config + + +def _make_request_processor( + weak_request: "weakref.ReferenceType[Request]", +) -> "EventProcessor": + def aiohttp_processor( + event: "Event", + hint: "dict[str, Tuple[type, BaseException, Any]]", + ) -> "Event": + request = weak_request() + if request is None: + return event + + with capture_internal_exceptions(): + request_info = event.setdefault("request", {}) + + request_info["url"] = "%s://%s%s" % ( + request.scheme, + request.host, + request.path, + ) + + request_info["query_string"] = request.query_string + request_info["method"] = request.method + request_info["env"] = {"REMOTE_ADDR": request.remote} + request_info["headers"] = _filter_headers(dict(request.headers)) + + # Just attach raw data here if it is within bounds, if available. + # Unfortunately there's no way to get structured data from aiohttp + # without awaiting on some coroutine. + request_info["data"] = get_aiohttp_request_data(request) + + return event + + return aiohttp_processor + + +def _capture_exception() -> "ExcInfo": + exc_info = sys.exc_info() + event, hint = event_from_exception( + exc_info, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "aiohttp", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + return exc_info + + +BODY_NOT_READ_MESSAGE = "[Can't show request body due to implementation details.]" + + +def get_aiohttp_request_data( + request: "Request", +) -> "Union[Optional[str], AnnotatedValue]": + bytes_body = request._read_bytes + + if bytes_body is not None: + # we have body to show + if not request_body_within_bounds(sentry_sdk.get_client(), len(bytes_body)): + return AnnotatedValue.removed_because_over_size_limit() + + encoding = request.charset or "utf-8" + return bytes_body.decode(encoding, "replace") + + if request.can_read_body: + # body exists but we can't show it + return BODY_NOT_READ_MESSAGE + + # request has no body + return None diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/anthropic.py b/lib/python3.12/site-packages/sentry_sdk/integrations/anthropic.py new file mode 100644 index 0000000000000000000000000000000000000000..b131e3381d0143c641642f02f224aaafaa94870f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/anthropic.py @@ -0,0 +1,607 @@ +import sys +import json +from collections.abc import Iterable +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import ( + GEN_AI_ALLOWED_MESSAGE_ROLES, + set_data_normalized, + normalize_message_roles, + truncate_and_annotate_messages, + get_start_span_function, + transform_anthropic_content_part, +) +from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import set_span_errored +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + package_version, + safe_serialize, + reraise, +) + +try: + try: + from anthropic import NotGiven + except ImportError: + NotGiven = None + + try: + from anthropic import Omit + except ImportError: + Omit = None + + from anthropic.resources import AsyncMessages, Messages + + if TYPE_CHECKING: + from anthropic.types import MessageStreamEvent, TextBlockParam +except ImportError: + raise DidNotEnable("Anthropic not installed") + +if TYPE_CHECKING: + from typing import Any, AsyncIterator, Iterator, List, Optional, Union + from sentry_sdk.tracing import Span + from sentry_sdk._types import TextPart + + +class _RecordedUsage: + output_tokens: int = 0 + input_tokens: int = 0 + cache_write_input_tokens: "Optional[int]" = 0 + cache_read_input_tokens: "Optional[int]" = 0 + + +class AnthropicIntegration(Integration): + identifier = "anthropic" + origin = f"auto.ai.{identifier}" + + def __init__(self: "AnthropicIntegration", include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + version = package_version("anthropic") + _check_minimum_version(AnthropicIntegration, version) + + Messages.create = _wrap_message_create(Messages.create) + AsyncMessages.create = _wrap_message_create_async(AsyncMessages.create) + + +def _capture_exception(exc: "Any") -> None: + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "anthropic", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _get_token_usage(result: "Messages") -> "tuple[int, int, int, int]": + """ + Get token usage from the Anthropic response. + Returns: (input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens) + """ + input_tokens = 0 + output_tokens = 0 + cache_read_input_tokens = 0 + cache_write_input_tokens = 0 + if hasattr(result, "usage"): + usage = result.usage + if hasattr(usage, "input_tokens") and isinstance(usage.input_tokens, int): + input_tokens = usage.input_tokens + if hasattr(usage, "output_tokens") and isinstance(usage.output_tokens, int): + output_tokens = usage.output_tokens + if hasattr(usage, "cache_read_input_tokens") and isinstance( + usage.cache_read_input_tokens, int + ): + cache_read_input_tokens = usage.cache_read_input_tokens + if hasattr(usage, "cache_creation_input_tokens") and isinstance( + usage.cache_creation_input_tokens, int + ): + cache_write_input_tokens = usage.cache_creation_input_tokens + + return ( + input_tokens, + output_tokens, + cache_read_input_tokens, + cache_write_input_tokens, + ) + + +def _collect_ai_data( + event: "MessageStreamEvent", + model: "str | None", + usage: "_RecordedUsage", + content_blocks: "list[str]", +) -> "tuple[str | None, _RecordedUsage, list[str]]": + """ + Collect model information, token usage, and collect content blocks from the AI streaming response. + """ + with capture_internal_exceptions(): + if hasattr(event, "type"): + if event.type == "content_block_start": + pass + elif event.type == "content_block_delta": + if hasattr(event.delta, "text"): + content_blocks.append(event.delta.text) + elif hasattr(event.delta, "partial_json"): + content_blocks.append(event.delta.partial_json) + elif event.type == "content_block_stop": + pass + + # Token counting logic mirrors anthropic SDK, which also extracts already accumulated tokens. + # https://github.com/anthropics/anthropic-sdk-python/blob/9c485f6966e10ae0ea9eabb3a921d2ea8145a25b/src/anthropic/lib/streaming/_messages.py#L433-L518 + if event.type == "message_start": + model = event.message.model or model + + incoming_usage = event.message.usage + usage.output_tokens = incoming_usage.output_tokens + usage.input_tokens = incoming_usage.input_tokens + + usage.cache_write_input_tokens = getattr( + incoming_usage, "cache_creation_input_tokens", None + ) + usage.cache_read_input_tokens = getattr( + incoming_usage, "cache_read_input_tokens", None + ) + + return ( + model, + usage, + content_blocks, + ) + + # Counterintuitive, but message_delta contains cumulative token counts :) + if event.type == "message_delta": + usage.output_tokens = event.usage.output_tokens + + # Update other usage fields if they exist in the event + input_tokens = getattr(event.usage, "input_tokens", None) + if input_tokens is not None: + usage.input_tokens = input_tokens + + cache_creation_input_tokens = getattr( + event.usage, "cache_creation_input_tokens", None + ) + if cache_creation_input_tokens is not None: + usage.cache_write_input_tokens = cache_creation_input_tokens + + cache_read_input_tokens = getattr( + event.usage, "cache_read_input_tokens", None + ) + if cache_read_input_tokens is not None: + usage.cache_read_input_tokens = cache_read_input_tokens + # TODO: Record event.usage.server_tool_use + + return ( + model, + usage, + content_blocks, + ) + + return ( + model, + usage, + content_blocks, + ) + + +def _transform_anthropic_content_block( + content_block: "dict[str, Any]", +) -> "dict[str, Any]": + """ + Transform an Anthropic content block using the Anthropic-specific transformer, + with special handling for Anthropic's text-type documents. + """ + # Handle Anthropic's text-type documents specially (not covered by shared function) + if content_block.get("type") == "document": + source = content_block.get("source") + if isinstance(source, dict) and source.get("type") == "text": + return { + "type": "text", + "text": source.get("data", ""), + } + + # Use Anthropic-specific transformation + result = transform_anthropic_content_part(content_block) + return result if result is not None else content_block + + +def _transform_system_instructions( + system_instructions: "Union[str, Iterable[TextBlockParam]]", +) -> "list[TextPart]": + if isinstance(system_instructions, str): + return [ + { + "type": "text", + "content": system_instructions, + } + ] + + return [ + { + "type": "text", + "content": instruction["text"], + } + for instruction in system_instructions + if isinstance(instruction, dict) and "text" in instruction + ] + + +def _set_input_data( + span: "Span", kwargs: "dict[str, Any]", integration: "AnthropicIntegration" +) -> None: + """ + Set input data for the span based on the provided keyword arguments for the anthropic message creation. + """ + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + system_instructions: "Union[str, Iterable[TextBlockParam]]" = kwargs.get("system") # type: ignore + messages = kwargs.get("messages") + if ( + messages is not None + and len(messages) > 0 + and should_send_default_pii() + and integration.include_prompts + ): + if isinstance(system_instructions, str) or isinstance( + system_instructions, Iterable + ): + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(_transform_system_instructions(system_instructions)), + ) + + normalized_messages = [] + for message in messages: + if ( + message.get("role") == GEN_AI_ALLOWED_MESSAGE_ROLES.USER + and "content" in message + and isinstance(message["content"], (list, tuple)) + ): + transformed_content = [] + for item in message["content"]: + # Skip tool_result items - they can contain images/documents + # with nested structures that are difficult to redact properly + if isinstance(item, dict) and item.get("type") == "tool_result": + continue + + # Transform content blocks (images, documents, etc.) + transformed_content.append( + _transform_anthropic_content_block(item) + if isinstance(item, dict) + else item + ) + + # If there are non-tool-result items, add them as a message + if transformed_content: + normalized_messages.append( + { + "role": message.get("role"), + "content": transformed_content, + } + ) + else: + # Transform content for non-list messages or assistant messages + transformed_message = message.copy() + if "content" in transformed_message: + content = transformed_message["content"] + if isinstance(content, (list, tuple)): + transformed_message["content"] = [ + _transform_anthropic_content_block(item) + if isinstance(item, dict) + else item + for item in content + ] + normalized_messages.append(transformed_message) + + role_normalized_messages = normalize_message_roles(normalized_messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + role_normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_STREAMING, kwargs.get("stream", False) + ) + + kwargs_keys_to_attributes = { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + } + for key, attribute in kwargs_keys_to_attributes.items(): + value = kwargs.get(key) + + if value is not None and _is_given(value): + set_data_normalized(span, attribute, value) + + # Input attributes: Tools + tools = kwargs.get("tools") + if tools is not None and _is_given(tools) and len(tools) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + + +def _set_output_data( + span: "Span", + integration: "AnthropicIntegration", + model: "str | None", + input_tokens: "int | None", + output_tokens: "int | None", + cache_read_input_tokens: "int | None", + cache_write_input_tokens: "int | None", + content_blocks: "list[Any]", + finish_span: bool = False, +) -> None: + """ + Set output data for the span based on the AI response.""" + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, model) + if should_send_default_pii() and integration.include_prompts: + output_messages: "dict[str, list[Any]]" = { + "response": [], + "tool": [], + } + + for output in content_blocks: + if output["type"] == "text": + output_messages["response"].append(output["text"]) + elif output["type"] == "tool_use": + output_messages["tool"].append(output) + + if len(output_messages["tool"]) > 0: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + output_messages["tool"], + unpack=False, + ) + + if len(output_messages["response"]) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"] + ) + + record_token_usage( + span, + input_tokens=input_tokens, + output_tokens=output_tokens, + input_tokens_cached=cache_read_input_tokens, + input_tokens_cache_write=cache_write_input_tokens, + ) + + if finish_span: + span.__exit__(None, None, None) + + +def _sentry_patched_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = kwargs.pop("integration") + if integration is None: + return f(*args, **kwargs) + + if "messages" not in kwargs: + return f(*args, **kwargs) + + try: + iter(kwargs["messages"]) + except TypeError: + return f(*args, **kwargs) + + model = kwargs.get("model", "") + + span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"chat {model}".strip(), + origin=AnthropicIntegration.origin, + ) + span.__enter__() + + _set_input_data(span, kwargs, integration) + + result = yield f, args, kwargs + + with capture_internal_exceptions(): + if hasattr(result, "content"): + ( + input_tokens, + output_tokens, + cache_read_input_tokens, + cache_write_input_tokens, + ) = _get_token_usage(result) + + content_blocks = [] + for content_block in result.content: + if hasattr(content_block, "to_dict"): + content_blocks.append(content_block.to_dict()) + elif hasattr(content_block, "model_dump"): + content_blocks.append(content_block.model_dump()) + elif hasattr(content_block, "text"): + content_blocks.append({"type": "text", "text": content_block.text}) + + _set_output_data( + span=span, + integration=integration, + model=getattr(result, "model", None), + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_read_input_tokens=cache_read_input_tokens, + cache_write_input_tokens=cache_write_input_tokens, + content_blocks=content_blocks, + finish_span=True, + ) + + # Streaming response + elif hasattr(result, "_iterator"): + old_iterator = result._iterator + + def new_iterator() -> "Iterator[MessageStreamEvent]": + model = None + usage = _RecordedUsage() + content_blocks: "list[str]" = [] + + for event in old_iterator: + ( + model, + usage, + content_blocks, + ) = _collect_ai_data( + event, + model, + usage, + content_blocks, + ) + yield event + + _set_output_data( + span=span, + integration=integration, + model=model, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + cache_read_input_tokens=usage.cache_read_input_tokens, + cache_write_input_tokens=usage.cache_write_input_tokens, + content_blocks=[{"text": "".join(content_blocks), "type": "text"}], + finish_span=True, + ) + + async def new_iterator_async() -> "AsyncIterator[MessageStreamEvent]": + model = None + usage = _RecordedUsage() + content_blocks: "list[str]" = [] + + async for event in old_iterator: + ( + model, + usage, + content_blocks, + ) = _collect_ai_data( + event, + model, + usage, + content_blocks, + ) + yield event + + _set_output_data( + span=span, + integration=integration, + model=model, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + cache_read_input_tokens=usage.cache_read_input_tokens, + cache_write_input_tokens=usage.cache_write_input_tokens, + content_blocks=[{"text": "".join(content_blocks), "type": "text"}], + finish_span=True, + ) + + if str(type(result._iterator)) == "": + result._iterator = new_iterator_async() + else: + result._iterator = new_iterator() + + else: + span.set_data("unknown_response", True) + span.__exit__(None, None, None) + + return result + + +def _wrap_message_create(f: "Any") -> "Any": + def _execute_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _sentry_patched_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return e.value + + try: + try: + result = f(*args, **kwargs) + except Exception as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(exc) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(AnthropicIntegration) + kwargs["integration"] = integration + + try: + return _execute_sync(f, *args, **kwargs) + finally: + span = sentry_sdk.get_current_span() + if span is not None and span.status == SPANSTATUS.INTERNAL_ERROR: + with capture_internal_exceptions(): + span.__exit__(None, None, None) + + return _sentry_patched_create_sync + + +def _wrap_message_create_async(f: "Any") -> "Any": + async def _execute_async(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _sentry_patched_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return await e.value + + try: + try: + result = await f(*args, **kwargs) + except Exception as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(exc) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + async def _sentry_patched_create_async(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(AnthropicIntegration) + kwargs["integration"] = integration + + try: + return await _execute_async(f, *args, **kwargs) + finally: + span = sentry_sdk.get_current_span() + if span is not None and span.status == SPANSTATUS.INTERNAL_ERROR: + with capture_internal_exceptions(): + span.__exit__(None, None, None) + + return _sentry_patched_create_async + + +def _is_given(obj: "Any") -> bool: + """ + Check for givenness safely across different anthropic versions. + """ + if NotGiven is not None and isinstance(obj, NotGiven): + return False + if Omit is not None and isinstance(obj, Omit): + return False + return True diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/argv.py b/lib/python3.12/site-packages/sentry_sdk/integrations/argv.py new file mode 100644 index 0000000000000000000000000000000000000000..b5b867c297f42b2c6ee9cec819f5dc7eba6fad24 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/argv.py @@ -0,0 +1,29 @@ +import sys + +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import add_global_event_processor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional + + from sentry_sdk._types import Event, Hint + + +class ArgvIntegration(Integration): + identifier = "argv" + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def processor(event: "Event", hint: "Optional[Hint]") -> "Optional[Event]": + if sentry_sdk.get_client().get_integration(ArgvIntegration) is not None: + extra = event.setdefault("extra", {}) + # If some event processor decided to set extra to e.g. an + # `int`, don't crash. Not here. + if isinstance(extra, dict): + extra["sys.argv"] = sys.argv + + return event diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/ariadne.py b/lib/python3.12/site-packages/sentry_sdk/integrations/ariadne.py new file mode 100644 index 0000000000000000000000000000000000000000..d353b62bea2d75175187645de45acbfbf3d02e6a --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/ariadne.py @@ -0,0 +1,160 @@ +from importlib import import_module + +import sentry_sdk +from sentry_sdk import get_client, capture_event +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.integrations._wsgi_common import request_body_within_bounds +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + package_version, +) + +try: + # importing like this is necessary due to name shadowing in ariadne + # (ariadne.graphql is also a function) + ariadne_graphql = import_module("ariadne.graphql") +except ImportError: + raise DidNotEnable("ariadne is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Dict, List, Optional + from ariadne.types import GraphQLError, GraphQLResult, GraphQLSchema, QueryParser # type: ignore + from graphql.language.ast import DocumentNode + from sentry_sdk._types import Event, EventProcessor + + +class AriadneIntegration(Integration): + identifier = "ariadne" + + @staticmethod + def setup_once() -> None: + version = package_version("ariadne") + _check_minimum_version(AriadneIntegration, version) + + ignore_logger("ariadne") + + _patch_graphql() + + +def _patch_graphql() -> None: + old_parse_query = ariadne_graphql.parse_query + old_handle_errors = ariadne_graphql.handle_graphql_errors + old_handle_query_result = ariadne_graphql.handle_query_result + + @ensure_integration_enabled(AriadneIntegration, old_parse_query) + def _sentry_patched_parse_query( + context_value: "Optional[Any]", + query_parser: "Optional[QueryParser]", + data: "Any", + ) -> "DocumentNode": + event_processor = _make_request_event_processor(data) + sentry_sdk.get_isolation_scope().add_event_processor(event_processor) + + result = old_parse_query(context_value, query_parser, data) + return result + + @ensure_integration_enabled(AriadneIntegration, old_handle_errors) + def _sentry_patched_handle_graphql_errors( + errors: "List[GraphQLError]", *args: "Any", **kwargs: "Any" + ) -> "GraphQLResult": + result = old_handle_errors(errors, *args, **kwargs) + + event_processor = _make_response_event_processor(result[1]) + sentry_sdk.get_isolation_scope().add_event_processor(event_processor) + + client = get_client() + if client.is_active(): + with capture_internal_exceptions(): + for error in errors: + event, hint = event_from_exception( + error, + client_options=client.options, + mechanism={ + "type": AriadneIntegration.identifier, + "handled": False, + }, + ) + capture_event(event, hint=hint) + + return result + + @ensure_integration_enabled(AriadneIntegration, old_handle_query_result) + def _sentry_patched_handle_query_result( + result: "Any", *args: "Any", **kwargs: "Any" + ) -> "GraphQLResult": + query_result = old_handle_query_result(result, *args, **kwargs) + + event_processor = _make_response_event_processor(query_result[1]) + sentry_sdk.get_isolation_scope().add_event_processor(event_processor) + + client = get_client() + if client.is_active(): + with capture_internal_exceptions(): + for error in result.errors or []: + event, hint = event_from_exception( + error, + client_options=client.options, + mechanism={ + "type": AriadneIntegration.identifier, + "handled": False, + }, + ) + capture_event(event, hint=hint) + + return query_result + + ariadne_graphql.parse_query = _sentry_patched_parse_query # type: ignore + ariadne_graphql.handle_graphql_errors = _sentry_patched_handle_graphql_errors # type: ignore + ariadne_graphql.handle_query_result = _sentry_patched_handle_query_result # type: ignore + + +def _make_request_event_processor(data: "GraphQLSchema") -> "EventProcessor": + """Add request data and api_target to events.""" + + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + if not isinstance(data, dict): + return event + + with capture_internal_exceptions(): + try: + content_length = int( + (data.get("headers") or {}).get("Content-Length", 0) + ) + except (TypeError, ValueError): + return event + + if should_send_default_pii() and request_body_within_bounds( + get_client(), content_length + ): + request_info = event.setdefault("request", {}) + request_info["api_target"] = "graphql" + request_info["data"] = data + + elif event.get("request", {}).get("data"): + del event["request"]["data"] + + return event + + return inner + + +def _make_response_event_processor(response: "Dict[str, Any]") -> "EventProcessor": + """Add response data to the event's response context.""" + + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + with capture_internal_exceptions(): + if should_send_default_pii() and response.get("errors"): + contexts = event.setdefault("contexts", {}) + contexts["response"] = { + "data": response, + } + + return event + + return inner diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/arq.py b/lib/python3.12/site-packages/sentry_sdk/integrations/arq.py new file mode 100644 index 0000000000000000000000000000000000000000..35c0a9c6f0780271522af061c0bb27c3232a87d1 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/arq.py @@ -0,0 +1,237 @@ +import sys + +import sentry_sdk +from sentry_sdk.consts import OP, SPANSTATUS +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import Transaction, TransactionSource +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + SENSITIVE_DATA_SUBSTITUTE, + parse_version, + reraise, +) + +try: + import arq.worker + from arq.version import VERSION as ARQ_VERSION + from arq.connections import ArqRedis + from arq.worker import JobExecutionFailed, Retry, RetryJob, Worker +except ImportError: + raise DidNotEnable("Arq is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Dict, Optional, Union + + from sentry_sdk._types import EventProcessor, Event, ExcInfo, Hint + + from arq.cron import CronJob + from arq.jobs import Job + from arq.typing import WorkerCoroutine + from arq.worker import Function + +ARQ_CONTROL_FLOW_EXCEPTIONS = (JobExecutionFailed, Retry, RetryJob) + + +class ArqIntegration(Integration): + identifier = "arq" + origin = f"auto.queue.{identifier}" + + @staticmethod + def setup_once() -> None: + try: + if isinstance(ARQ_VERSION, str): + version = parse_version(ARQ_VERSION) + else: + version = ARQ_VERSION.version[:2] + + except (TypeError, ValueError): + version = None + + _check_minimum_version(ArqIntegration, version) + + patch_enqueue_job() + patch_run_job() + patch_create_worker() + + ignore_logger("arq.worker") + + +def patch_enqueue_job() -> None: + old_enqueue_job = ArqRedis.enqueue_job + original_kwdefaults = old_enqueue_job.__kwdefaults__ + + async def _sentry_enqueue_job( + self: "ArqRedis", function: str, *args: "Any", **kwargs: "Any" + ) -> "Optional[Job]": + integration = sentry_sdk.get_client().get_integration(ArqIntegration) + if integration is None: + return await old_enqueue_job(self, function, *args, **kwargs) + + with sentry_sdk.start_span( + op=OP.QUEUE_SUBMIT_ARQ, name=function, origin=ArqIntegration.origin + ): + return await old_enqueue_job(self, function, *args, **kwargs) + + _sentry_enqueue_job.__kwdefaults__ = original_kwdefaults + ArqRedis.enqueue_job = _sentry_enqueue_job + + +def patch_run_job() -> None: + old_run_job = Worker.run_job + + async def _sentry_run_job(self: "Worker", job_id: str, score: int) -> None: + integration = sentry_sdk.get_client().get_integration(ArqIntegration) + if integration is None: + return await old_run_job(self, job_id, score) + + with sentry_sdk.isolation_scope() as scope: + scope._name = "arq" + scope.clear_breadcrumbs() + + transaction = Transaction( + name="unknown arq task", + status="ok", + op=OP.QUEUE_TASK_ARQ, + source=TransactionSource.TASK, + origin=ArqIntegration.origin, + ) + + with sentry_sdk.start_transaction(transaction): + return await old_run_job(self, job_id, score) + + Worker.run_job = _sentry_run_job + + +def _capture_exception(exc_info: "ExcInfo") -> None: + scope = sentry_sdk.get_current_scope() + + if scope.transaction is not None: + if exc_info[0] in ARQ_CONTROL_FLOW_EXCEPTIONS: + scope.transaction.set_status(SPANSTATUS.ABORTED) + return + + scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR) + + event, hint = event_from_exception( + exc_info, + client_options=sentry_sdk.get_client().options, + mechanism={"type": ArqIntegration.identifier, "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _make_event_processor( + ctx: "Dict[Any, Any]", *args: "Any", **kwargs: "Any" +) -> "EventProcessor": + def event_processor(event: "Event", hint: "Hint") -> "Optional[Event]": + with capture_internal_exceptions(): + scope = sentry_sdk.get_current_scope() + if scope.transaction is not None: + scope.transaction.name = ctx["job_name"] + event["transaction"] = ctx["job_name"] + + tags = event.setdefault("tags", {}) + tags["arq_task_id"] = ctx["job_id"] + tags["arq_task_retry"] = ctx["job_try"] > 1 + extra = event.setdefault("extra", {}) + extra["arq-job"] = { + "task": ctx["job_name"], + "args": ( + args if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE + ), + "kwargs": ( + kwargs if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE + ), + "retry": ctx["job_try"], + } + + return event + + return event_processor + + +def _wrap_coroutine(name: str, coroutine: "WorkerCoroutine") -> "WorkerCoroutine": + async def _sentry_coroutine( + ctx: "Dict[Any, Any]", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(ArqIntegration) + if integration is None: + return await coroutine(ctx, *args, **kwargs) + + sentry_sdk.get_isolation_scope().add_event_processor( + _make_event_processor({**ctx, "job_name": name}, *args, **kwargs) + ) + + try: + result = await coroutine(ctx, *args, **kwargs) + except Exception: + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result + + return _sentry_coroutine + + +def patch_create_worker() -> None: + old_create_worker = arq.worker.create_worker + + @ensure_integration_enabled(ArqIntegration, old_create_worker) + def _sentry_create_worker(*args: "Any", **kwargs: "Any") -> "Worker": + settings_cls = args[0] if args else kwargs.get("settings_cls") + + if isinstance(settings_cls, dict): + if "functions" in settings_cls: + settings_cls["functions"] = [ + _get_arq_function(func) + for func in settings_cls.get("functions", []) + ] + if "cron_jobs" in settings_cls: + settings_cls["cron_jobs"] = [ + _get_arq_cron_job(cron_job) + for cron_job in settings_cls.get("cron_jobs", []) + ] + + if hasattr(settings_cls, "functions"): + settings_cls.functions = [ # type: ignore[union-attr] + _get_arq_function(func) + for func in settings_cls.functions # type: ignore[union-attr] + ] + if hasattr(settings_cls, "cron_jobs"): + settings_cls.cron_jobs = [ # type: ignore[union-attr] + _get_arq_cron_job(cron_job) + for cron_job in (settings_cls.cron_jobs or []) # type: ignore[union-attr] + ] + + if "functions" in kwargs: + kwargs["functions"] = [ + _get_arq_function(func) for func in kwargs.get("functions", []) + ] + if "cron_jobs" in kwargs: + kwargs["cron_jobs"] = [ + _get_arq_cron_job(cron_job) for cron_job in kwargs.get("cron_jobs", []) + ] + + return old_create_worker(*args, **kwargs) + + arq.worker.create_worker = _sentry_create_worker + + +def _get_arq_function(func: "Union[str, Function, WorkerCoroutine]") -> "Function": + arq_func = arq.worker.func(func) + arq_func.coroutine = _wrap_coroutine(arq_func.name, arq_func.coroutine) + + return arq_func + + +def _get_arq_cron_job(cron_job: "CronJob") -> "CronJob": + cron_job.coroutine = _wrap_coroutine(cron_job.name, cron_job.coroutine) + + return cron_job diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/asgi.py b/lib/python3.12/site-packages/sentry_sdk/integrations/asgi.py new file mode 100644 index 0000000000000000000000000000000000000000..6983af89ede9e633be6a254bde89cf9589dfff87 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/asgi.py @@ -0,0 +1,336 @@ +""" +An ASGI middleware. + +Based on Tom Christie's `sentry-asgi `. +""" + +import asyncio +import inspect +from copy import deepcopy +from functools import partial + +import sentry_sdk +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.integrations._asgi_common import ( + _get_headers, + _get_request_data, + _get_url, +) +from sentry_sdk.integrations._wsgi_common import ( + DEFAULT_HTTP_METHODS_TO_CAPTURE, + nullcontext, +) +from sentry_sdk.sessions import track_session +from sentry_sdk.tracing import ( + SOURCE_FOR_STYLE, + TransactionSource, +) +from sentry_sdk.utils import ( + ContextVar, + event_from_exception, + HAS_REAL_CONTEXTVARS, + CONTEXTVARS_ERROR_MESSAGE, + logger, + transaction_from_function, + _get_installed_modules, +) +from sentry_sdk.tracing import Transaction + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Optional + from typing import Tuple + + from sentry_sdk._types import Event, Hint + + +_asgi_middleware_applied = ContextVar("sentry_asgi_middleware_applied") + +_DEFAULT_TRANSACTION_NAME = "generic ASGI request" + +TRANSACTION_STYLE_VALUES = ("endpoint", "url") + + +def _capture_exception(exc: "Any", mechanism_type: str = "asgi") -> None: + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": mechanism_type, "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _looks_like_asgi3(app: "Any") -> bool: + """ + Try to figure out if an application object supports ASGI3. + + This is how uvicorn figures out the application version as well. + """ + if inspect.isclass(app): + return hasattr(app, "__await__") + elif inspect.isfunction(app): + return asyncio.iscoroutinefunction(app) + else: + call = getattr(app, "__call__", None) # noqa + return asyncio.iscoroutinefunction(call) + + +class SentryAsgiMiddleware: + __slots__ = ( + "app", + "__call__", + "transaction_style", + "mechanism_type", + "span_origin", + "http_methods_to_capture", + ) + + def __init__( + self, + app: "Any", + unsafe_context_data: bool = False, + transaction_style: str = "endpoint", + mechanism_type: str = "asgi", + span_origin: str = "manual", + http_methods_to_capture: "Tuple[str, ...]" = DEFAULT_HTTP_METHODS_TO_CAPTURE, + asgi_version: "Optional[int]" = None, + ) -> None: + """ + Instrument an ASGI application with Sentry. Provides HTTP/websocket + data to sent events and basic handling for exceptions bubbling up + through the middleware. + + :param unsafe_context_data: Disable errors when a proper contextvars installation could not be found. We do not recommend changing this from the default. + """ + if not unsafe_context_data and not HAS_REAL_CONTEXTVARS: + # We better have contextvars or we're going to leak state between + # requests. + raise RuntimeError( + "The ASGI middleware for Sentry requires Python 3.7+ " + "or the aiocontextvars package." + CONTEXTVARS_ERROR_MESSAGE + ) + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + + asgi_middleware_while_using_starlette_or_fastapi = ( + mechanism_type == "asgi" and "starlette" in _get_installed_modules() + ) + if asgi_middleware_while_using_starlette_or_fastapi: + logger.warning( + "The Sentry Python SDK can now automatically support ASGI frameworks like Starlette and FastAPI. " + "Please remove 'SentryAsgiMiddleware' from your project. " + "See https://docs.sentry.io/platforms/python/guides/asgi/ for more information." + ) + + self.transaction_style = transaction_style + self.mechanism_type = mechanism_type + self.span_origin = span_origin + self.app = app + self.http_methods_to_capture = http_methods_to_capture + + if asgi_version is None: + if _looks_like_asgi3(app): + asgi_version = 3 + else: + asgi_version = 2 + + if asgi_version == 3: + self.__call__ = self._run_asgi3 + elif asgi_version == 2: + self.__call__ = self._run_asgi2 # type: ignore + + def _capture_lifespan_exception(self, exc: Exception) -> None: + """Capture exceptions raise in application lifespan handlers. + + The separate function is needed to support overriding in derived integrations that use different catching mechanisms. + """ + return _capture_exception(exc=exc, mechanism_type=self.mechanism_type) + + def _capture_request_exception(self, exc: Exception) -> None: + """Capture exceptions raised in incoming request handlers. + + The separate function is needed to support overriding in derived integrations that use different catching mechanisms. + """ + return _capture_exception(exc=exc, mechanism_type=self.mechanism_type) + + def _run_asgi2(self, scope: "Any") -> "Any": + async def inner(receive: "Any", send: "Any") -> "Any": + return await self._run_app(scope, receive, send, asgi_version=2) + + return inner + + async def _run_asgi3(self, scope: "Any", receive: "Any", send: "Any") -> "Any": + return await self._run_app(scope, receive, send, asgi_version=3) + + async def _run_app( + self, scope: "Any", receive: "Any", send: "Any", asgi_version: int + ) -> "Any": + is_recursive_asgi_middleware = _asgi_middleware_applied.get(False) + is_lifespan = scope["type"] == "lifespan" + if is_recursive_asgi_middleware or is_lifespan: + try: + if asgi_version == 2: + return await self.app(scope)(receive, send) + else: + return await self.app(scope, receive, send) + + except Exception as exc: + self._capture_lifespan_exception(exc) + raise exc from None + + _asgi_middleware_applied.set(True) + try: + with sentry_sdk.isolation_scope() as sentry_scope: + with track_session(sentry_scope, session_mode="request"): + sentry_scope.clear_breadcrumbs() + sentry_scope._name = "asgi" + processor = partial(self.event_processor, asgi_scope=scope) + sentry_scope.add_event_processor(processor) + + ty = scope["type"] + ( + transaction_name, + transaction_source, + ) = self._get_transaction_name_and_source( + self.transaction_style, + scope, + ) + + method = scope.get("method", "").upper() + transaction = None + if ty in ("http", "websocket"): + if ty == "websocket" or method in self.http_methods_to_capture: + transaction = continue_trace( + _get_headers(scope), + op="{}.server".format(ty), + name=transaction_name, + source=transaction_source, + origin=self.span_origin, + ) + else: + transaction = Transaction( + op=OP.HTTP_SERVER, + name=transaction_name, + source=transaction_source, + origin=self.span_origin, + ) + + if transaction: + transaction.set_tag("asgi.type", ty) + + transaction_context = ( + sentry_sdk.start_transaction( + transaction, + custom_sampling_context={"asgi_scope": scope}, + ) + if transaction is not None + else nullcontext() + ) + with transaction_context: + try: + + async def _sentry_wrapped_send( + event: "Dict[str, Any]", + ) -> "Any": + if transaction is not None: + is_http_response = ( + event.get("type") == "http.response.start" + and "status" in event + ) + if is_http_response: + transaction.set_http_status(event["status"]) + + return await send(event) + + if asgi_version == 2: + return await self.app(scope)( + receive, _sentry_wrapped_send + ) + else: + return await self.app( + scope, receive, _sentry_wrapped_send + ) + except Exception as exc: + self._capture_request_exception(exc) + raise exc from None + finally: + _asgi_middleware_applied.set(False) + + def event_processor( + self, event: "Event", hint: "Hint", asgi_scope: "Any" + ) -> "Optional[Event]": + request_data = event.get("request", {}) + request_data.update(_get_request_data(asgi_scope)) + event["request"] = deepcopy(request_data) + + # Only set transaction name if not already set by Starlette or FastAPI (or other frameworks) + transaction = event.get("transaction") + transaction_source = (event.get("transaction_info") or {}).get("source") + already_set = ( + transaction is not None + and transaction != _DEFAULT_TRANSACTION_NAME + and transaction_source + in [ + TransactionSource.COMPONENT, + TransactionSource.ROUTE, + TransactionSource.CUSTOM, + ] + ) + if not already_set: + name, source = self._get_transaction_name_and_source( + self.transaction_style, asgi_scope + ) + event["transaction"] = name + event["transaction_info"] = {"source": source} + + return event + + # Helper functions. + # + # Note: Those functions are not public API. If you want to mutate request + # data to your liking it's recommended to use the `before_send` callback + # for that. + + def _get_transaction_name_and_source( + self: "SentryAsgiMiddleware", transaction_style: str, asgi_scope: "Any" + ) -> "Tuple[str, str]": + name = None + source = SOURCE_FOR_STYLE[transaction_style] + ty = asgi_scope.get("type") + + if transaction_style == "endpoint": + endpoint = asgi_scope.get("endpoint") + # Webframeworks like Starlette mutate the ASGI env once routing is + # done, which is sometime after the request has started. If we have + # an endpoint, overwrite our generic transaction name. + if endpoint: + name = transaction_from_function(endpoint) or "" + else: + name = _get_url(asgi_scope, "http" if ty == "http" else "ws", host=None) + source = TransactionSource.URL + + elif transaction_style == "url": + # FastAPI includes the route object in the scope to let Sentry extract the + # path from it for the transaction name + route = asgi_scope.get("route") + if route: + path = getattr(route, "path", None) + if path is not None: + name = path + else: + name = _get_url(asgi_scope, "http" if ty == "http" else "ws", host=None) + source = TransactionSource.URL + + if name is None: + name = _DEFAULT_TRANSACTION_NAME + source = TransactionSource.ROUTE + return name, source + + return name, source diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/asyncio.py b/lib/python3.12/site-packages/sentry_sdk/integrations/asyncio.py new file mode 100644 index 0000000000000000000000000000000000000000..b7aa0a720270de80d2675111696c4698bb0c2530 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/asyncio.py @@ -0,0 +1,203 @@ +import sys +import functools + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import nullcontext +from sentry_sdk.utils import event_from_exception, logger, reraise + +try: + import asyncio + from asyncio.tasks import Task +except ImportError: + raise DidNotEnable("asyncio not available") + +from typing import cast, TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, TypeVar + from collections.abc import Coroutine + + from sentry_sdk._types import ExcInfo + + T = TypeVar("T", bound=Callable[..., Any]) + + +def get_name(coro: "Any") -> str: + return ( + getattr(coro, "__qualname__", None) + or getattr(coro, "__name__", None) + or "coroutine without __name__" + ) + + +def _wrap_coroutine(wrapped: "Coroutine[Any, Any, Any]") -> "Callable[[T], T]": + # Only __name__ and __qualname__ are copied from function to coroutine in CPython + return functools.partial( + functools.update_wrapper, + wrapped=wrapped, # type: ignore + assigned=("__name__", "__qualname__"), + updated=(), + ) + + +def patch_asyncio() -> None: + orig_task_factory = None + try: + loop = asyncio.get_running_loop() + orig_task_factory = loop.get_task_factory() + + # Check if already patched + if getattr(orig_task_factory, "_is_sentry_task_factory", False): + return + + def _sentry_task_factory( + loop: "asyncio.AbstractEventLoop", + coro: "Coroutine[Any, Any, Any]", + **kwargs: "Any", + ) -> "asyncio.Future[Any]": + @_wrap_coroutine(coro) + async def _task_with_sentry_span_creation() -> "Any": + result = None + + integration = sentry_sdk.get_client().get_integration( + AsyncioIntegration + ) + task_spans = integration.task_spans if integration else False + + with sentry_sdk.isolation_scope(): + with ( + sentry_sdk.start_span( + op=OP.FUNCTION, + name=get_name(coro), + origin=AsyncioIntegration.origin, + ) + if task_spans + else nullcontext() + ): + try: + result = await coro + except StopAsyncIteration as e: + raise e from None + except Exception: + reraise(*_capture_exception()) + + return result + + task = None + + # Trying to use user set task factory (if there is one) + if orig_task_factory: + task = orig_task_factory( + loop, _task_with_sentry_span_creation(), **kwargs + ) + + if task is None: + # The default task factory in `asyncio` does not have its own function + # but is just a couple of lines in `asyncio.base_events.create_task()` + # Those lines are copied here. + + # WARNING: + # If the default behavior of the task creation in asyncio changes, + # this will break! + task = Task(_task_with_sentry_span_creation(), loop=loop, **kwargs) + if task._source_traceback: # type: ignore + del task._source_traceback[-1] # type: ignore + + # Set the task name to include the original coroutine's name + try: + cast("asyncio.Task[Any]", task).set_name( + f"{get_name(coro)} (Sentry-wrapped)" + ) + except AttributeError: + # set_name might not be available in all Python versions + pass + + return task + + _sentry_task_factory._is_sentry_task_factory = True # type: ignore + loop.set_task_factory(_sentry_task_factory) # type: ignore + + except RuntimeError: + # When there is no running loop, we have nothing to patch. + logger.warning( + "There is no running asyncio loop so there is nothing Sentry can patch. " + "Please make sure you call sentry_sdk.init() within a running " + "asyncio loop for the AsyncioIntegration to work. " + "See https://docs.sentry.io/platforms/python/integrations/asyncio/" + ) + + +def _capture_exception() -> "ExcInfo": + exc_info = sys.exc_info() + + client = sentry_sdk.get_client() + + integration = client.get_integration(AsyncioIntegration) + if integration is not None: + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "asyncio", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return exc_info + + +class AsyncioIntegration(Integration): + identifier = "asyncio" + origin = f"auto.function.{identifier}" + + def __init__(self, task_spans: bool = True) -> None: + self.task_spans = task_spans + + @staticmethod + def setup_once() -> None: + patch_asyncio() + + +def enable_asyncio_integration(*args: "Any", **kwargs: "Any") -> None: + """ + Enable AsyncioIntegration with the provided options. + + This is useful in scenarios where Sentry needs to be initialized before + an event loop is set up, but you still want to instrument asyncio once there + is an event loop. In that case, you can sentry_sdk.init() early on without + the AsyncioIntegration and then, once the event loop has been set up, + execute: + + ```python + from sentry_sdk.integrations.asyncio import enable_asyncio_integration + + async def async_entrypoint(): + enable_asyncio_integration() + ``` + + Any arguments provided will be passed to AsyncioIntegration() as is. + + If AsyncioIntegration has already patched the current event loop, this + function won't have any effect. + + If AsyncioIntegration was provided in + sentry_sdk.init(disabled_integrations=[...]), this function will ignore that + and the integration will be enabled. + """ + client = sentry_sdk.get_client() + if not client.is_active(): + return + + # This function purposefully bypasses the integration machinery in + # integrations/__init__.py. _installed_integrations/_processed_integrations + # is used to prevent double patching the same module, but in the case of + # the AsyncioIntegration, we don't monkeypatch the standard library directly, + # we patch the currently running event loop, and we keep the record of doing + # that on the loop itself. + logger.debug("Setting up integration asyncio") + + integration = AsyncioIntegration(*args, **kwargs) + integration.setup_once() + + if "asyncio" not in client.integrations: + client.integrations["asyncio"] = integration diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/asyncpg.py b/lib/python3.12/site-packages/sentry_sdk/integrations/asyncpg.py new file mode 100644 index 0000000000000000000000000000000000000000..7f3591154a3943d353350e189294de9da0a32211 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/asyncpg.py @@ -0,0 +1,210 @@ +from __future__ import annotations +import contextlib +from typing import Any, TypeVar, Callable, Awaitable, Iterator + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.tracing import Span +from sentry_sdk.tracing_utils import add_query_source, record_sql_queries +from sentry_sdk.utils import ( + ensure_integration_enabled, + parse_version, + capture_internal_exceptions, +) + +try: + import asyncpg # type: ignore[import-not-found] + from asyncpg.cursor import BaseCursor # type: ignore + +except ImportError: + raise DidNotEnable("asyncpg not installed.") + + +class AsyncPGIntegration(Integration): + identifier = "asyncpg" + origin = f"auto.db.{identifier}" + _record_params = False + + def __init__(self, *, record_params: bool = False): + AsyncPGIntegration._record_params = record_params + + @staticmethod + def setup_once() -> None: + # asyncpg.__version__ is a string containing the semantic version in the form of ".." + asyncpg_version = parse_version(asyncpg.__version__) + _check_minimum_version(AsyncPGIntegration, asyncpg_version) + + asyncpg.Connection.execute = _wrap_execute( + asyncpg.Connection.execute, + ) + + asyncpg.Connection._execute = _wrap_connection_method( + asyncpg.Connection._execute + ) + asyncpg.Connection._executemany = _wrap_connection_method( + asyncpg.Connection._executemany, executemany=True + ) + asyncpg.Connection.cursor = _wrap_cursor_creation(asyncpg.Connection.cursor) + asyncpg.Connection.prepare = _wrap_connection_method(asyncpg.Connection.prepare) + asyncpg.connect_utils._connect_addr = _wrap_connect_addr( + asyncpg.connect_utils._connect_addr + ) + + +T = TypeVar("T") + + +def _wrap_execute(f: "Callable[..., Awaitable[T]]") -> "Callable[..., Awaitable[T]]": + async def _inner(*args: "Any", **kwargs: "Any") -> "T": + if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None: + return await f(*args, **kwargs) + + # Avoid recording calls to _execute twice. + # Calls to Connection.execute with args also call + # Connection._execute, which is recorded separately + # args[0] = the connection object, args[1] is the query + if len(args) > 2: + return await f(*args, **kwargs) + + query = args[1] + with record_sql_queries( + cursor=None, + query=query, + params_list=None, + paramstyle=None, + executemany=False, + span_origin=AsyncPGIntegration.origin, + ) as span: + res = await f(*args, **kwargs) + + with capture_internal_exceptions(): + add_query_source(span) + + return res + + return _inner + + +SubCursor = TypeVar("SubCursor", bound=BaseCursor) + + +@contextlib.contextmanager +def _record( + cursor: "SubCursor | None", + query: str, + params_list: "tuple[Any, ...] | None", + *, + executemany: bool = False, +) -> "Iterator[Span]": + integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration) + if integration is not None and not integration._record_params: + params_list = None + + param_style = "pyformat" if params_list else None + + with record_sql_queries( + cursor=cursor, + query=query, + params_list=params_list, + paramstyle=param_style, + executemany=executemany, + record_cursor_repr=cursor is not None, + span_origin=AsyncPGIntegration.origin, + ) as span: + yield span + + +def _wrap_connection_method( + f: "Callable[..., Awaitable[T]]", *, executemany: bool = False +) -> "Callable[..., Awaitable[T]]": + async def _inner(*args: "Any", **kwargs: "Any") -> "T": + if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None: + return await f(*args, **kwargs) + query = args[1] + params_list = args[2] if len(args) > 2 else None + with _record(None, query, params_list, executemany=executemany) as span: + _set_db_data(span, args[0]) + res = await f(*args, **kwargs) + + return res + + return _inner + + +def _wrap_cursor_creation(f: "Callable[..., T]") -> "Callable[..., T]": + @ensure_integration_enabled(AsyncPGIntegration, f) + def _inner(*args: "Any", **kwargs: "Any") -> "T": # noqa: N807 + query = args[1] + params_list = args[2] if len(args) > 2 else None + + with _record( + None, + query, + params_list, + executemany=False, + ) as span: + _set_db_data(span, args[0]) + res = f(*args, **kwargs) + span.set_data("db.cursor", res) + + return res + + return _inner + + +def _wrap_connect_addr( + f: "Callable[..., Awaitable[T]]", +) -> "Callable[..., Awaitable[T]]": + async def _inner(*args: "Any", **kwargs: "Any") -> "T": + if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None: + return await f(*args, **kwargs) + + user = kwargs["params"].user + database = kwargs["params"].database + + with sentry_sdk.start_span( + op=OP.DB, + name="connect", + origin=AsyncPGIntegration.origin, + ) as span: + span.set_data(SPANDATA.DB_SYSTEM, "postgresql") + addr = kwargs.get("addr") + if addr: + try: + span.set_data(SPANDATA.SERVER_ADDRESS, addr[0]) + span.set_data(SPANDATA.SERVER_PORT, addr[1]) + except IndexError: + pass + span.set_data(SPANDATA.DB_NAME, database) + span.set_data(SPANDATA.DB_USER, user) + + with capture_internal_exceptions(): + sentry_sdk.add_breadcrumb( + message="connect", category="query", data=span._data + ) + res = await f(*args, **kwargs) + + return res + + return _inner + + +def _set_db_data(span: "Span", conn: "Any") -> None: + span.set_data(SPANDATA.DB_SYSTEM, "postgresql") + + addr = conn._addr + if addr: + try: + span.set_data(SPANDATA.SERVER_ADDRESS, addr[0]) + span.set_data(SPANDATA.SERVER_PORT, addr[1]) + except IndexError: + pass + + database = conn._params.database + if database: + span.set_data(SPANDATA.DB_NAME, database) + + user = conn._params.user + if user: + span.set_data(SPANDATA.DB_USER, user) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/atexit.py b/lib/python3.12/site-packages/sentry_sdk/integrations/atexit.py new file mode 100644 index 0000000000000000000000000000000000000000..efa4c74af01f665ad479afb8d75797d0af63c4cc --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/atexit.py @@ -0,0 +1,52 @@ +import os +import sys +import atexit + +import sentry_sdk +from sentry_sdk.utils import logger +from sentry_sdk.integrations import Integration +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Optional + + +def default_callback(pending: int, timeout: int) -> None: + """This is the default shutdown callback that is set on the options. + It prints out a message to stderr that informs the user that some events + are still pending and the process is waiting for them to flush out. + """ + + def echo(msg: str) -> None: + sys.stderr.write(msg + "\n") + + echo("Sentry is attempting to send %i pending events" % pending) + echo("Waiting up to %s seconds" % timeout) + echo("Press Ctrl-%s to quit" % (os.name == "nt" and "Break" or "C")) + sys.stderr.flush() + + +class AtexitIntegration(Integration): + identifier = "atexit" + + def __init__(self, callback: "Optional[Any]" = None) -> None: + if callback is None: + callback = default_callback + self.callback = callback + + @staticmethod + def setup_once() -> None: + @atexit.register + def _shutdown() -> None: + client = sentry_sdk.get_client() + integration = client.get_integration(AtexitIntegration) + + if integration is None: + return + + logger.debug("atexit: got shutdown signal") + logger.debug("atexit: shutting down client") + sentry_sdk.get_isolation_scope().end_session() + + client.close(callback=integration.callback) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/aws_lambda.py b/lib/python3.12/site-packages/sentry_sdk/integrations/aws_lambda.py new file mode 100644 index 0000000000000000000000000000000000000000..22893313ae2983148d6b4ecf5b722dbad4a4e637 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/aws_lambda.py @@ -0,0 +1,489 @@ +import functools +import json +import re +import sys +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from os import environ + +import sentry_sdk +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + AnnotatedValue, + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + logger, + TimeoutThread, + reraise, +) +from sentry_sdk.integrations import Integration +from sentry_sdk.integrations._wsgi_common import _filter_headers + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import TypeVar + from typing import Callable + from typing import Optional + + from sentry_sdk._types import EventProcessor, Event, Hint + + F = TypeVar("F", bound=Callable[..., Any]) + +# Constants +TIMEOUT_WARNING_BUFFER = 1500 # Buffer time required to send timeout warning to Sentry +MILLIS_TO_SECONDS = 1000.0 + + +def _wrap_init_error(init_error: "F") -> "F": + @ensure_integration_enabled(AwsLambdaIntegration, init_error) + def sentry_init_error(*args: "Any", **kwargs: "Any") -> "Any": + client = sentry_sdk.get_client() + + with capture_internal_exceptions(): + sentry_sdk.get_isolation_scope().clear_breadcrumbs() + + exc_info = sys.exc_info() + if exc_info and all(exc_info): + sentry_event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "aws_lambda", "handled": False}, + ) + sentry_sdk.capture_event(sentry_event, hint=hint) + + else: + # Fall back to AWS lambdas JSON representation of the error + error_info = args[1] + if isinstance(error_info, str): + error_info = json.loads(error_info) + sentry_event = _event_from_error_json(error_info) + sentry_sdk.capture_event(sentry_event) + + return init_error(*args, **kwargs) + + return sentry_init_error # type: ignore + + +def _wrap_handler(handler: "F") -> "F": + @functools.wraps(handler) + def sentry_handler( + aws_event: "Any", aws_context: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + # Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html, + # `event` here is *likely* a dictionary, but also might be a number of + # other types (str, int, float, None). + # + # In some cases, it is a list (if the user is batch-invoking their + # function, for example), in which case we'll use the first entry as a + # representative from which to try pulling request data. (Presumably it + # will be the same for all events in the list, since they're all hitting + # the lambda in the same request.) + + client = sentry_sdk.get_client() + integration = client.get_integration(AwsLambdaIntegration) + + if integration is None: + return handler(aws_event, aws_context, *args, **kwargs) + + if isinstance(aws_event, list) and len(aws_event) >= 1: + request_data = aws_event[0] + batch_size = len(aws_event) + else: + request_data = aws_event + batch_size = 1 + + if not isinstance(request_data, dict): + # If we're not dealing with a dictionary, we won't be able to get + # headers, path, http method, etc in any case, so it's fine that + # this is empty + request_data = {} + + configured_time = aws_context.get_remaining_time_in_millis() + + with sentry_sdk.isolation_scope() as scope: + timeout_thread = None + with capture_internal_exceptions(): + scope.clear_breadcrumbs() + scope.add_event_processor( + _make_request_event_processor( + request_data, aws_context, configured_time + ) + ) + scope.set_tag( + "aws_region", aws_context.invoked_function_arn.split(":")[3] + ) + if batch_size > 1: + scope.set_tag("batch_request", True) + scope.set_tag("batch_size", batch_size) + + # Starting the Timeout thread only if the configured time is greater than Timeout warning + # buffer and timeout_warning parameter is set True. + if ( + integration.timeout_warning + and configured_time > TIMEOUT_WARNING_BUFFER + ): + waiting_time = ( + configured_time - TIMEOUT_WARNING_BUFFER + ) / MILLIS_TO_SECONDS + + timeout_thread = TimeoutThread( + waiting_time, + configured_time / MILLIS_TO_SECONDS, + isolation_scope=scope, + current_scope=sentry_sdk.get_current_scope(), + ) + + # Starting the thread to raise timeout warning exception + timeout_thread.start() + + headers = request_data.get("headers", {}) + # Some AWS Services (ie. EventBridge) set headers as a list + # or None, so we must ensure it is a dict + if not isinstance(headers, dict): + headers = {} + + transaction = continue_trace( + headers, + op=OP.FUNCTION_AWS, + name=aws_context.function_name, + source=TransactionSource.COMPONENT, + origin=AwsLambdaIntegration.origin, + ) + with sentry_sdk.start_transaction( + transaction, + custom_sampling_context={ + "aws_event": aws_event, + "aws_context": aws_context, + }, + ): + try: + return handler(aws_event, aws_context, *args, **kwargs) + except Exception: + exc_info = sys.exc_info() + sentry_event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "aws_lambda", "handled": False}, + ) + sentry_sdk.capture_event(sentry_event, hint=hint) + reraise(*exc_info) + finally: + if timeout_thread: + timeout_thread.stop() + + return sentry_handler # type: ignore + + +def _drain_queue() -> None: + with capture_internal_exceptions(): + client = sentry_sdk.get_client() + integration = client.get_integration(AwsLambdaIntegration) + if integration is not None: + # Flush out the event queue before AWS kills the + # process. + client.flush() + + +class AwsLambdaIntegration(Integration): + identifier = "aws_lambda" + origin = f"auto.function.{identifier}" + + def __init__(self, timeout_warning: bool = False) -> None: + self.timeout_warning = timeout_warning + + @staticmethod + def setup_once() -> None: + lambda_bootstrap = get_lambda_bootstrap() + if not lambda_bootstrap: + logger.warning( + "Not running in AWS Lambda environment, " + "AwsLambdaIntegration disabled (could not find bootstrap module)" + ) + return + + if not hasattr(lambda_bootstrap, "handle_event_request"): + logger.warning( + "Not running in AWS Lambda environment, " + "AwsLambdaIntegration disabled (could not find handle_event_request)" + ) + return + + pre_37 = hasattr(lambda_bootstrap, "handle_http_request") # Python 3.6 + + if pre_37: + old_handle_event_request = lambda_bootstrap.handle_event_request + + def sentry_handle_event_request( + request_handler: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + request_handler = _wrap_handler(request_handler) + return old_handle_event_request(request_handler, *args, **kwargs) + + lambda_bootstrap.handle_event_request = sentry_handle_event_request + + old_handle_http_request = lambda_bootstrap.handle_http_request + + def sentry_handle_http_request( + request_handler: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + request_handler = _wrap_handler(request_handler) + return old_handle_http_request(request_handler, *args, **kwargs) + + lambda_bootstrap.handle_http_request = sentry_handle_http_request + + # Patch to_json to drain the queue. This should work even when the + # SDK is initialized inside of the handler + + old_to_json = lambda_bootstrap.to_json + + def sentry_to_json(*args: "Any", **kwargs: "Any") -> "Any": + _drain_queue() + return old_to_json(*args, **kwargs) + + lambda_bootstrap.to_json = sentry_to_json + else: + lambda_bootstrap.LambdaRuntimeClient.post_init_error = _wrap_init_error( + lambda_bootstrap.LambdaRuntimeClient.post_init_error + ) + + old_handle_event_request = lambda_bootstrap.handle_event_request + + def sentry_handle_event_request( # type: ignore + lambda_runtime_client, request_handler, *args, **kwargs + ): + request_handler = _wrap_handler(request_handler) + return old_handle_event_request( + lambda_runtime_client, request_handler, *args, **kwargs + ) + + lambda_bootstrap.handle_event_request = sentry_handle_event_request + + # Patch the runtime client to drain the queue. This should work + # even when the SDK is initialized inside of the handler + + def _wrap_post_function(f: "F") -> "F": + def inner(*args: "Any", **kwargs: "Any") -> "Any": + _drain_queue() + return f(*args, **kwargs) + + return inner # type: ignore + + lambda_bootstrap.LambdaRuntimeClient.post_invocation_result = ( + _wrap_post_function( + lambda_bootstrap.LambdaRuntimeClient.post_invocation_result + ) + ) + lambda_bootstrap.LambdaRuntimeClient.post_invocation_error = ( + _wrap_post_function( + lambda_bootstrap.LambdaRuntimeClient.post_invocation_error + ) + ) + + +def get_lambda_bootstrap() -> "Optional[Any]": + # Python 3.7: If the bootstrap module is *already imported*, it is the + # one we actually want to use (no idea what's in __main__) + # + # Python 3.8: bootstrap is also importable, but will be the same file + # as __main__ imported under a different name: + # + # sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__ + # sys.modules['__main__'] is not sys.modules['bootstrap'] + # + # Python 3.9: bootstrap is in __main__.awslambdaricmain + # + # On container builds using the `aws-lambda-python-runtime-interface-client` + # (awslamdaric) module, bootstrap is located in sys.modules['__main__'].bootstrap + # + # Such a setup would then make all monkeypatches useless. + if "bootstrap" in sys.modules: + return sys.modules["bootstrap"] + elif "__main__" in sys.modules: + module = sys.modules["__main__"] + # python3.9 runtime + if hasattr(module, "awslambdaricmain") and hasattr( + module.awslambdaricmain, "bootstrap" + ): + return module.awslambdaricmain.bootstrap + elif hasattr(module, "bootstrap"): + # awslambdaric python module in container builds + return module.bootstrap + + # python3.8 runtime + return module + else: + return None + + +def _make_request_event_processor( + aws_event: "Any", aws_context: "Any", configured_timeout: "Any" +) -> "EventProcessor": + start_time = datetime.now(timezone.utc) + + def event_processor( + sentry_event: "Event", hint: "Hint", start_time: "datetime" = start_time + ) -> "Optional[Event]": + remaining_time_in_milis = aws_context.get_remaining_time_in_millis() + exec_duration = configured_timeout - remaining_time_in_milis + + extra = sentry_event.setdefault("extra", {}) + extra["lambda"] = { + "function_name": aws_context.function_name, + "function_version": aws_context.function_version, + "invoked_function_arn": aws_context.invoked_function_arn, + "aws_request_id": aws_context.aws_request_id, + "execution_duration_in_millis": exec_duration, + "remaining_time_in_millis": remaining_time_in_milis, + } + + extra["cloudwatch logs"] = { + "url": _get_cloudwatch_logs_url(aws_context, start_time), + "log_group": aws_context.log_group_name, + "log_stream": aws_context.log_stream_name, + } + + request = sentry_event.get("request", {}) + + if "httpMethod" in aws_event: + request["method"] = aws_event["httpMethod"] + + request["url"] = _get_url(aws_event, aws_context) + + if "queryStringParameters" in aws_event: + request["query_string"] = aws_event["queryStringParameters"] + + if "headers" in aws_event: + request["headers"] = _filter_headers(aws_event["headers"]) + + if should_send_default_pii(): + user_info = sentry_event.setdefault("user", {}) + + identity = aws_event.get("identity") + if identity is None: + identity = {} + + id = identity.get("userArn") + if id is not None: + user_info.setdefault("id", id) + + ip = identity.get("sourceIp") + if ip is not None: + user_info.setdefault("ip_address", ip) + + if "body" in aws_event: + request["data"] = aws_event.get("body", "") + else: + if aws_event.get("body", None): + # Unfortunately couldn't find a way to get structured body from AWS + # event. Meaning every body is unstructured to us. + request["data"] = AnnotatedValue.removed_because_raw_data() + + sentry_event["request"] = deepcopy(request) + + return sentry_event + + return event_processor + + +def _get_url(aws_event: "Any", aws_context: "Any") -> str: + path = aws_event.get("path", None) + + headers = aws_event.get("headers") + if headers is None: + headers = {} + + host = headers.get("Host", None) + proto = headers.get("X-Forwarded-Proto", None) + if proto and host and path: + return "{}://{}{}".format(proto, host, path) + return "awslambda:///{}".format(aws_context.function_name) + + +def _get_cloudwatch_logs_url(aws_context: "Any", start_time: "datetime") -> str: + """ + Generates a CloudWatchLogs console URL based on the context object + + Arguments: + aws_context {Any} -- context from lambda handler + + Returns: + str -- AWS Console URL to logs. + """ + formatstring = "%Y-%m-%dT%H:%M:%SZ" + region = environ.get("AWS_REGION", "") + + url = ( + "https://console.{domain}/cloudwatch/home?region={region}" + "#logEventViewer:group={log_group};stream={log_stream}" + ";start={start_time};end={end_time}" + ).format( + domain="amazonaws.cn" if region.startswith("cn-") else "aws.amazon.com", + region=region, + log_group=aws_context.log_group_name, + log_stream=aws_context.log_stream_name, + start_time=(start_time - timedelta(seconds=1)).strftime(formatstring), + end_time=(datetime.now(timezone.utc) + timedelta(seconds=2)).strftime( + formatstring + ), + ) + + return url + + +def _parse_formatted_traceback(formatted_tb: "list[str]") -> "list[dict[str, Any]]": + frames = [] + for frame in formatted_tb: + match = re.match(r'File "(.+)", line (\d+), in (.+)', frame.strip()) + if match: + file_name, line_number, func_name = match.groups() + line_number = int(line_number) + frames.append( + { + "filename": file_name, + "function": func_name, + "lineno": line_number, + "vars": None, + "pre_context": None, + "context_line": None, + "post_context": None, + } + ) + return frames + + +def _event_from_error_json(error_json: "dict[str, Any]") -> "Event": + """ + Converts the error JSON from AWS Lambda into a Sentry error event. + This is not a full fletched event, but better than nothing. + + This is an example of where AWS creates the error JSON: + https://github.com/aws/aws-lambda-python-runtime-interface-client/blob/2.2.1/awslambdaric/bootstrap.py#L479 + """ + event: "Event" = { + "level": "error", + "exception": { + "values": [ + { + "type": error_json.get("errorType"), + "value": error_json.get("errorMessage"), + "stacktrace": { + "frames": _parse_formatted_traceback( + error_json.get("stackTrace", []) + ), + }, + "mechanism": { + "type": "aws_lambda", + "handled": False, + }, + } + ], + }, + } + + return event diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/beam.py b/lib/python3.12/site-packages/sentry_sdk/integrations/beam.py new file mode 100644 index 0000000000000000000000000000000000000000..6496e6293d05a76de8ddf9144eb5183bfba9b083 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/beam.py @@ -0,0 +1,168 @@ +import sys +import types +from functools import wraps + +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Iterator + from typing import TypeVar + from typing import Callable + + from sentry_sdk._types import ExcInfo + + T = TypeVar("T") + F = TypeVar("F", bound=Callable[..., Any]) + + +WRAPPED_FUNC = "_wrapped_{}_" +INSPECT_FUNC = "_inspect_{}" # Required format per apache_beam/transforms/core.py +USED_FUNC = "_sentry_used_" + + +class BeamIntegration(Integration): + identifier = "beam" + + @staticmethod + def setup_once() -> None: + from apache_beam.transforms.core import DoFn, ParDo # type: ignore + + ignore_logger("root") + ignore_logger("bundle_processor.create") + + function_patches = ["process", "start_bundle", "finish_bundle", "setup"] + for func_name in function_patches: + setattr( + DoFn, + INSPECT_FUNC.format(func_name), + _wrap_inspect_call(DoFn, func_name), + ) + + old_init = ParDo.__init__ + + def sentry_init_pardo( + self: "ParDo", fn: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + # Do not monkey patch init twice + if not getattr(self, "_sentry_is_patched", False): + for func_name in function_patches: + if not hasattr(fn, func_name): + continue + wrapped_func = WRAPPED_FUNC.format(func_name) + + # Check to see if inspect is set and process is not + # to avoid monkey patching process twice. + # Check to see if function is part of object for + # backwards compatibility. + process_func = getattr(fn, func_name) + inspect_func = getattr(fn, INSPECT_FUNC.format(func_name)) + if not getattr(inspect_func, USED_FUNC, False) and not getattr( + process_func, USED_FUNC, False + ): + setattr(fn, wrapped_func, process_func) + setattr(fn, func_name, _wrap_task_call(process_func)) + + self._sentry_is_patched = True + old_init(self, fn, *args, **kwargs) + + ParDo.__init__ = sentry_init_pardo + + +def _wrap_inspect_call(cls: "Any", func_name: "Any") -> "Any": + if not hasattr(cls, func_name): + return None + + def _inspect(self: "Any") -> "Any": + """ + Inspect function overrides the way Beam gets argspec. + """ + wrapped_func = WRAPPED_FUNC.format(func_name) + if hasattr(self, wrapped_func): + process_func = getattr(self, wrapped_func) + else: + process_func = getattr(self, func_name) + setattr(self, func_name, _wrap_task_call(process_func)) + setattr(self, wrapped_func, process_func) + + # getfullargspec is deprecated in more recent beam versions and get_function_args_defaults + # (which uses Signatures internally) should be used instead. + try: + from apache_beam.transforms.core import get_function_args_defaults + + return get_function_args_defaults(process_func) + except ImportError: + from apache_beam.typehints.decorators import getfullargspec # type: ignore + + return getfullargspec(process_func) + + setattr(_inspect, USED_FUNC, True) + return _inspect + + +def _wrap_task_call(func: "F") -> "F": + """ + Wrap task call with a try catch to get exceptions. + """ + + @wraps(func) + def _inner(*args: "Any", **kwargs: "Any") -> "Any": + try: + gen = func(*args, **kwargs) + except Exception: + raise_exception() + + if not isinstance(gen, types.GeneratorType): + return gen + return _wrap_generator_call(gen) + + setattr(_inner, USED_FUNC, True) + return _inner # type: ignore + + +@ensure_integration_enabled(BeamIntegration) +def _capture_exception(exc_info: "ExcInfo") -> None: + """ + Send Beam exception to Sentry. + """ + client = sentry_sdk.get_client() + + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "beam", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def raise_exception() -> None: + """ + Raise an exception. + """ + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(exc_info) + reraise(*exc_info) + + +def _wrap_generator_call(gen: "Iterator[T]") -> "Iterator[T]": + """ + Wrap the generator to handle any failures. + """ + while True: + try: + yield next(gen) + except StopIteration: + break + except Exception: + raise_exception() diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/boto3.py b/lib/python3.12/site-packages/sentry_sdk/integrations/boto3.py new file mode 100644 index 0000000000000000000000000000000000000000..b65e2c6b69e7d54bc1bb23a2ae055f69f1c1dd18 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/boto3.py @@ -0,0 +1,138 @@ +from functools import partial + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.tracing import Span +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + parse_url, + parse_version, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Optional + from typing import Type + +try: + from botocore import __version__ as BOTOCORE_VERSION # type: ignore + from botocore.client import BaseClient # type: ignore + from botocore.response import StreamingBody # type: ignore + from botocore.awsrequest import AWSRequest # type: ignore +except ImportError: + raise DidNotEnable("botocore is not installed") + + +class Boto3Integration(Integration): + identifier = "boto3" + origin = f"auto.http.{identifier}" + + @staticmethod + def setup_once() -> None: + version = parse_version(BOTOCORE_VERSION) + _check_minimum_version(Boto3Integration, version, "botocore") + + orig_init = BaseClient.__init__ + + def sentry_patched_init( + self: "Type[BaseClient]", *args: "Any", **kwargs: "Any" + ) -> None: + orig_init(self, *args, **kwargs) + meta = self.meta + service_id = meta.service_model.service_id.hyphenize() + meta.events.register( + "request-created", + partial(_sentry_request_created, service_id=service_id), + ) + meta.events.register("after-call", _sentry_after_call) + meta.events.register("after-call-error", _sentry_after_call_error) + + BaseClient.__init__ = sentry_patched_init + + +@ensure_integration_enabled(Boto3Integration) +def _sentry_request_created( + service_id: str, request: "AWSRequest", operation_name: str, **kwargs: "Any" +) -> None: + description = "aws.%s.%s" % (service_id, operation_name) + span = sentry_sdk.start_span( + op=OP.HTTP_CLIENT, + name=description, + origin=Boto3Integration.origin, + ) + + with capture_internal_exceptions(): + parsed_url = parse_url(request.url, sanitize=False) + span.set_data("aws.request.url", parsed_url.url) + span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query) + span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment) + + span.set_tag("aws.service_id", service_id) + span.set_tag("aws.operation_name", operation_name) + span.set_data(SPANDATA.HTTP_METHOD, request.method) + + # We do it in order for subsequent http calls/retries be + # attached to this span. + span.__enter__() + + # request.context is an open-ended data-structure + # where we can add anything useful in request life cycle. + request.context["_sentrysdk_span"] = span + + +def _sentry_after_call( + context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any" +) -> None: + span: "Optional[Span]" = context.pop("_sentrysdk_span", None) + + # Span could be absent if the integration is disabled. + if span is None: + return + span.__exit__(None, None, None) + + body = parsed.get("Body") + if not isinstance(body, StreamingBody): + return + + streaming_span = span.start_child( + op=OP.HTTP_CLIENT_STREAM, + name=span.description, + origin=Boto3Integration.origin, + ) + + orig_read = body.read + orig_close = body.close + + def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes: + try: + ret = orig_read(*args, **kwargs) + if not ret: + streaming_span.finish() + return ret + except Exception: + streaming_span.finish() + raise + + body.read = sentry_streaming_body_read + + def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None: + streaming_span.finish() + orig_close(*args, **kwargs) + + body.close = sentry_streaming_body_close + + +def _sentry_after_call_error( + context: "Dict[str, Any]", exception: "Type[BaseException]", **kwargs: "Any" +) -> None: + span: "Optional[Span]" = context.pop("_sentrysdk_span", None) + + # Span could be absent if the integration is disabled. + if span is None: + return + span.__exit__(type(exception), exception, None) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/bottle.py b/lib/python3.12/site-packages/sentry_sdk/integrations/bottle.py new file mode 100644 index 0000000000000000000000000000000000000000..29862c6d6c94c482b81dc890634fea80ac903e5c --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/bottle.py @@ -0,0 +1,211 @@ +import functools + +import sentry_sdk +from sentry_sdk.tracing import SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + parse_version, + transaction_from_function, +) +from sentry_sdk.integrations import ( + Integration, + DidNotEnable, + _DEFAULT_FAILED_REQUEST_STATUS_CODES, + _check_minimum_version, +) +from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware +from sentry_sdk.integrations._wsgi_common import RequestExtractor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Set + + from sentry_sdk.integrations.wsgi import _ScopedResponse + from typing import Any + from typing import Dict + from typing import Callable + from typing import Optional + from bottle import FileUpload, FormsDict, LocalRequest # type: ignore + + from sentry_sdk._types import EventProcessor, Event + +try: + from bottle import ( + Bottle, + HTTPResponse, + Route, + request as bottle_request, + __version__ as BOTTLE_VERSION, + ) +except ImportError: + raise DidNotEnable("Bottle not installed") + + +TRANSACTION_STYLE_VALUES = ("endpoint", "url") + + +class BottleIntegration(Integration): + identifier = "bottle" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__( + self, + transaction_style: str = "endpoint", + *, + failed_request_status_codes: "Set[int]" = _DEFAULT_FAILED_REQUEST_STATUS_CODES, + ) -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + self.failed_request_status_codes = failed_request_status_codes + + @staticmethod + def setup_once() -> None: + version = parse_version(BOTTLE_VERSION) + _check_minimum_version(BottleIntegration, version) + + old_app = Bottle.__call__ + + @ensure_integration_enabled(BottleIntegration, old_app) + def sentry_patched_wsgi_app( + self: "Any", environ: "Dict[str, str]", start_response: "Callable[..., Any]" + ) -> "_ScopedResponse": + middleware = SentryWsgiMiddleware( + lambda *a, **kw: old_app(self, *a, **kw), + span_origin=BottleIntegration.origin, + ) + + return middleware(environ, start_response) + + Bottle.__call__ = sentry_patched_wsgi_app + + old_handle = Bottle._handle + + @functools.wraps(old_handle) + def _patched_handle(self: "Bottle", environ: "Dict[str, Any]") -> "Any": + integration = sentry_sdk.get_client().get_integration(BottleIntegration) + if integration is None: + return old_handle(self, environ) + + scope = sentry_sdk.get_isolation_scope() + scope._name = "bottle" + scope.add_event_processor( + _make_request_event_processor(self, bottle_request, integration) + ) + res = old_handle(self, environ) + + return res + + Bottle._handle = _patched_handle + + old_make_callback = Route._make_callback + + @functools.wraps(old_make_callback) + def patched_make_callback( + self: "Route", *args: object, **kwargs: object + ) -> "Any": + prepared_callback = old_make_callback(self, *args, **kwargs) + + integration = sentry_sdk.get_client().get_integration(BottleIntegration) + if integration is None: + return prepared_callback + + def wrapped_callback(*args: object, **kwargs: object) -> "Any": + try: + res = prepared_callback(*args, **kwargs) + except Exception as exception: + _capture_exception(exception, handled=False) + raise exception + + if ( + isinstance(res, HTTPResponse) + and res.status_code in integration.failed_request_status_codes + ): + _capture_exception(res, handled=True) + + return res + + return wrapped_callback + + Route._make_callback = patched_make_callback + + +class BottleRequestExtractor(RequestExtractor): + def env(self) -> "Dict[str, str]": + return self.request.environ + + def cookies(self) -> "Dict[str, str]": + return self.request.cookies + + def raw_data(self) -> bytes: + return self.request.body.read() + + def form(self) -> "FormsDict": + if self.is_json(): + return None + return self.request.forms.decode() + + def files(self) -> "Optional[Dict[str, str]]": + if self.is_json(): + return None + + return self.request.files + + def size_of_file(self, file: "FileUpload") -> int: + return file.content_length + + +def _set_transaction_name_and_source( + event: "Event", transaction_style: str, request: "Any" +) -> None: + name = "" + + if transaction_style == "url": + try: + name = request.route.rule or "" + except RuntimeError: + pass + + elif transaction_style == "endpoint": + try: + name = ( + request.route.name + or transaction_from_function(request.route.callback) + or "" + ) + except RuntimeError: + pass + + event["transaction"] = name + event["transaction_info"] = {"source": SOURCE_FOR_STYLE[transaction_style]} + + +def _make_request_event_processor( + app: "Bottle", request: "LocalRequest", integration: "BottleIntegration" +) -> "EventProcessor": + def event_processor(event: "Event", hint: "dict[str, Any]") -> "Event": + _set_transaction_name_and_source(event, integration.transaction_style, request) + + with capture_internal_exceptions(): + BottleRequestExtractor(request).extract_into_event(event) + + return event + + return event_processor + + +def _capture_exception(exception: BaseException, handled: bool) -> None: + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "bottle", "handled": handled}, + ) + sentry_sdk.capture_event(event, hint=hint) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/chalice.py b/lib/python3.12/site-packages/sentry_sdk/integrations/chalice.py new file mode 100644 index 0000000000000000000000000000000000000000..89911dc1abc2d33661849c1e6f817e06d9aafb77 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/chalice.py @@ -0,0 +1,132 @@ +import sys +from functools import wraps + +import sentry_sdk +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.integrations.aws_lambda import _make_request_event_processor +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + parse_version, + reraise, +) + +try: + import chalice # type: ignore + from chalice import __version__ as CHALICE_VERSION + from chalice import Chalice, ChaliceViewError + from chalice.app import EventSourceHandler as ChaliceEventSourceHandler # type: ignore +except ImportError: + raise DidNotEnable("Chalice is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import TypeVar + from typing import Callable + + F = TypeVar("F", bound=Callable[..., Any]) + + +class EventSourceHandler(ChaliceEventSourceHandler): # type: ignore + def __call__(self, event: "Any", context: "Any") -> "Any": + client = sentry_sdk.get_client() + + with sentry_sdk.isolation_scope() as scope: + with capture_internal_exceptions(): + configured_time = context.get_remaining_time_in_millis() + scope.add_event_processor( + _make_request_event_processor(event, context, configured_time) + ) + try: + return ChaliceEventSourceHandler.__call__(self, event, context) + except Exception: + exc_info = sys.exc_info() + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "chalice", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + client.flush() + reraise(*exc_info) + + +def _get_view_function_response( + app: "Any", view_function: "F", function_args: "Any" +) -> "F": + @wraps(view_function) + def wrapped_view_function(**function_args: "Any") -> "Any": + client = sentry_sdk.get_client() + with sentry_sdk.isolation_scope() as scope: + with capture_internal_exceptions(): + configured_time = app.lambda_context.get_remaining_time_in_millis() + scope.set_transaction_name( + app.lambda_context.function_name, + source=TransactionSource.COMPONENT, + ) + + scope.add_event_processor( + _make_request_event_processor( + app.current_request.to_dict(), + app.lambda_context, + configured_time, + ) + ) + try: + return view_function(**function_args) + except Exception as exc: + if isinstance(exc, ChaliceViewError): + raise + exc_info = sys.exc_info() + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "chalice", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + client.flush() + raise + + return wrapped_view_function # type: ignore + + +class ChaliceIntegration(Integration): + identifier = "chalice" + + @staticmethod + def setup_once() -> None: + version = parse_version(CHALICE_VERSION) + + if version is None: + raise DidNotEnable("Unparsable Chalice version: {}".format(CHALICE_VERSION)) + + if version < (1, 20): + old_get_view_function_response = Chalice._get_view_function_response + else: + from chalice.app import RestAPIEventHandler + + old_get_view_function_response = ( + RestAPIEventHandler._get_view_function_response + ) + + def sentry_event_response( + app: "Any", view_function: "F", function_args: "Dict[str, Any]" + ) -> "Any": + wrapped_view_function = _get_view_function_response( + app, view_function, function_args + ) + + return old_get_view_function_response( + app, wrapped_view_function, function_args + ) + + if version < (1, 20): + Chalice._get_view_function_response = sentry_event_response + else: + RestAPIEventHandler._get_view_function_response = sentry_event_response + # for everything else (like events) + chalice.app.EventSourceHandler = EventSourceHandler diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/clickhouse_driver.py b/lib/python3.12/site-packages/sentry_sdk/integrations/clickhouse_driver.py new file mode 100644 index 0000000000000000000000000000000000000000..7bbea9421050e7e34014a154b102eb9f64ee3dfc --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/clickhouse_driver.py @@ -0,0 +1,171 @@ +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.tracing import Span +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled + +from typing import TYPE_CHECKING, TypeVar + +# Hack to get new Python features working in older versions +# without introducing a hard dependency on `typing_extensions` +# from: https://stackoverflow.com/a/71944042/300572 +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import Any, ParamSpec, Callable +else: + # Fake ParamSpec + class ParamSpec: + def __init__(self, _): + self.args = None + self.kwargs = None + + # Callable[anything] will return None + class _Callable: + def __getitem__(self, _): + return None + + # Make instances + Callable = _Callable() + + +try: + from clickhouse_driver import VERSION # type: ignore[import-not-found] + from clickhouse_driver.client import Client # type: ignore[import-not-found] + from clickhouse_driver.connection import Connection # type: ignore[import-not-found] + +except ImportError: + raise DidNotEnable("clickhouse-driver not installed.") + + +class ClickhouseDriverIntegration(Integration): + identifier = "clickhouse_driver" + origin = f"auto.db.{identifier}" + + @staticmethod + def setup_once() -> None: + _check_minimum_version(ClickhouseDriverIntegration, VERSION) + + # Every query is done using the Connection's `send_query` function + Connection.send_query = _wrap_start(Connection.send_query) + + # If the query contains parameters then the send_data function is used to send those parameters to clickhouse + _wrap_send_data() + + # Every query ends either with the Client's `receive_end_of_query` (no result expected) + # or its `receive_result` (result expected) + Client.receive_end_of_query = _wrap_end(Client.receive_end_of_query) + if hasattr(Client, "receive_end_of_insert_query"): + # In 0.2.7, insert queries are handled separately via `receive_end_of_insert_query` + Client.receive_end_of_insert_query = _wrap_end( + Client.receive_end_of_insert_query + ) + Client.receive_result = _wrap_end(Client.receive_result) + + +P = ParamSpec("P") +T = TypeVar("T") + + +def _wrap_start(f: "Callable[P, T]") -> "Callable[P, T]": + @ensure_integration_enabled(ClickhouseDriverIntegration, f) + def _inner(*args: "P.args", **kwargs: "P.kwargs") -> "T": + connection = args[0] + query = args[1] + query_id = args[2] if len(args) > 2 else kwargs.get("query_id") + params = args[3] if len(args) > 3 else kwargs.get("params") + + span = sentry_sdk.start_span( + op=OP.DB, + name=query, + origin=ClickhouseDriverIntegration.origin, + ) + + connection._sentry_span = span # type: ignore[attr-defined] + + _set_db_data(span, connection) + + span.set_data("query", query) + + if query_id: + span.set_data("db.query_id", query_id) + + if params and should_send_default_pii(): + span.set_data("db.params", params) + + # run the original code + ret = f(*args, **kwargs) + + return ret + + return _inner + + +def _wrap_end(f: "Callable[P, T]") -> "Callable[P, T]": + def _inner_end(*args: "P.args", **kwargs: "P.kwargs") -> "T": + res = f(*args, **kwargs) + instance = args[0] + span = getattr(instance.connection, "_sentry_span", None) # type: ignore[attr-defined] + + if span is not None: + if res is not None and should_send_default_pii(): + span.set_data("db.result", res) + + with capture_internal_exceptions(): + span.scope.add_breadcrumb( + message=span._data.pop("query"), category="query", data=span._data + ) + + span.finish() + + return res + + return _inner_end + + +def _wrap_send_data() -> None: + original_send_data = Client.send_data + + def _inner_send_data( # type: ignore[no-untyped-def] # clickhouse-driver does not type send_data + self, sample_block, data, types_check=False, columnar=False, *args, **kwargs + ): + span = getattr(self.connection, "_sentry_span", None) + + if span is not None: + _set_db_data(span, self.connection) + + if should_send_default_pii(): + db_params = span._data.get("db.params", []) + + if isinstance(data, (list, tuple)): + db_params.extend(data) + + else: # data is a generic iterator + orig_data = data + + # Wrap the generator to add items to db.params as they are yielded. + # This allows us to send the params to Sentry without needing to allocate + # memory for the entire generator at once. + def wrapped_generator() -> "Iterator[Any]": + for item in orig_data: + db_params.append(item) + yield item + + # Replace the original iterator with the wrapped one. + data = wrapped_generator() + + span.set_data("db.params", db_params) + + return original_send_data( + self, sample_block, data, types_check, columnar, *args, **kwargs + ) + + Client.send_data = _inner_send_data + + +def _set_db_data(span: "Span", connection: "Connection") -> None: + span.set_data(SPANDATA.DB_SYSTEM, "clickhouse") + span.set_data(SPANDATA.SERVER_ADDRESS, connection.host) + span.set_data(SPANDATA.SERVER_PORT, connection.port) + span.set_data(SPANDATA.DB_NAME, connection.database) + span.set_data(SPANDATA.DB_USER, connection.user) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/cloud_resource_context.py b/lib/python3.12/site-packages/sentry_sdk/integrations/cloud_resource_context.py new file mode 100644 index 0000000000000000000000000000000000000000..09d55ac1199f57c6991997dcba86effe89741b82 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/cloud_resource_context.py @@ -0,0 +1,272 @@ +import json +import urllib3 + +from sentry_sdk.integrations import Integration +from sentry_sdk.api import set_context +from sentry_sdk.utils import logger + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Dict + + +CONTEXT_TYPE = "cloud_resource" + +HTTP_TIMEOUT = 2.0 + +AWS_METADATA_HOST = "169.254.169.254" +AWS_TOKEN_URL = "http://{}/latest/api/token".format(AWS_METADATA_HOST) +AWS_METADATA_URL = "http://{}/latest/dynamic/instance-identity/document".format( + AWS_METADATA_HOST +) + +GCP_METADATA_HOST = "metadata.google.internal" +GCP_METADATA_URL = "http://{}/computeMetadata/v1/?recursive=true".format( + GCP_METADATA_HOST +) + + +class CLOUD_PROVIDER: # noqa: N801 + """ + Name of the cloud provider. + see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/ + """ + + ALIBABA = "alibaba_cloud" + AWS = "aws" + AZURE = "azure" + GCP = "gcp" + IBM = "ibm_cloud" + TENCENT = "tencent_cloud" + + +class CLOUD_PLATFORM: # noqa: N801 + """ + The cloud platform. + see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/ + """ + + AWS_EC2 = "aws_ec2" + GCP_COMPUTE_ENGINE = "gcp_compute_engine" + + +class CloudResourceContextIntegration(Integration): + """ + Adds cloud resource context to the Senty scope + """ + + identifier = "cloudresourcecontext" + + cloud_provider = "" + + aws_token = "" + http = urllib3.PoolManager(timeout=HTTP_TIMEOUT) + + gcp_metadata = None + + def __init__(self, cloud_provider: str = "") -> None: + CloudResourceContextIntegration.cloud_provider = cloud_provider + + @classmethod + def _is_aws(cls) -> bool: + try: + r = cls.http.request( + "PUT", + AWS_TOKEN_URL, + headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"}, + ) + + if r.status != 200: + return False + + cls.aws_token = r.data.decode() + return True + + except urllib3.exceptions.TimeoutError: + logger.debug( + "AWS metadata service timed out after %s seconds", HTTP_TIMEOUT + ) + return False + except Exception as e: + logger.debug("Error checking AWS metadata service: %s", str(e)) + return False + + @classmethod + def _get_aws_context(cls) -> "Dict[str, str]": + ctx = { + "cloud.provider": CLOUD_PROVIDER.AWS, + "cloud.platform": CLOUD_PLATFORM.AWS_EC2, + } + + try: + r = cls.http.request( + "GET", + AWS_METADATA_URL, + headers={"X-aws-ec2-metadata-token": cls.aws_token}, + ) + + if r.status != 200: + return ctx + + data = json.loads(r.data.decode("utf-8")) + + try: + ctx["cloud.account.id"] = data["accountId"] + except Exception: + pass + + try: + ctx["cloud.availability_zone"] = data["availabilityZone"] + except Exception: + pass + + try: + ctx["cloud.region"] = data["region"] + except Exception: + pass + + try: + ctx["host.id"] = data["instanceId"] + except Exception: + pass + + try: + ctx["host.type"] = data["instanceType"] + except Exception: + pass + + except urllib3.exceptions.TimeoutError: + logger.debug( + "AWS metadata service timed out after %s seconds", HTTP_TIMEOUT + ) + except Exception as e: + logger.debug("Error fetching AWS metadata: %s", str(e)) + + return ctx + + @classmethod + def _is_gcp(cls) -> bool: + try: + r = cls.http.request( + "GET", + GCP_METADATA_URL, + headers={"Metadata-Flavor": "Google"}, + ) + + if r.status != 200: + return False + + cls.gcp_metadata = json.loads(r.data.decode("utf-8")) + return True + + except urllib3.exceptions.TimeoutError: + logger.debug( + "GCP metadata service timed out after %s seconds", HTTP_TIMEOUT + ) + return False + except Exception as e: + logger.debug("Error checking GCP metadata service: %s", str(e)) + return False + + @classmethod + def _get_gcp_context(cls) -> "Dict[str, str]": + ctx = { + "cloud.provider": CLOUD_PROVIDER.GCP, + "cloud.platform": CLOUD_PLATFORM.GCP_COMPUTE_ENGINE, + } + + try: + if cls.gcp_metadata is None: + r = cls.http.request( + "GET", + GCP_METADATA_URL, + headers={"Metadata-Flavor": "Google"}, + ) + + if r.status != 200: + return ctx + + cls.gcp_metadata = json.loads(r.data.decode("utf-8")) + + try: + ctx["cloud.account.id"] = cls.gcp_metadata["project"]["projectId"] + except Exception: + pass + + try: + ctx["cloud.availability_zone"] = cls.gcp_metadata["instance"][ + "zone" + ].split("/")[-1] + except Exception: + pass + + try: + # only populated in google cloud run + ctx["cloud.region"] = cls.gcp_metadata["instance"]["region"].split("/")[ + -1 + ] + except Exception: + pass + + try: + ctx["host.id"] = cls.gcp_metadata["instance"]["id"] + except Exception: + pass + + except urllib3.exceptions.TimeoutError: + logger.debug( + "GCP metadata service timed out after %s seconds", HTTP_TIMEOUT + ) + except Exception as e: + logger.debug("Error fetching GCP metadata: %s", str(e)) + + return ctx + + @classmethod + def _get_cloud_provider(cls) -> str: + if cls._is_aws(): + return CLOUD_PROVIDER.AWS + + if cls._is_gcp(): + return CLOUD_PROVIDER.GCP + + return "" + + @classmethod + def _get_cloud_resource_context(cls) -> "Dict[str, str]": + cloud_provider = ( + cls.cloud_provider + if cls.cloud_provider != "" + else CloudResourceContextIntegration._get_cloud_provider() + ) + if cloud_provider in context_getters.keys(): + return context_getters[cloud_provider]() + + return {} + + @staticmethod + def setup_once() -> None: + cloud_provider = CloudResourceContextIntegration.cloud_provider + unsupported_cloud_provider = ( + cloud_provider != "" and cloud_provider not in context_getters.keys() + ) + + if unsupported_cloud_provider: + logger.warning( + "Invalid value for cloud_provider: %s (must be in %s). Falling back to autodetection...", + CloudResourceContextIntegration.cloud_provider, + list(context_getters.keys()), + ) + + context = CloudResourceContextIntegration._get_cloud_resource_context() + if context != {}: + set_context(CONTEXT_TYPE, context) + + +# Map with the currently supported cloud providers +# mapping to functions extracting the context +context_getters = { + CLOUD_PROVIDER.AWS: CloudResourceContextIntegration._get_aws_context, + CLOUD_PROVIDER.GCP: CloudResourceContextIntegration._get_gcp_context, +} diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/cohere.py b/lib/python3.12/site-packages/sentry_sdk/integrations/cohere.py new file mode 100644 index 0000000000000000000000000000000000000000..f45a02f2b53a395613156d2e9e458a7ca37b8a99 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/cohere.py @@ -0,0 +1,269 @@ +import sys +from functools import wraps + +from sentry_sdk import consts +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.consts import SPANDATA +from sentry_sdk.ai.utils import set_data_normalized + +from typing import TYPE_CHECKING + +from sentry_sdk.tracing_utils import set_span_errored + +if TYPE_CHECKING: + from typing import Any, Callable, Iterator + from sentry_sdk.tracing import Span + +import sentry_sdk +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise + +try: + from cohere.client import Client + from cohere.base_client import BaseCohere + from cohere import ( + ChatStreamEndEvent, + NonStreamedChatResponse, + ) + + if TYPE_CHECKING: + from cohere import StreamedChatResponse +except ImportError: + raise DidNotEnable("Cohere not installed") + +try: + # cohere 5.9.3+ + from cohere import StreamEndStreamedChatResponse +except ImportError: + from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse + + +COLLECTED_CHAT_PARAMS = { + "model": SPANDATA.AI_MODEL_ID, + "k": SPANDATA.AI_TOP_K, + "p": SPANDATA.AI_TOP_P, + "seed": SPANDATA.AI_SEED, + "frequency_penalty": SPANDATA.AI_FREQUENCY_PENALTY, + "presence_penalty": SPANDATA.AI_PRESENCE_PENALTY, + "raw_prompting": SPANDATA.AI_RAW_PROMPTING, +} + +COLLECTED_PII_CHAT_PARAMS = { + "tools": SPANDATA.AI_TOOLS, + "preamble": SPANDATA.AI_PREAMBLE, +} + +COLLECTED_CHAT_RESP_ATTRS = { + "generation_id": SPANDATA.AI_GENERATION_ID, + "is_search_required": SPANDATA.AI_SEARCH_REQUIRED, + "finish_reason": SPANDATA.AI_FINISH_REASON, +} + +COLLECTED_PII_CHAT_RESP_ATTRS = { + "citations": SPANDATA.AI_CITATIONS, + "documents": SPANDATA.AI_DOCUMENTS, + "search_queries": SPANDATA.AI_SEARCH_QUERIES, + "search_results": SPANDATA.AI_SEARCH_RESULTS, + "tool_calls": SPANDATA.AI_TOOL_CALLS, +} + + +class CohereIntegration(Integration): + identifier = "cohere" + origin = f"auto.ai.{identifier}" + + def __init__(self: "CohereIntegration", include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False) + Client.embed = _wrap_embed(Client.embed) + BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True) + + +def _capture_exception(exc: "Any") -> None: + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "cohere", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _wrap_chat(f: "Callable[..., Any]", streaming: bool) -> "Callable[..., Any]": + def collect_chat_response_fields( + span: "Span", res: "NonStreamedChatResponse", include_pii: bool + ) -> None: + if include_pii: + if hasattr(res, "text"): + set_data_normalized( + span, + SPANDATA.AI_RESPONSES, + [res.text], + ) + for pii_attr in COLLECTED_PII_CHAT_RESP_ATTRS: + if hasattr(res, pii_attr): + set_data_normalized(span, "ai." + pii_attr, getattr(res, pii_attr)) + + for attr in COLLECTED_CHAT_RESP_ATTRS: + if hasattr(res, attr): + set_data_normalized(span, "ai." + attr, getattr(res, attr)) + + if hasattr(res, "meta"): + if hasattr(res.meta, "billed_units"): + record_token_usage( + span, + input_tokens=res.meta.billed_units.input_tokens, + output_tokens=res.meta.billed_units.output_tokens, + ) + elif hasattr(res.meta, "tokens"): + record_token_usage( + span, + input_tokens=res.meta.tokens.input_tokens, + output_tokens=res.meta.tokens.output_tokens, + ) + + if hasattr(res.meta, "warnings"): + set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings) + + @wraps(f) + def new_chat(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(CohereIntegration) + + if ( + integration is None + or "message" not in kwargs + or not isinstance(kwargs.get("message"), str) + ): + return f(*args, **kwargs) + + message = kwargs.get("message") + + span = sentry_sdk.start_span( + op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE, + name="cohere.client.Chat", + origin=CohereIntegration.origin, + ) + span.__enter__() + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + span.__exit__(None, None, None) + reraise(*exc_info) + + with capture_internal_exceptions(): + if should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, + SPANDATA.AI_INPUT_MESSAGES, + list( + map( + lambda x: { + "role": getattr(x, "role", "").lower(), + "content": getattr(x, "message", ""), + }, + kwargs.get("chat_history", []), + ) + ) + + [{"role": "user", "content": message}], + ) + for k, v in COLLECTED_PII_CHAT_PARAMS.items(): + if k in kwargs: + set_data_normalized(span, v, kwargs[k]) + + for k, v in COLLECTED_CHAT_PARAMS.items(): + if k in kwargs: + set_data_normalized(span, v, kwargs[k]) + set_data_normalized(span, SPANDATA.AI_STREAMING, False) + + if streaming: + old_iterator = res + + def new_iterator() -> "Iterator[StreamedChatResponse]": + with capture_internal_exceptions(): + for x in old_iterator: + if isinstance(x, ChatStreamEndEvent) or isinstance( + x, StreamEndStreamedChatResponse + ): + collect_chat_response_fields( + span, + x.response, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) + yield x + + span.__exit__(None, None, None) + + return new_iterator() + elif isinstance(res, NonStreamedChatResponse): + collect_chat_response_fields( + span, + res, + include_pii=should_send_default_pii() + and integration.include_prompts, + ) + span.__exit__(None, None, None) + else: + set_data_normalized(span, "unknown_response", True) + span.__exit__(None, None, None) + return res + + return new_chat + + +def _wrap_embed(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_embed(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(CohereIntegration) + if integration is None: + return f(*args, **kwargs) + + with sentry_sdk.start_span( + op=consts.OP.COHERE_EMBEDDINGS_CREATE, + name="Cohere Embedding Creation", + origin=CohereIntegration.origin, + ) as span: + if "texts" in kwargs and ( + should_send_default_pii() and integration.include_prompts + ): + if isinstance(kwargs["texts"], str): + set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]]) + elif ( + isinstance(kwargs["texts"], list) + and len(kwargs["texts"]) > 0 + and isinstance(kwargs["texts"][0], str) + ): + set_data_normalized( + span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"] + ) + + if "model" in kwargs: + set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"]) + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + if ( + hasattr(res, "meta") + and hasattr(res.meta, "billed_units") + and hasattr(res.meta.billed_units, "input_tokens") + ): + record_token_usage( + span, + input_tokens=res.meta.billed_units.input_tokens, + total_tokens=res.meta.billed_units.input_tokens, + ) + return res + + return new_embed diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/dedupe.py b/lib/python3.12/site-packages/sentry_sdk/integrations/dedupe.py new file mode 100644 index 0000000000000000000000000000000000000000..09e60e4be6550d49ef37bf3e44eadce1ce52cf2f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/dedupe.py @@ -0,0 +1,63 @@ +import weakref + +import sentry_sdk +from sentry_sdk.utils import ContextVar, logger +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import add_global_event_processor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional + + from sentry_sdk._types import Event, Hint + + +class DedupeIntegration(Integration): + identifier = "dedupe" + + def __init__(self) -> None: + self._last_seen = ContextVar("last-seen") + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def processor(event: "Event", hint: "Optional[Hint]") -> "Optional[Event]": + if hint is None: + return event + + integration = sentry_sdk.get_client().get_integration(DedupeIntegration) + if integration is None: + return event + + exc_info = hint.get("exc_info", None) + if exc_info is None: + return event + + last_seen = integration._last_seen.get(None) + if last_seen is not None: + # last_seen is either a weakref or the original instance + last_seen = ( + last_seen() if isinstance(last_seen, weakref.ref) else last_seen + ) + + exc = exc_info[1] + if last_seen is exc: + logger.info("DedupeIntegration dropped duplicated error event %s", exc) + return None + + # we can only weakref non builtin types + try: + integration._last_seen.set(weakref.ref(exc)) + except TypeError: + integration._last_seen.set(exc) + + return event + + @staticmethod + def reset_last_seen() -> None: + integration = sentry_sdk.get_client().get_integration(DedupeIntegration) + if integration is None: + return + + integration._last_seen.set(None) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/dramatiq.py b/lib/python3.12/site-packages/sentry_sdk/integrations/dramatiq.py new file mode 100644 index 0000000000000000000000000000000000000000..f954d4fb98f7fc6113475e13d6f77caaae6a136c --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/dramatiq.py @@ -0,0 +1,228 @@ +import json + +import sentry_sdk +from sentry_sdk.consts import OP, SPANSTATUS +from sentry_sdk.api import continue_trace, get_baggage, get_traceparent +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import request_body_within_bounds +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SENTRY_TRACE_HEADER_NAME, + TransactionSource, +) +from sentry_sdk.utils import ( + AnnotatedValue, + capture_internal_exceptions, + event_from_exception, +) +from typing import TypeVar + +R = TypeVar("R") + +try: + from dramatiq.broker import Broker + from dramatiq.middleware import Middleware, default_middleware + from dramatiq.errors import Retry + from dramatiq.message import Message +except ImportError: + raise DidNotEnable("Dramatiq is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Dict, Optional, Union + from sentry_sdk._types import Event, Hint + + +class DramatiqIntegration(Integration): + """ + Dramatiq integration for Sentry + + Please make sure that you call `sentry_sdk.init` *before* initializing + your broker, as it monkey patches `Broker.__init__`. + + This integration was originally developed and maintained + by https://github.com/jacobsvante and later donated to the Sentry + project. + """ + + identifier = "dramatiq" + origin = f"auto.queue.{identifier}" + + @staticmethod + def setup_once() -> None: + _patch_dramatiq_broker() + + +def _patch_dramatiq_broker() -> None: + original_broker__init__ = Broker.__init__ + + def sentry_patched_broker__init__( + self: "Broker", *args: "Any", **kw: "Any" + ) -> None: + integration = sentry_sdk.get_client().get_integration(DramatiqIntegration) + + try: + middleware = kw.pop("middleware") + except KeyError: + # Unfortunately Broker and StubBroker allows middleware to be + # passed in as positional arguments, whilst RabbitmqBroker and + # RedisBroker does not. + if len(args) == 1: + middleware = args[0] + args = [] # type: ignore + else: + middleware = None + + if middleware is None: + middleware = list(m() for m in default_middleware) + else: + middleware = list(middleware) + + if integration is not None: + middleware = [m for m in middleware if not isinstance(m, SentryMiddleware)] + middleware.insert(0, SentryMiddleware()) + + kw["middleware"] = middleware + original_broker__init__(self, *args, **kw) + + Broker.__init__ = sentry_patched_broker__init__ + + +class SentryMiddleware(Middleware): # type: ignore[misc] + """ + A Dramatiq middleware that automatically captures and sends + exceptions to Sentry. + + This is automatically added to every instantiated broker via the + DramatiqIntegration. + """ + + SENTRY_HEADERS_NAME = "_sentry_headers" + + def before_enqueue( + self, broker: "Broker", message: "Message[R]", delay: int + ) -> None: + integration = sentry_sdk.get_client().get_integration(DramatiqIntegration) + if integration is None: + return + + message.options[self.SENTRY_HEADERS_NAME] = { + BAGGAGE_HEADER_NAME: get_baggage(), + SENTRY_TRACE_HEADER_NAME: get_traceparent(), + } + + def before_process_message(self, broker: "Broker", message: "Message[R]") -> None: + integration = sentry_sdk.get_client().get_integration(DramatiqIntegration) + if integration is None: + return + + message._scope_manager = sentry_sdk.isolation_scope() + scope = message._scope_manager.__enter__() + scope.clear_breadcrumbs() + scope.set_extra("dramatiq_message_id", message.message_id) + scope.add_event_processor(_make_message_event_processor(message, integration)) + + sentry_headers = message.options.get(self.SENTRY_HEADERS_NAME) or {} + if "retries" in message.options: + # start new trace in case of retrying + sentry_headers = {} + + transaction = continue_trace( + sentry_headers, + name=message.actor_name, + op=OP.QUEUE_TASK_DRAMATIQ, + source=TransactionSource.TASK, + origin=DramatiqIntegration.origin, + ) + transaction.set_status(SPANSTATUS.OK) + sentry_sdk.start_transaction( + transaction, + name=message.actor_name, + op=OP.QUEUE_TASK_DRAMATIQ, + source=TransactionSource.TASK, + ) + transaction.__enter__() + + def after_process_message( + self, + broker: "Broker", + message: "Message[R]", + *, + result: "Optional[Any]" = None, + exception: "Optional[Exception]" = None, + ) -> None: + integration = sentry_sdk.get_client().get_integration(DramatiqIntegration) + if integration is None: + return + + actor = broker.get_actor(message.actor_name) + throws = message.options.get("throws") or actor.options.get("throws") + + scope_manager = message._scope_manager + transaction = sentry_sdk.get_current_scope().transaction + if not transaction: + return None + + is_event_capture_required = ( + exception is not None + and not (throws and isinstance(exception, throws)) + and not isinstance(exception, Retry) + ) + if not is_event_capture_required: + # normal transaction finish + transaction.__exit__(None, None, None) + scope_manager.__exit__(None, None, None) + return + + event, hint = event_from_exception( + exception, # type: ignore[arg-type] + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": DramatiqIntegration.identifier, + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + # transaction error + transaction.__exit__(type(exception), exception, None) + scope_manager.__exit__(type(exception), exception, None) + + after_skip_message = after_process_message + + +def _make_message_event_processor( + message: "Message[R]", integration: "DramatiqIntegration" +) -> "Callable[[Event, Hint], Optional[Event]]": + def inner(event: "Event", hint: "Hint") -> "Optional[Event]": + with capture_internal_exceptions(): + DramatiqMessageExtractor(message).extract_into_event(event) + + return event + + return inner + + +class DramatiqMessageExtractor: + def __init__(self, message: "Message[R]") -> None: + self.message_data = dict(message.asdict()) + + def content_length(self) -> int: + return len(json.dumps(self.message_data)) + + def extract_into_event(self, event: "Event") -> None: + client = sentry_sdk.get_client() + if not client.is_active(): + return + + contexts = event.setdefault("contexts", {}) + request_info = contexts.setdefault("dramatiq", {}) + request_info["type"] = "dramatiq" + + data: "Optional[Union[AnnotatedValue, Dict[str, Any]]]" = None + if not request_body_within_bounds(client, self.content_length()): + data = AnnotatedValue.removed_because_over_size_limit() + else: + data = self.message_data + + request_info["data"] = data diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/excepthook.py b/lib/python3.12/site-packages/sentry_sdk/integrations/excepthook.py new file mode 100644 index 0000000000000000000000000000000000000000..640931999014e3c74d961fcafb4ab7a8184afd15 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/excepthook.py @@ -0,0 +1,81 @@ +import sys + +import sentry_sdk +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, +) +from sentry_sdk.integrations import Integration + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Callable + from typing import Any + from typing import Type + from typing import Optional + + from types import TracebackType + + Excepthook = Callable[ + [Type[BaseException], BaseException, Optional[TracebackType]], + Any, + ] + + +class ExcepthookIntegration(Integration): + identifier = "excepthook" + + always_run = False + + def __init__(self, always_run: bool = False) -> None: + if not isinstance(always_run, bool): + raise ValueError( + "Invalid value for always_run: %s (must be type boolean)" + % (always_run,) + ) + self.always_run = always_run + + @staticmethod + def setup_once() -> None: + sys.excepthook = _make_excepthook(sys.excepthook) + + +def _make_excepthook(old_excepthook: "Excepthook") -> "Excepthook": + def sentry_sdk_excepthook( + type_: "Type[BaseException]", + value: BaseException, + traceback: "Optional[TracebackType]", + ) -> None: + integration = sentry_sdk.get_client().get_integration(ExcepthookIntegration) + + # Note: If we replace this with ensure_integration_enabled then + # we break the exceptiongroup backport; + # See: https://github.com/getsentry/sentry-python/issues/3097 + if integration is None: + return old_excepthook(type_, value, traceback) + + if _should_send(integration.always_run): + with capture_internal_exceptions(): + event, hint = event_from_exception( + (type_, value, traceback), + client_options=sentry_sdk.get_client().options, + mechanism={"type": "excepthook", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return old_excepthook(type_, value, traceback) + + return sentry_sdk_excepthook + + +def _should_send(always_run: bool = False) -> bool: + if always_run: + return True + + if hasattr(sys, "ps1"): + # Disable the excepthook for interactive Python shells, otherwise + # every typo gets sent to Sentry. + return False + + return True diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/executing.py b/lib/python3.12/site-packages/sentry_sdk/integrations/executing.py new file mode 100644 index 0000000000000000000000000000000000000000..c5aa5226677d2dff3842e944a528cd1c84f42c49 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/executing.py @@ -0,0 +1,66 @@ +import sentry_sdk +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.scope import add_global_event_processor +from sentry_sdk.utils import walk_exception_chain, iter_stacks + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional + + from sentry_sdk._types import Event, Hint + +try: + import executing +except ImportError: + raise DidNotEnable("executing is not installed") + + +class ExecutingIntegration(Integration): + identifier = "executing" + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def add_executing_info( + event: "Event", hint: "Optional[Hint]" + ) -> "Optional[Event]": + if sentry_sdk.get_client().get_integration(ExecutingIntegration) is None: + return event + + if hint is None: + return event + + exc_info = hint.get("exc_info", None) + + if exc_info is None: + return event + + exception = event.get("exception", None) + + if exception is None: + return event + + values = exception.get("values", None) + + if values is None: + return event + + for exception, (_exc_type, _exc_value, exc_tb) in zip( + reversed(values), walk_exception_chain(exc_info) + ): + sentry_frames = [ + frame + for frame in exception.get("stacktrace", {}).get("frames", []) + if frame.get("function") + ] + tbs = list(iter_stacks(exc_tb)) + if len(sentry_frames) != len(tbs): + continue + + for sentry_frame, tb in zip(sentry_frames, tbs): + frame = tb.tb_frame + source = executing.Source.for_frame(frame) + sentry_frame["function"] = source.code_qualname(frame.f_code) + + return event diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/falcon.py b/lib/python3.12/site-packages/sentry_sdk/integrations/falcon.py new file mode 100644 index 0000000000000000000000000000000000000000..158b4e61aa255b7f0c8036649406111bb4a7e460 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/falcon.py @@ -0,0 +1,259 @@ +import sentry_sdk +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import RequestExtractor +from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware +from sentry_sdk.tracing import SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + parse_version, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Optional + + from sentry_sdk._types import Event, EventProcessor + +# In Falcon 3.0 `falcon.api_helpers` is renamed to `falcon.app_helpers` +# and `falcon.API` to `falcon.App` + +try: + import falcon # type: ignore + + from falcon import __version__ as FALCON_VERSION +except ImportError: + raise DidNotEnable("Falcon not installed") + +try: + import falcon.app_helpers # type: ignore + + falcon_helpers = falcon.app_helpers + falcon_app_class = falcon.App + FALCON3 = True +except ImportError: + import falcon.api_helpers # type: ignore + + falcon_helpers = falcon.api_helpers + falcon_app_class = falcon.API + FALCON3 = False + + +_FALCON_UNSET: "Optional[object]" = None +if FALCON3: # falcon.request._UNSET is only available in Falcon 3.0+ + with capture_internal_exceptions(): + from falcon.request import _UNSET as _FALCON_UNSET # type: ignore[import-not-found, no-redef] + + +class FalconRequestExtractor(RequestExtractor): + def env(self) -> "Dict[str, Any]": + return self.request.env + + def cookies(self) -> "Dict[str, Any]": + return self.request.cookies + + def form(self) -> None: + return None # No such concept in Falcon + + def files(self) -> None: + return None # No such concept in Falcon + + def raw_data(self) -> "Optional[str]": + # As request data can only be read once we won't make this available + # to Sentry. Just send back a dummy string in case there was a + # content length. + # TODO(jmagnusson): Figure out if there's a way to support this + content_length = self.content_length() + if content_length > 0: + return "[REQUEST_CONTAINING_RAW_DATA]" + else: + return None + + def json(self) -> "Optional[Dict[str, Any]]": + # fallback to cached_media = None if self.request._media is not available + cached_media = None + with capture_internal_exceptions(): + # self.request._media is the cached self.request.media + # value. It is only available if self.request.media + # has already been accessed. Therefore, reading + # self.request._media will not exhaust the raw request + # stream (self.request.bounded_stream) because it has + # already been read if self.request._media is set. + cached_media = self.request._media + + if cached_media is not _FALCON_UNSET: + return cached_media + + return None + + +class SentryFalconMiddleware: + """Captures exceptions in Falcon requests and send to Sentry""" + + def process_request( + self, req: "Any", resp: "Any", *args: "Any", **kwargs: "Any" + ) -> None: + integration = sentry_sdk.get_client().get_integration(FalconIntegration) + if integration is None: + return + + scope = sentry_sdk.get_isolation_scope() + scope._name = "falcon" + scope.add_event_processor(_make_request_event_processor(req, integration)) + + +TRANSACTION_STYLE_VALUES = ("uri_template", "path") + + +class FalconIntegration(Integration): + identifier = "falcon" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__(self, transaction_style: str = "uri_template") -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + + @staticmethod + def setup_once() -> None: + version = parse_version(FALCON_VERSION) + _check_minimum_version(FalconIntegration, version) + + _patch_wsgi_app() + _patch_handle_exception() + _patch_prepare_middleware() + + +def _patch_wsgi_app() -> None: + original_wsgi_app = falcon_app_class.__call__ + + def sentry_patched_wsgi_app( + self: "falcon.API", env: "Any", start_response: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(FalconIntegration) + if integration is None: + return original_wsgi_app(self, env, start_response) + + sentry_wrapped = SentryWsgiMiddleware( + lambda envi, start_resp: original_wsgi_app(self, envi, start_resp), + span_origin=FalconIntegration.origin, + ) + + return sentry_wrapped(env, start_response) + + falcon_app_class.__call__ = sentry_patched_wsgi_app + + +def _patch_handle_exception() -> None: + original_handle_exception = falcon_app_class._handle_exception + + @ensure_integration_enabled(FalconIntegration, original_handle_exception) + def sentry_patched_handle_exception(self: "falcon.API", *args: "Any") -> "Any": + # NOTE(jmagnusson): falcon 2.0 changed falcon.API._handle_exception + # method signature from `(ex, req, resp, params)` to + # `(req, resp, ex, params)` + ex = response = None + with capture_internal_exceptions(): + ex = next(argument for argument in args if isinstance(argument, Exception)) + response = next( + argument for argument in args if isinstance(argument, falcon.Response) + ) + + was_handled = original_handle_exception(self, *args) + + if ex is None or response is None: + # Both ex and response should have a non-None value at this point; otherwise, + # there is an error with the SDK that will have been captured in the + # capture_internal_exceptions block above. + return was_handled + + if _exception_leads_to_http_5xx(ex, response): + event, hint = event_from_exception( + ex, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "falcon", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return was_handled + + falcon_app_class._handle_exception = sentry_patched_handle_exception + + +def _patch_prepare_middleware() -> None: + original_prepare_middleware = falcon_helpers.prepare_middleware + + def sentry_patched_prepare_middleware( + middleware: "Any" = None, + independent_middleware: "Any" = False, + asgi: bool = False, + ) -> "Any": + if asgi: + # We don't support ASGI Falcon apps, so we don't patch anything here + return original_prepare_middleware(middleware, independent_middleware, asgi) + + integration = sentry_sdk.get_client().get_integration(FalconIntegration) + if integration is not None: + middleware = [SentryFalconMiddleware()] + (middleware or []) + + # We intentionally omit the asgi argument here, since the default is False anyways, + # and this way, we remain backwards-compatible with pre-3.0.0 Falcon versions. + return original_prepare_middleware(middleware, independent_middleware) + + falcon_helpers.prepare_middleware = sentry_patched_prepare_middleware + + +def _exception_leads_to_http_5xx(ex: Exception, response: "falcon.Response") -> bool: + is_server_error = isinstance(ex, falcon.HTTPError) and (ex.status or "").startswith( + "5" + ) + is_unhandled_error = not isinstance( + ex, (falcon.HTTPError, falcon.http_status.HTTPStatus) + ) + + # We only check the HTTP status on Falcon 3 because in Falcon 2, the status on the response + # at the stage where we capture it is listed as 200, even though we would expect to see a 500 + # status. Since at the time of this change, Falcon 2 is ca. 4 years old, we have decided to + # only perform this check on Falcon 3+, despite the risk that some handled errors might be + # reported to Sentry as unhandled on Falcon 2. + return (is_server_error or is_unhandled_error) and ( + not FALCON3 or _has_http_5xx_status(response) + ) + + +def _has_http_5xx_status(response: "falcon.Response") -> bool: + return response.status.startswith("5") + + +def _set_transaction_name_and_source( + event: "Event", transaction_style: str, request: "falcon.Request" +) -> None: + name_for_style = { + "uri_template": request.uri_template, + "path": request.path, + } + event["transaction"] = name_for_style[transaction_style] + event["transaction_info"] = {"source": SOURCE_FOR_STYLE[transaction_style]} + + +def _make_request_event_processor( + req: "falcon.Request", integration: "FalconIntegration" +) -> "EventProcessor": + def event_processor(event: "Event", hint: "dict[str, Any]") -> "Event": + _set_transaction_name_and_source(event, integration.transaction_style, req) + + with capture_internal_exceptions(): + FalconRequestExtractor(req).extract_into_event(event) + + return event + + return event_processor diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/fastapi.py b/lib/python3.12/site-packages/sentry_sdk/integrations/fastapi.py new file mode 100644 index 0000000000000000000000000000000000000000..66f73ea4e0147d11f94ec1569f37415dd21d984c --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/fastapi.py @@ -0,0 +1,136 @@ +import asyncio +from copy import deepcopy +from functools import wraps + +import sentry_sdk +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource +from sentry_sdk.utils import transaction_from_function + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Dict + from sentry_sdk._types import Event + +try: + from sentry_sdk.integrations.starlette import ( + StarletteIntegration, + StarletteRequestExtractor, + ) +except DidNotEnable: + raise DidNotEnable("Starlette is not installed") + +try: + import fastapi # type: ignore +except ImportError: + raise DidNotEnable("FastAPI is not installed") + + +_DEFAULT_TRANSACTION_NAME = "generic FastAPI request" + + +class FastApiIntegration(StarletteIntegration): + identifier = "fastapi" + + @staticmethod + def setup_once() -> None: + patch_get_request_handler() + + +def _set_transaction_name_and_source( + scope: "sentry_sdk.Scope", transaction_style: str, request: "Any" +) -> None: + name = "" + + if transaction_style == "endpoint": + endpoint = request.scope.get("endpoint") + if endpoint: + name = transaction_from_function(endpoint) or "" + + elif transaction_style == "url": + route = request.scope.get("route") + if route: + path = getattr(route, "path", None) + if path is not None: + name = path + + if not name: + name = _DEFAULT_TRANSACTION_NAME + source = TransactionSource.ROUTE + else: + source = SOURCE_FOR_STYLE[transaction_style] + + scope.set_transaction_name(name, source=source) + + +def patch_get_request_handler() -> None: + old_get_request_handler = fastapi.routing.get_request_handler + + def _sentry_get_request_handler(*args: "Any", **kwargs: "Any") -> "Any": + dependant = kwargs.get("dependant") + if ( + dependant + and dependant.call is not None + and not asyncio.iscoroutinefunction(dependant.call) + ): + old_call = dependant.call + + @wraps(old_call) + def _sentry_call(*args: "Any", **kwargs: "Any") -> "Any": + current_scope = sentry_sdk.get_current_scope() + if current_scope.transaction is not None: + current_scope.transaction.update_active_thread() + + sentry_scope = sentry_sdk.get_isolation_scope() + if sentry_scope.profile is not None: + sentry_scope.profile.update_active_thread_id() + + return old_call(*args, **kwargs) + + dependant.call = _sentry_call + + old_app = old_get_request_handler(*args, **kwargs) + + async def _sentry_app(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(FastApiIntegration) + if integration is None: + return await old_app(*args, **kwargs) + + request = args[0] + + _set_transaction_name_and_source( + sentry_sdk.get_current_scope(), integration.transaction_style, request + ) + sentry_scope = sentry_sdk.get_isolation_scope() + extractor = StarletteRequestExtractor(request) + info = await extractor.extract_request_info() + + def _make_request_event_processor( + req: "Any", integration: "Any" + ) -> "Callable[[Event, Dict[str, Any]], Event]": + def event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event": + # Extract information from request + request_info = event.get("request", {}) + if info: + if "cookies" in info and should_send_default_pii(): + request_info["cookies"] = info["cookies"] + if "data" in info: + request_info["data"] = info["data"] + event["request"] = deepcopy(request_info) + + return event + + return event_processor + + sentry_scope._name = FastApiIntegration.identifier + sentry_scope.add_event_processor( + _make_request_event_processor(request, integration) + ) + + return await old_app(*args, **kwargs) + + return _sentry_app + + fastapi.routing.get_request_handler = _sentry_get_request_handler diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/flask.py b/lib/python3.12/site-packages/sentry_sdk/integrations/flask.py new file mode 100644 index 0000000000000000000000000000000000000000..9adf8d51e8b4822c20916a827d2c1c0a17b2a46e --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/flask.py @@ -0,0 +1,265 @@ +import sentry_sdk +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.integrations._wsgi_common import ( + DEFAULT_HTTP_METHODS_TO_CAPTURE, + RequestExtractor, +) +from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + package_version, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Dict, Union + + from sentry_sdk._types import Event, EventProcessor + from sentry_sdk.integrations.wsgi import _ScopedResponse + from werkzeug.datastructures import FileStorage, ImmutableMultiDict + + +try: + import flask_login # type: ignore +except ImportError: + flask_login = None + +try: + from flask import Flask, Request # type: ignore + from flask import request as flask_request + from flask.signals import ( + before_render_template, + got_request_exception, + request_started, + ) + from markupsafe import Markup +except ImportError: + raise DidNotEnable("Flask is not installed") + +try: + import blinker # noqa +except ImportError: + raise DidNotEnable("blinker is not installed") + +TRANSACTION_STYLE_VALUES = ("endpoint", "url") + + +class FlaskIntegration(Integration): + identifier = "flask" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__( + self, + transaction_style: str = "endpoint", + http_methods_to_capture: "tuple[str, ...]" = DEFAULT_HTTP_METHODS_TO_CAPTURE, + ) -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + self.http_methods_to_capture = tuple(map(str.upper, http_methods_to_capture)) + + @staticmethod + def setup_once() -> None: + try: + from quart import Quart # type: ignore + + if Flask == Quart: + # This is Quart masquerading as Flask, don't enable the Flask + # integration. See https://github.com/getsentry/sentry-python/issues/2709 + raise DidNotEnable( + "This is not a Flask app but rather Quart pretending to be Flask" + ) + except ImportError: + pass + + version = package_version("flask") + _check_minimum_version(FlaskIntegration, version) + + before_render_template.connect(_add_sentry_trace) + request_started.connect(_request_started) + got_request_exception.connect(_capture_exception) + + old_app = Flask.__call__ + + def sentry_patched_wsgi_app( + self: "Any", environ: "Dict[str, str]", start_response: "Callable[..., Any]" + ) -> "_ScopedResponse": + if sentry_sdk.get_client().get_integration(FlaskIntegration) is None: + return old_app(self, environ, start_response) + + integration = sentry_sdk.get_client().get_integration(FlaskIntegration) + + middleware = SentryWsgiMiddleware( + lambda *a, **kw: old_app(self, *a, **kw), + span_origin=FlaskIntegration.origin, + http_methods_to_capture=( + integration.http_methods_to_capture + if integration + else DEFAULT_HTTP_METHODS_TO_CAPTURE + ), + ) + return middleware(environ, start_response) + + Flask.__call__ = sentry_patched_wsgi_app + + +def _add_sentry_trace( + sender: "Flask", template: "Any", context: "Dict[str, Any]", **extra: "Any" +) -> None: + if "sentry_trace" in context: + return + + scope = sentry_sdk.get_current_scope() + trace_meta = Markup(scope.trace_propagation_meta()) + context["sentry_trace"] = trace_meta # for backwards compatibility + context["sentry_trace_meta"] = trace_meta + + +def _set_transaction_name_and_source( + scope: "sentry_sdk.Scope", transaction_style: str, request: "Request" +) -> None: + try: + name_for_style = { + "url": request.url_rule.rule, + "endpoint": request.url_rule.endpoint, + } + scope.set_transaction_name( + name_for_style[transaction_style], + source=SOURCE_FOR_STYLE[transaction_style], + ) + except Exception: + pass + + +def _request_started(app: "Flask", **kwargs: "Any") -> None: + integration = sentry_sdk.get_client().get_integration(FlaskIntegration) + if integration is None: + return + + request = flask_request._get_current_object() + + # Set the transaction name and source here, + # but rely on WSGI middleware to actually start the transaction + _set_transaction_name_and_source( + sentry_sdk.get_current_scope(), integration.transaction_style, request + ) + + scope = sentry_sdk.get_isolation_scope() + evt_processor = _make_request_event_processor(app, request, integration) + scope.add_event_processor(evt_processor) + + +class FlaskRequestExtractor(RequestExtractor): + def env(self) -> "Dict[str, str]": + return self.request.environ + + def cookies(self) -> "Dict[Any, Any]": + return { + k: v[0] if isinstance(v, list) and len(v) == 1 else v + for k, v in self.request.cookies.items() + } + + def raw_data(self) -> bytes: + return self.request.get_data() + + def form(self) -> "ImmutableMultiDict[str, Any]": + return self.request.form + + def files(self) -> "ImmutableMultiDict[str, Any]": + return self.request.files + + def is_json(self) -> bool: + return self.request.is_json + + def json(self) -> "Any": + return self.request.get_json(silent=True) + + def size_of_file(self, file: "FileStorage") -> int: + return file.content_length + + +def _make_request_event_processor( + app: "Flask", request: "Callable[[], Request]", integration: "FlaskIntegration" +) -> "EventProcessor": + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + # if the request is gone we are fine not logging the data from + # it. This might happen if the processor is pushed away to + # another thread. + if request is None: + return event + + with capture_internal_exceptions(): + FlaskRequestExtractor(request).extract_into_event(event) + + if should_send_default_pii(): + with capture_internal_exceptions(): + _add_user_to_event(event) + + return event + + return inner + + +@ensure_integration_enabled(FlaskIntegration) +def _capture_exception( + sender: "Flask", exception: "Union[ValueError, BaseException]", **kwargs: "Any" +) -> None: + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "flask", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) + + +def _add_user_to_event(event: "Event") -> None: + if flask_login is None: + return + + user = flask_login.current_user + if user is None: + return + + with capture_internal_exceptions(): + # Access this object as late as possible as accessing the user + # is relatively costly + + user_info = event.setdefault("user", {}) + + try: + user_info.setdefault("id", user.get_id()) + # TODO: more configurable user attrs here + except AttributeError: + # might happen if: + # - flask_login could not be imported + # - flask_login is not configured + # - no user is logged in + pass + + # The following attribute accesses are ineffective for the general + # Flask-Login case, because the User interface of Flask-Login does not + # care about anything but the ID. However, Flask-User (based on + # Flask-Login) documents a few optional extra attributes. + # + # https://github.com/lingthio/Flask-User/blob/a379fa0a281789618c484b459cb41236779b95b1/docs/source/data_models.rst#fixed-data-model-property-names + + try: + user_info.setdefault("email", user.email) + except Exception: + pass + + try: + user_info.setdefault("username", user.username) + except Exception: + pass diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/gcp.py b/lib/python3.12/site-packages/sentry_sdk/integrations/gcp.py new file mode 100644 index 0000000000000000000000000000000000000000..994d38f932615e2ee11c13ece8d8086bc396a7b7 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/gcp.py @@ -0,0 +1,234 @@ +import functools +import sys +from copy import deepcopy +from datetime import datetime, timedelta, timezone +from os import environ + +import sentry_sdk +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.integrations import Integration +from sentry_sdk.integrations._wsgi_common import _filter_headers +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + AnnotatedValue, + capture_internal_exceptions, + event_from_exception, + logger, + TimeoutThread, + reraise, +) + +from typing import TYPE_CHECKING + +# Constants +TIMEOUT_WARNING_BUFFER = 1.5 # Buffer time required to send timeout warning to Sentry +MILLIS_TO_SECONDS = 1000.0 + +if TYPE_CHECKING: + from typing import Any + from typing import TypeVar + from typing import Callable + from typing import Optional + + from sentry_sdk._types import EventProcessor, Event, Hint + + F = TypeVar("F", bound=Callable[..., Any]) + + +def _wrap_func(func: "F") -> "F": + @functools.wraps(func) + def sentry_func( + functionhandler: "Any", gcp_event: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + client = sentry_sdk.get_client() + + integration = client.get_integration(GcpIntegration) + if integration is None: + return func(functionhandler, gcp_event, *args, **kwargs) + + configured_time = environ.get("FUNCTION_TIMEOUT_SEC") + if not configured_time: + logger.debug( + "The configured timeout could not be fetched from Cloud Functions configuration." + ) + return func(functionhandler, gcp_event, *args, **kwargs) + + configured_time = int(configured_time) + + initial_time = datetime.now(timezone.utc) + + with sentry_sdk.isolation_scope() as scope: + with capture_internal_exceptions(): + scope.clear_breadcrumbs() + scope.add_event_processor( + _make_request_event_processor( + gcp_event, configured_time, initial_time + ) + ) + scope.set_tag("gcp_region", environ.get("FUNCTION_REGION")) + timeout_thread = None + if ( + integration.timeout_warning + and configured_time > TIMEOUT_WARNING_BUFFER + ): + waiting_time = configured_time - TIMEOUT_WARNING_BUFFER + + timeout_thread = TimeoutThread( + waiting_time, + configured_time, + isolation_scope=scope, + current_scope=sentry_sdk.get_current_scope(), + ) + + # Starting the thread to raise timeout warning exception + timeout_thread.start() + + headers = {} + if hasattr(gcp_event, "headers"): + headers = gcp_event.headers + + transaction = continue_trace( + headers, + op=OP.FUNCTION_GCP, + name=environ.get("FUNCTION_NAME", ""), + source=TransactionSource.COMPONENT, + origin=GcpIntegration.origin, + ) + sampling_context = { + "gcp_env": { + "function_name": environ.get("FUNCTION_NAME"), + "function_entry_point": environ.get("ENTRY_POINT"), + "function_identity": environ.get("FUNCTION_IDENTITY"), + "function_region": environ.get("FUNCTION_REGION"), + "function_project": environ.get("GCP_PROJECT"), + }, + "gcp_event": gcp_event, + } + with sentry_sdk.start_transaction( + transaction, custom_sampling_context=sampling_context + ): + try: + return func(functionhandler, gcp_event, *args, **kwargs) + except Exception: + exc_info = sys.exc_info() + sentry_event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "gcp", "handled": False}, + ) + sentry_sdk.capture_event(sentry_event, hint=hint) + reraise(*exc_info) + finally: + if timeout_thread: + timeout_thread.stop() + # Flush out the event queue + client.flush() + + return sentry_func # type: ignore + + +class GcpIntegration(Integration): + identifier = "gcp" + origin = f"auto.function.{identifier}" + + def __init__(self, timeout_warning: bool = False) -> None: + self.timeout_warning = timeout_warning + + @staticmethod + def setup_once() -> None: + import __main__ as gcp_functions + + if not hasattr(gcp_functions, "worker_v1"): + logger.warning( + "GcpIntegration currently supports only Python 3.7 runtime environment." + ) + return + + worker1 = gcp_functions.worker_v1 + + worker1.FunctionHandler.invoke_user_function = _wrap_func( + worker1.FunctionHandler.invoke_user_function + ) + + +def _make_request_event_processor( + gcp_event: "Any", configured_timeout: "Any", initial_time: "Any" +) -> "EventProcessor": + def event_processor(event: "Event", hint: "Hint") -> "Optional[Event]": + final_time = datetime.now(timezone.utc) + time_diff = final_time - initial_time + + execution_duration_in_millis = time_diff / timedelta(milliseconds=1) + + extra = event.setdefault("extra", {}) + extra["google cloud functions"] = { + "function_name": environ.get("FUNCTION_NAME"), + "function_entry_point": environ.get("ENTRY_POINT"), + "function_identity": environ.get("FUNCTION_IDENTITY"), + "function_region": environ.get("FUNCTION_REGION"), + "function_project": environ.get("GCP_PROJECT"), + "execution_duration_in_millis": execution_duration_in_millis, + "configured_timeout_in_seconds": configured_timeout, + } + + extra["google cloud logs"] = { + "url": _get_google_cloud_logs_url(final_time), + } + + request = event.get("request", {}) + + request["url"] = "gcp:///{}".format(environ.get("FUNCTION_NAME")) + + if hasattr(gcp_event, "method"): + request["method"] = gcp_event.method + + if hasattr(gcp_event, "query_string"): + request["query_string"] = gcp_event.query_string.decode("utf-8") + + if hasattr(gcp_event, "headers"): + request["headers"] = _filter_headers(gcp_event.headers) + + if should_send_default_pii(): + if hasattr(gcp_event, "data"): + request["data"] = gcp_event.data + else: + if hasattr(gcp_event, "data"): + # Unfortunately couldn't find a way to get structured body from GCP + # event. Meaning every body is unstructured to us. + request["data"] = AnnotatedValue.removed_because_raw_data() + + event["request"] = deepcopy(request) + + return event + + return event_processor + + +def _get_google_cloud_logs_url(final_time: "datetime") -> str: + """ + Generates a Google Cloud Logs console URL based on the environment variables + Arguments: + final_time {datetime} -- Final time + Returns: + str -- Google Cloud Logs Console URL to logs. + """ + hour_ago = final_time - timedelta(hours=1) + formatstring = "%Y-%m-%dT%H:%M:%SZ" + + url = ( + "https://console.cloud.google.com/logs/viewer?project={project}&resource=cloud_function" + "%2Ffunction_name%2F{function_name}%2Fregion%2F{region}&minLogLevel=0&expandAll=false" + "×tamp={timestamp_end}&customFacets=&limitCustomFacetWidth=true" + "&dateRangeStart={timestamp_start}&dateRangeEnd={timestamp_end}" + "&interval=PT1H&scrollTimestamp={timestamp_end}" + ).format( + project=environ.get("GCP_PROJECT"), + function_name=environ.get("FUNCTION_NAME"), + region=environ.get("FUNCTION_REGION"), + timestamp_end=final_time.strftime(formatstring), + timestamp_start=hour_ago.strftime(formatstring), + ) + + return url diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/gnu_backtrace.py b/lib/python3.12/site-packages/sentry_sdk/integrations/gnu_backtrace.py new file mode 100644 index 0000000000000000000000000000000000000000..dbadf42088ef0e50b7408c86327fa85593ac8642 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/gnu_backtrace.py @@ -0,0 +1,96 @@ +import re + +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import add_global_event_processor +from sentry_sdk.utils import capture_internal_exceptions + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from sentry_sdk._types import Event + +# function is everything between index at @ +# and then we match on the @ plus the hex val +FUNCTION_RE = r"[^@]+?" +HEX_ADDRESS = r"\s+@\s+0x[0-9a-fA-F]+" + +FRAME_RE = r""" +^(?P\d+)\.\s+(?P{FUNCTION_RE}){HEX_ADDRESS}(?:\s+in\s+(?P.+))?$ +""".format( + FUNCTION_RE=FUNCTION_RE, + HEX_ADDRESS=HEX_ADDRESS, +) + +FRAME_RE = re.compile(FRAME_RE, re.MULTILINE | re.VERBOSE) + + +class GnuBacktraceIntegration(Integration): + identifier = "gnu_backtrace" + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def process_gnu_backtrace(event: "Event", hint: "dict[str, Any]") -> "Event": + with capture_internal_exceptions(): + return _process_gnu_backtrace(event, hint) + + +def _process_gnu_backtrace(event: "Event", hint: "dict[str, Any]") -> "Event": + if sentry_sdk.get_client().get_integration(GnuBacktraceIntegration) is None: + return event + + exc_info = hint.get("exc_info", None) + + if exc_info is None: + return event + + exception = event.get("exception", None) + + if exception is None: + return event + + values = exception.get("values", None) + + if values is None: + return event + + for exception in values: + frames = exception.get("stacktrace", {}).get("frames", []) + if not frames: + continue + + msg = exception.get("value", None) + if not msg: + continue + + additional_frames = [] + new_msg = [] + + for line in msg.splitlines(): + match = FRAME_RE.match(line) + if match: + additional_frames.append( + ( + int(match.group("index")), + { + "package": match.group("package") or None, + "function": match.group("function") or None, + "platform": "native", + }, + ) + ) + else: + # Put garbage lines back into message, not sure what to do with them. + new_msg.append(line) + + if additional_frames: + additional_frames.sort(key=lambda x: -x[0]) + for _, frame in additional_frames: + frames.append(frame) + + new_msg.append("") + exception["value"] = "\n".join(new_msg) + + return event diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__init__.py b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..27a42f4f6a0ba2e47968a1e76a4d28cef870d2bc --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__init__.py @@ -0,0 +1,370 @@ +from functools import wraps +from typing import ( + Any, + AsyncIterator, + Callable, + Iterator, + List, +) + +import sentry_sdk +from sentry_sdk.ai.utils import get_start_span_function +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.tracing import SPANSTATUS + + +try: + from google.genai.models import Models, AsyncModels +except ImportError: + raise DidNotEnable("google-genai not installed") + + +from .consts import IDENTIFIER, ORIGIN, GEN_AI_SYSTEM +from .utils import ( + set_span_data_for_request, + set_span_data_for_response, + _capture_exception, + prepare_generate_content_args, + prepare_embed_content_args, + set_span_data_for_embed_request, + set_span_data_for_embed_response, +) +from .streaming import ( + set_span_data_for_streaming_response, + accumulate_streaming_response, +) + + +class GoogleGenAIIntegration(Integration): + identifier = IDENTIFIER + origin = ORIGIN + + def __init__(self: "GoogleGenAIIntegration", include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + # Patch sync methods + Models.generate_content = _wrap_generate_content(Models.generate_content) + Models.generate_content_stream = _wrap_generate_content_stream( + Models.generate_content_stream + ) + Models.embed_content = _wrap_embed_content(Models.embed_content) + + # Patch async methods + AsyncModels.generate_content = _wrap_async_generate_content( + AsyncModels.generate_content + ) + AsyncModels.generate_content_stream = _wrap_async_generate_content_stream( + AsyncModels.generate_content_stream + ) + AsyncModels.embed_content = _wrap_async_embed_content(AsyncModels.embed_content) + + +def _wrap_generate_content_stream(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_generate_content_stream( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return f(self, *args, **kwargs) + + _model, contents, model_name = prepare_generate_content_args(args, kwargs) + + span = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name="invoke_agent", + origin=ORIGIN, + ) + span.__enter__() + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + set_span_data_for_request(span, integration, model_name, contents, kwargs) + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + + chat_span = sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) + chat_span.__enter__() + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) + chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + + try: + stream = f(self, *args, **kwargs) + + # Create wrapper iterator to accumulate responses + def new_iterator() -> "Iterator[Any]": + chunks: "List[Any]" = [] + try: + for chunk in stream: + chunks.append(chunk) + yield chunk + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + finally: + # Accumulate all chunks and set final response data on spans + if chunks: + accumulated_response = accumulate_streaming_response(chunks) + set_span_data_for_streaming_response( + chat_span, integration, accumulated_response + ) + set_span_data_for_streaming_response( + span, integration, accumulated_response + ) + chat_span.__exit__(None, None, None) + span.__exit__(None, None, None) + + return new_iterator() + + except Exception as exc: + _capture_exception(exc) + chat_span.__exit__(None, None, None) + span.__exit__(None, None, None) + raise + + return new_generate_content_stream + + +def _wrap_async_generate_content_stream( + f: "Callable[..., Any]", +) -> "Callable[..., Any]": + @wraps(f) + async def new_async_generate_content_stream( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + _model, contents, model_name = prepare_generate_content_args(args, kwargs) + + span = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name="invoke_agent", + origin=ORIGIN, + ) + span.__enter__() + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + set_span_data_for_request(span, integration, model_name, contents, kwargs) + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + + chat_span = sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) + chat_span.__enter__() + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) + chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + + try: + stream = await f(self, *args, **kwargs) + + # Create wrapper async iterator to accumulate responses + async def new_async_iterator() -> "AsyncIterator[Any]": + chunks: "List[Any]" = [] + try: + async for chunk in stream: + chunks.append(chunk) + yield chunk + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + finally: + # Accumulate all chunks and set final response data on spans + if chunks: + accumulated_response = accumulate_streaming_response(chunks) + set_span_data_for_streaming_response( + chat_span, integration, accumulated_response + ) + set_span_data_for_streaming_response( + span, integration, accumulated_response + ) + chat_span.__exit__(None, None, None) + span.__exit__(None, None, None) + + return new_async_iterator() + + except Exception as exc: + _capture_exception(exc) + chat_span.__exit__(None, None, None) + span.__exit__(None, None, None) + raise + + return new_async_generate_content_stream + + +def _wrap_generate_content(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return f(self, *args, **kwargs) + + model, contents, model_name = prepare_generate_content_args(args, kwargs) + + with get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name="invoke_agent", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + set_span_data_for_request(span, integration, model_name, contents, kwargs) + + with sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) as chat_span: + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + set_span_data_for_request( + chat_span, integration, model_name, contents, kwargs + ) + + try: + response = f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + + set_span_data_for_response(chat_span, integration, response) + set_span_data_for_response(span, integration, response) + + return response + + return new_generate_content + + +def _wrap_async_generate_content(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + async def new_async_generate_content( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + model, contents, model_name = prepare_generate_content_args(args, kwargs) + + with get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name="invoke_agent", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + set_span_data_for_request(span, integration, model_name, contents, kwargs) + + with sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) as chat_span: + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_request( + chat_span, integration, model_name, contents, kwargs + ) + try: + response = await f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + + set_span_data_for_response(chat_span, integration, response) + set_span_data_for_response(span, integration, response) + + return response + + return new_async_generate_content + + +def _wrap_embed_content(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_embed_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return f(self, *args, **kwargs) + + model_name, contents = prepare_embed_content_args(args, kwargs) + + with sentry_sdk.start_span( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) + + try: + response = f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + + set_span_data_for_embed_response(span, integration, response) + + return response + + return new_embed_content + + +def _wrap_async_embed_content(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + async def new_async_embed_content( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + model_name, contents = prepare_embed_content_args(args, kwargs) + + with sentry_sdk.start_span( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) + + try: + response = await f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + + set_span_data_for_embed_response(span, integration, response) + + return response + + return new_async_embed_content diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6fcfb6041c0239966367fd2fd3625879e51e621 Binary files /dev/null and b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/consts.cpython-312.pyc b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/consts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a91d175b4b9f3e8d9e60ae38ffcd9c6a1af0b7af Binary files /dev/null and b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/consts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/streaming.cpython-312.pyc b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/streaming.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d576e190b1771b0b2858df3817460a65d7a7cb3 Binary files /dev/null and b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/streaming.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f91c46b7d8f6b8fb5217b290653b7f6558cf2145 Binary files /dev/null and b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/consts.py b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/consts.py new file mode 100644 index 0000000000000000000000000000000000000000..5b53ebf0e2a3c40d07c35adfc1bbbcfcd47d7576 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/consts.py @@ -0,0 +1,16 @@ +GEN_AI_SYSTEM = "gcp.gemini" + +# Mapping of tool attributes to their descriptions +# These are all tools that are available in the Google GenAI API +TOOL_ATTRIBUTES_MAP = { + "google_search_retrieval": "Google Search retrieval tool", + "google_search": "Google Search tool", + "retrieval": "Retrieval tool", + "enterprise_web_search": "Enterprise web search tool", + "google_maps": "Google Maps tool", + "code_execution": "Code execution tool", + "computer_use": "Computer use tool", +} + +IDENTIFIER = "google_genai" +ORIGIN = f"auto.ai.{IDENTIFIER}" diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/streaming.py b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/streaming.py new file mode 100644 index 0000000000000000000000000000000000000000..8649ce2ac093cd6fb3184d3885e7d685a7bdabe4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/streaming.py @@ -0,0 +1,160 @@ +from typing import TYPE_CHECKING, Any, List, TypedDict, Optional, Union + +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.consts import SPANDATA +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + safe_serialize, +) +from .utils import ( + extract_tool_calls, + extract_finish_reasons, + extract_contents_text, + extract_usage_data, + UsageData, +) + +if TYPE_CHECKING: + from sentry_sdk.tracing import Span + from google.genai.types import GenerateContentResponse + + +class AccumulatedResponse(TypedDict): + id: "Optional[str]" + model: "Optional[str]" + text: str + finish_reasons: "List[str]" + tool_calls: "List[dict[str, Any]]" + usage_metadata: "Optional[UsageData]" + + +def element_wise_usage_max(self: "UsageData", other: "UsageData") -> "UsageData": + return UsageData( + input_tokens=max(self["input_tokens"], other["input_tokens"]), + output_tokens=max(self["output_tokens"], other["output_tokens"]), + input_tokens_cached=max( + self["input_tokens_cached"], other["input_tokens_cached"] + ), + output_tokens_reasoning=max( + self["output_tokens_reasoning"], other["output_tokens_reasoning"] + ), + total_tokens=max(self["total_tokens"], other["total_tokens"]), + ) + + +def accumulate_streaming_response( + chunks: "List[GenerateContentResponse]", +) -> "AccumulatedResponse": + """Accumulate streaming chunks into a single response-like object.""" + accumulated_text = [] + finish_reasons = [] + tool_calls = [] + usage_data = None + response_id = None + model = None + + for chunk in chunks: + # Extract text and tool calls + if getattr(chunk, "candidates", None): + for candidate in getattr(chunk, "candidates", []): + if hasattr(candidate, "content") and getattr( + candidate.content, "parts", [] + ): + extracted_text = extract_contents_text(candidate.content) + if extracted_text: + accumulated_text.append(extracted_text) + + extracted_finish_reasons = extract_finish_reasons(chunk) + if extracted_finish_reasons: + finish_reasons.extend(extracted_finish_reasons) + + extracted_tool_calls = extract_tool_calls(chunk) + if extracted_tool_calls: + tool_calls.extend(extracted_tool_calls) + + # Use last possible chunk, in case of interruption, and + # gracefully handle missing intermediate tokens by taking maximum + # with previous token reporting. + chunk_usage_data = extract_usage_data(chunk) + usage_data = ( + chunk_usage_data + if usage_data is None + else element_wise_usage_max(usage_data, chunk_usage_data) + ) + + accumulated_response = AccumulatedResponse( + text="".join(accumulated_text), + finish_reasons=finish_reasons, + tool_calls=tool_calls, + usage_metadata=usage_data, + id=response_id, + model=model, + ) + + return accumulated_response + + +def set_span_data_for_streaming_response( + span: "Span", integration: "Any", accumulated_response: "AccumulatedResponse" +) -> None: + """Set span data for accumulated streaming response.""" + if ( + should_send_default_pii() + and integration.include_prompts + and accumulated_response.get("text") + ): + span.set_data( + SPANDATA.GEN_AI_RESPONSE_TEXT, + safe_serialize([accumulated_response["text"]]), + ) + + if accumulated_response.get("finish_reasons"): + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, + accumulated_response["finish_reasons"], + ) + + if accumulated_response.get("tool_calls"): + span.set_data( + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + safe_serialize(accumulated_response["tool_calls"]), + ) + + if accumulated_response.get("id"): + span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"]) + if accumulated_response.get("model"): + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"]) + + if accumulated_response["usage_metadata"] is None: + return + + if accumulated_response["usage_metadata"]["input_tokens"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, + accumulated_response["usage_metadata"]["input_tokens"], + ) + + if accumulated_response["usage_metadata"]["input_tokens_cached"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, + accumulated_response["usage_metadata"]["input_tokens_cached"], + ) + + if accumulated_response["usage_metadata"]["output_tokens"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + accumulated_response["usage_metadata"]["output_tokens"], + ) + + if accumulated_response["usage_metadata"]["output_tokens_reasoning"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, + accumulated_response["usage_metadata"]["output_tokens_reasoning"], + ) + + if accumulated_response["usage_metadata"]["total_tokens"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + accumulated_response["usage_metadata"]["total_tokens"], + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/utils.py b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b2d64998435691bc1a36e718fb761abba6694262 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/utils.py @@ -0,0 +1,1001 @@ +import copy +import json +import inspect +from functools import wraps +from .consts import ORIGIN, TOOL_ATTRIBUTES_MAP, GEN_AI_SYSTEM +from sentry_sdk._types import BLOB_DATA_SUBSTITUTE +from typing import ( + cast, + TYPE_CHECKING, + Iterable, + Any, + Callable, + List, + Optional, + Union, + TypedDict, + Dict, +) + +import sentry_sdk +from sentry_sdk.ai.utils import ( + set_data_normalized, + truncate_and_annotate_messages, + normalize_message_roles, + redact_blob_message_parts, + transform_google_content_part, + get_modality_from_mime_type, +) +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + safe_serialize, +) +from google.genai.types import GenerateContentConfig, Part, Content +from itertools import chain + +if TYPE_CHECKING: + from sentry_sdk.tracing import Span + from sentry_sdk._types import TextPart + from google.genai.types import ( + GenerateContentResponse, + ContentListUnion, + ContentUnionDict, + Tool, + Model, + EmbedContentResponse, + ContentUnion, + ) + + +class UsageData(TypedDict): + """Structure for token usage data.""" + + input_tokens: int + input_tokens_cached: int + output_tokens: int + output_tokens_reasoning: int + total_tokens: int + + +def extract_usage_data( + response: "Union[GenerateContentResponse, dict[str, Any]]", +) -> "UsageData": + """Extract usage data from response into a structured format. + + Args: + response: The GenerateContentResponse object or dictionary containing usage metadata + + Returns: + UsageData: Dictionary with input_tokens, input_tokens_cached, + output_tokens, and output_tokens_reasoning fields + """ + usage_data = UsageData( + input_tokens=0, + input_tokens_cached=0, + output_tokens=0, + output_tokens_reasoning=0, + total_tokens=0, + ) + + # Handle dictionary response (from streaming) + if isinstance(response, dict): + usage = response.get("usage_metadata", {}) + if not usage: + return usage_data + + prompt_tokens = usage.get("prompt_token_count", 0) or 0 + tool_use_prompt_tokens = usage.get("tool_use_prompt_token_count", 0) or 0 + usage_data["input_tokens"] = prompt_tokens + tool_use_prompt_tokens + + cached_tokens = usage.get("cached_content_token_count", 0) or 0 + usage_data["input_tokens_cached"] = cached_tokens + + reasoning_tokens = usage.get("thoughts_token_count", 0) or 0 + usage_data["output_tokens_reasoning"] = reasoning_tokens + + candidates_tokens = usage.get("candidates_token_count", 0) or 0 + # python-genai reports output and reasoning tokens separately + # reasoning should be sub-category of output tokens + usage_data["output_tokens"] = candidates_tokens + reasoning_tokens + + total_tokens = usage.get("total_token_count", 0) or 0 + usage_data["total_tokens"] = total_tokens + + return usage_data + + if not hasattr(response, "usage_metadata"): + return usage_data + + usage = response.usage_metadata + + # Input tokens include both prompt and tool use prompt tokens + prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0 + tool_use_prompt_tokens = getattr(usage, "tool_use_prompt_token_count", 0) or 0 + usage_data["input_tokens"] = prompt_tokens + tool_use_prompt_tokens + + # Cached input tokens + cached_tokens = getattr(usage, "cached_content_token_count", 0) or 0 + usage_data["input_tokens_cached"] = cached_tokens + + # Reasoning tokens + reasoning_tokens = getattr(usage, "thoughts_token_count", 0) or 0 + usage_data["output_tokens_reasoning"] = reasoning_tokens + + # output_tokens = candidates_tokens + reasoning_tokens + # google-genai reports output and reasoning tokens separately + candidates_tokens = getattr(usage, "candidates_token_count", 0) or 0 + usage_data["output_tokens"] = candidates_tokens + reasoning_tokens + + total_tokens = getattr(usage, "total_token_count", 0) or 0 + usage_data["total_tokens"] = total_tokens + + return usage_data + + +def _capture_exception(exc: "Any") -> None: + """Capture exception with Google GenAI mechanism.""" + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "google_genai", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def get_model_name(model: "Union[str, Model]") -> str: + """Extract model name from model parameter.""" + if isinstance(model, str): + return model + # Handle case where model might be an object with a name attribute + if hasattr(model, "name"): + return str(model.name) + return str(model) + + +def extract_contents_messages(contents: "ContentListUnion") -> "List[Dict[str, Any]]": + """Extract messages from contents parameter which can have various formats. + + Returns a list of message dictionaries in the format: + - System: {"role": "system", "content": "string"} + - User/Assistant: {"role": "user"|"assistant", "content": [{"text": "...", "type": "text"}, ...]} + """ + if contents is None: + return [] + + messages = [] + + # Handle string case + if isinstance(contents, str): + return [{"role": "user", "content": contents}] + + # Handle list case - process each item (non-recursive, flatten at top level) + if isinstance(contents, list): + for item in contents: + item_messages = extract_contents_messages(item) + messages.extend(item_messages) + return messages + + # Handle dictionary case (ContentDict) + if isinstance(contents, dict): + role = contents.get("role", "user") + parts = contents.get("parts") + + if parts: + content_parts = [] + tool_messages = [] + + for part in parts: + part_result = _extract_part_content(part) + if part_result is None: + continue + + if isinstance(part_result, dict) and part_result.get("role") == "tool": + # Tool message - add separately + tool_messages.append(part_result) + else: + # Regular content part + content_parts.append(part_result) + + # Add main message if we have content parts + if content_parts: + # Normalize role: "model" -> "assistant" + normalized_role = "assistant" if role == "model" else role or "user" + messages.append({"role": normalized_role, "content": content_parts}) + + # Add tool messages + messages.extend(tool_messages) + elif "text" in contents: + # Simple text in dict + messages.append( + { + "role": role or "user", + "content": [{"text": contents["text"], "type": "text"}], + } + ) + + return messages + + # Handle Content object + if hasattr(contents, "parts") and contents.parts: + role = getattr(contents, "role", None) or "user" + content_parts = [] + tool_messages = [] + + for part in contents.parts: + part_result = _extract_part_content(part) + if part_result is None: + continue + + if isinstance(part_result, dict) and part_result.get("role") == "tool": + tool_messages.append(part_result) + else: + content_parts.append(part_result) + + if content_parts: + normalized_role = "assistant" if role == "model" else role + messages.append({"role": normalized_role, "content": content_parts}) + + messages.extend(tool_messages) + return messages + + # Handle Part object directly + part_result = _extract_part_content(contents) + if part_result: + if isinstance(part_result, dict) and part_result.get("role") == "tool": + return [part_result] + else: + return [{"role": "user", "content": [part_result]}] + + # Handle PIL.Image.Image + try: + from PIL import Image as PILImage # type: ignore[import-not-found] + + if isinstance(contents, PILImage.Image): + blob_part = _extract_pil_image(contents) + if blob_part: + return [{"role": "user", "content": [blob_part]}] + except ImportError: + pass + + # Handle File object + if hasattr(contents, "uri") and hasattr(contents, "mime_type"): + # File object + file_uri = getattr(contents, "uri", None) + mime_type = getattr(contents, "mime_type", None) + # Process if we have file_uri, even if mime_type is missing + if file_uri is not None: + # Default to empty string if mime_type is None + if mime_type is None: + mime_type = "" + + blob_part = { + "type": "uri", + "modality": get_modality_from_mime_type(mime_type), + "mime_type": mime_type, + "uri": file_uri, + } + return [{"role": "user", "content": [blob_part]}] + + # Handle direct text attribute + if hasattr(contents, "text") and contents.text: + return [ + {"role": "user", "content": [{"text": str(contents.text), "type": "text"}]} + ] + + return [] + + +def _extract_part_content(part: "Any") -> "Optional[dict[str, Any]]": + """Extract content from a Part object or dict. + + Returns: + - dict for content part (text/blob) or tool message + - None if part should be skipped + """ + if part is None: + return None + + # Handle dict Part + if isinstance(part, dict): + # Check for function_response first (tool message) + if "function_response" in part: + return _extract_tool_message_from_part(part) + + if part.get("text"): + return {"text": part["text"], "type": "text"} + + # Try using Google-specific transform for dict formats (inline_data, file_data) + result = transform_google_content_part(part) + if result is not None: + # For inline_data with bytes data, substitute the content + if "inline_data" in part: + inline_data = part["inline_data"] + if isinstance(inline_data, dict) and isinstance( + inline_data.get("data"), bytes + ): + result["content"] = BLOB_DATA_SUBSTITUTE + return result + + return None + + # Handle Part object + # Check for function_response (tool message) + if hasattr(part, "function_response") and part.function_response: + return _extract_tool_message_from_part(part) + + # Handle text + if hasattr(part, "text") and part.text: + return {"text": part.text, "type": "text"} + + # Handle file_data + if hasattr(part, "file_data") and part.file_data: + file_data = part.file_data + file_uri = getattr(file_data, "file_uri", None) + mime_type = getattr(file_data, "mime_type", None) + # Process if we have file_uri, even if mime_type is missing (consistent with dict handling) + if file_uri is not None: + # Default to empty string if mime_type is None (consistent with transform_google_content_part) + if mime_type is None: + mime_type = "" + + return { + "type": "uri", + "modality": get_modality_from_mime_type(mime_type), + "mime_type": mime_type, + "uri": file_uri, + } + + # Handle inline_data + if hasattr(part, "inline_data") and part.inline_data: + inline_data = part.inline_data + data = getattr(inline_data, "data", None) + mime_type = getattr(inline_data, "mime_type", None) + # Process if we have data, even if mime_type is missing/empty (consistent with dict handling) + if data is not None: + # Default to empty string if mime_type is None (consistent with transform_google_content_part) + if mime_type is None: + mime_type = "" + + # Handle both bytes (binary data) and str (base64-encoded data) + if isinstance(data, bytes): + content = BLOB_DATA_SUBSTITUTE + else: + # For non-bytes data (e.g., base64 strings), use as-is + content = data + + return { + "type": "blob", + "modality": get_modality_from_mime_type(mime_type), + "mime_type": mime_type, + "content": content, + } + + return None + + +def _extract_tool_message_from_part(part: "Any") -> "Optional[dict[str, Any]]": + """Extract tool message from a Part with function_response. + + Returns: + {"role": "tool", "content": {"toolCallId": "...", "toolName": "...", "output": "..."}} + or None if not a valid tool message + """ + function_response = None + + if isinstance(part, dict): + function_response = part.get("function_response") + elif hasattr(part, "function_response"): + function_response = part.function_response + + if not function_response: + return None + + # Extract fields from function_response + tool_call_id = None + tool_name = None + output = None + + if isinstance(function_response, dict): + tool_call_id = function_response.get("id") + tool_name = function_response.get("name") + response_dict = function_response.get("response", {}) + # Prefer "output" key if present, otherwise use entire response + output = response_dict.get("output", response_dict) + else: + # FunctionResponse object + tool_call_id = getattr(function_response, "id", None) + tool_name = getattr(function_response, "name", None) + response_obj = getattr(function_response, "response", None) + if response_obj is None: + response_obj = {} + if isinstance(response_obj, dict): + output = response_obj.get("output", response_obj) + else: + output = response_obj + + if not tool_name: + return None + + return { + "role": "tool", + "content": { + "toolCallId": str(tool_call_id) if tool_call_id else None, + "toolName": str(tool_name), + "output": safe_serialize(output) if output is not None else None, + }, + } + + +def _extract_pil_image(image: "Any") -> "Optional[dict[str, Any]]": + """Extract blob part from PIL.Image.Image.""" + try: + from PIL import Image as PILImage + + if not isinstance(image, PILImage.Image): + return None + + # Get format, default to JPEG + format_str = image.format or "JPEG" + suffix = format_str.lower() + mime_type = f"image/{suffix}" + + return { + "type": "blob", + "modality": get_modality_from_mime_type(mime_type), + "mime_type": mime_type, + "content": BLOB_DATA_SUBSTITUTE, + } + except Exception: + return None + + +def extract_contents_text(contents: "ContentListUnion") -> "Optional[str]": + """Extract text from contents parameter which can have various formats. + + This is a compatibility function that extracts text from messages. + For new code, use extract_contents_messages instead. + """ + messages = extract_contents_messages(contents) + if not messages: + return None + + texts = [] + for message in messages: + content = message.get("content") + if isinstance(content, str): + texts.append(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + texts.append(part.get("text", "")) + + return " ".join(texts) if texts else None + + +def _format_tools_for_span( + tools: "Iterable[Tool | Callable[..., Any]]", +) -> "Optional[List[dict[str, Any]]]": + """Format tools parameter for span data.""" + formatted_tools = [] + for tool in tools: + if callable(tool): + # Handle callable functions passed directly + formatted_tools.append( + { + "name": getattr(tool, "__name__", "unknown"), + "description": getattr(tool, "__doc__", None), + } + ) + elif ( + hasattr(tool, "function_declarations") + and tool.function_declarations is not None + ): + # Tool object with function declarations + for func_decl in tool.function_declarations: + formatted_tools.append( + { + "name": getattr(func_decl, "name", None), + "description": getattr(func_decl, "description", None), + } + ) + else: + # Check for predefined tool attributes - each of these tools + # is an attribute of the tool object, by default set to None + for attr_name, description in TOOL_ATTRIBUTES_MAP.items(): + if getattr(tool, attr_name, None): + formatted_tools.append( + { + "name": attr_name, + "description": description, + } + ) + break + + return formatted_tools if formatted_tools else None + + +def extract_tool_calls( + response: "GenerateContentResponse", +) -> "Optional[List[dict[str, Any]]]": + """Extract tool/function calls from response candidates and automatic function calling history.""" + + tool_calls = [] + + # Extract from candidates, sometimes tool calls are nested under the content.parts object + if getattr(response, "candidates", []): + for candidate in response.candidates: + if not hasattr(candidate, "content") or not getattr( + candidate.content, "parts", [] + ): + continue + + for part in candidate.content.parts: + if getattr(part, "function_call", None): + function_call = part.function_call + tool_call = { + "name": getattr(function_call, "name", None), + "type": "function_call", + } + + # Extract arguments if available + if getattr(function_call, "args", None): + tool_call["arguments"] = safe_serialize(function_call.args) + + tool_calls.append(tool_call) + + # Extract from automatic_function_calling_history + # This is the history of tool calls made by the model + if getattr(response, "automatic_function_calling_history", None): + for content in response.automatic_function_calling_history: + if not getattr(content, "parts", None): + continue + + for part in getattr(content, "parts", []): + if getattr(part, "function_call", None): + function_call = part.function_call + tool_call = { + "name": getattr(function_call, "name", None), + "type": "function_call", + } + + # Extract arguments if available + if hasattr(function_call, "args"): + tool_call["arguments"] = safe_serialize(function_call.args) + + tool_calls.append(tool_call) + + return tool_calls if tool_calls else None + + +def _capture_tool_input( + args: "tuple[Any, ...]", kwargs: "dict[str, Any]", tool: "Tool" +) -> "dict[str, Any]": + """Capture tool input from args and kwargs.""" + tool_input = kwargs.copy() if kwargs else {} + + # If we have positional args, try to map them to the function signature + if args: + try: + sig = inspect.signature(tool) + param_names = list(sig.parameters.keys()) + for i, arg in enumerate(args): + if i < len(param_names): + tool_input[param_names[i]] = arg + except Exception: + # Fallback if we can't get the signature + tool_input["args"] = args + + return tool_input + + +def _create_tool_span(tool_name: str, tool_doc: "Optional[str]") -> "Span": + """Create a span for tool execution.""" + span = sentry_sdk.start_span( + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {tool_name}", + origin=ORIGIN, + ) + span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + span.set_data(SPANDATA.GEN_AI_TOOL_TYPE, "function") + if tool_doc: + span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_doc) + return span + + +def wrapped_tool(tool: "Tool | Callable[..., Any]") -> "Tool | Callable[..., Any]": + """Wrap a tool to emit execute_tool spans when called.""" + if not callable(tool): + # Not a callable function, return as-is (predefined tools) + return tool + + tool_name = getattr(tool, "__name__", "unknown") + tool_doc = tool.__doc__ + + if inspect.iscoroutinefunction(tool): + # Async function + @wraps(tool) + async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": + with _create_tool_span(tool_name, tool_doc) as span: + # Capture tool input + tool_input = _capture_tool_input(args, kwargs, tool) + with capture_internal_exceptions(): + span.set_data( + SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) + ) + + try: + result = await tool(*args, **kwargs) + + # Capture tool output + with capture_internal_exceptions(): + span.set_data( + SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) + ) + + return result + except Exception as exc: + _capture_exception(exc) + raise + + return async_wrapped + else: + # Sync function + @wraps(tool) + def sync_wrapped(*args: "Any", **kwargs: "Any") -> "Any": + with _create_tool_span(tool_name, tool_doc) as span: + # Capture tool input + tool_input = _capture_tool_input(args, kwargs, tool) + with capture_internal_exceptions(): + span.set_data( + SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) + ) + + try: + result = tool(*args, **kwargs) + + # Capture tool output + with capture_internal_exceptions(): + span.set_data( + SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) + ) + + return result + except Exception as exc: + _capture_exception(exc) + raise + + return sync_wrapped + + +def wrapped_config_with_tools( + config: "GenerateContentConfig", +) -> "GenerateContentConfig": + """Wrap tools in config to emit execute_tool spans. Tools are sometimes passed directly as + callable functions as a part of the config object.""" + + if not config or not getattr(config, "tools", None): + return config + + result = copy.copy(config) + result.tools = [wrapped_tool(tool) for tool in config.tools] + + return result + + +def _extract_response_text( + response: "GenerateContentResponse", +) -> "Optional[List[str]]": + """Extract text from response candidates.""" + + if not response or not getattr(response, "candidates", []): + return None + + texts = [] + for candidate in response.candidates: + if not hasattr(candidate, "content") or not hasattr(candidate.content, "parts"): + continue + + for part in candidate.content.parts: + if getattr(part, "text", None): + texts.append(part.text) + + return texts if texts else None + + +def extract_finish_reasons( + response: "GenerateContentResponse", +) -> "Optional[List[str]]": + """Extract finish reasons from response candidates.""" + if not response or not getattr(response, "candidates", []): + return None + + finish_reasons = [] + for candidate in response.candidates: + if getattr(candidate, "finish_reason", None): + # Convert enum value to string if necessary + reason = str(candidate.finish_reason) + # Remove enum prefix if present (e.g., "FinishReason.STOP" -> "STOP") + if "." in reason: + reason = reason.split(".")[-1] + finish_reasons.append(reason) + + return finish_reasons if finish_reasons else None + + +def _transform_system_instruction_one_level( + system_instructions: "Union[ContentUnionDict, ContentUnion]", + can_be_content: bool, +) -> "list[TextPart]": + text_parts: "list[TextPart]" = [] + + if isinstance(system_instructions, str): + return [{"type": "text", "content": system_instructions}] + + if isinstance(system_instructions, Part) and system_instructions.text: + return [{"type": "text", "content": system_instructions.text}] + + if can_be_content and isinstance(system_instructions, Content): + if isinstance(system_instructions.parts, list): + for part in system_instructions.parts: + if isinstance(part.text, str): + text_parts.append({"type": "text", "content": part.text}) + return text_parts + + if isinstance(system_instructions, dict) and system_instructions.get("text"): + return [{"type": "text", "content": system_instructions["text"]}] + + elif can_be_content and isinstance(system_instructions, dict): + parts = system_instructions.get("parts", []) + for part in parts: + if isinstance(part, Part) and isinstance(part.text, str): + text_parts.append({"type": "text", "content": part.text}) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + text_parts.append({"type": "text", "content": part["text"]}) + return text_parts + + return text_parts + + +def _transform_system_instructions( + system_instructions: "Union[ContentUnionDict, ContentUnion]", +) -> "list[TextPart]": + text_parts: "list[TextPart]" = [] + + if isinstance(system_instructions, list): + text_parts = list( + chain.from_iterable( + _transform_system_instruction_one_level( + instructions, can_be_content=False + ) + for instructions in system_instructions + ) + ) + + return text_parts + + return _transform_system_instruction_one_level( + system_instructions, can_be_content=True + ) + + +def set_span_data_for_request( + span: "Span", + integration: "Any", + model: str, + contents: "ContentListUnion", + kwargs: "dict[str, Any]", +) -> None: + """Set span data for the request.""" + span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model) + + if kwargs.get("stream", False): + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + + config: "Optional[GenerateContentConfig]" = kwargs.get("config") + + # Set input messages/prompts if PII is allowed + if should_send_default_pii() and integration.include_prompts: + messages = [] + + # Add system instruction if present + system_instructions = None + if config and hasattr(config, "system_instruction"): + system_instructions = config.system_instruction + elif isinstance(config, dict) and "system_instruction" in config: + system_instructions = config.get("system_instruction") + + if system_instructions is not None: + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(_transform_system_instructions(system_instructions)), + ) + + # Extract messages from contents + contents_messages = extract_contents_messages(contents) + messages.extend(contents_messages) + + if messages: + normalized_messages = normalize_message_roles(messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + # Extract parameters directly from config (not nested under generation_config) + for param, span_key in [ + ("temperature", SPANDATA.GEN_AI_REQUEST_TEMPERATURE), + ("top_p", SPANDATA.GEN_AI_REQUEST_TOP_P), + ("top_k", SPANDATA.GEN_AI_REQUEST_TOP_K), + ("max_output_tokens", SPANDATA.GEN_AI_REQUEST_MAX_TOKENS), + ("presence_penalty", SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY), + ("frequency_penalty", SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY), + ("seed", SPANDATA.GEN_AI_REQUEST_SEED), + ]: + if hasattr(config, param): + value = getattr(config, param) + if value is not None: + span.set_data(span_key, value) + + # Set tools if available + if config is not None and hasattr(config, "tools"): + tools = config.tools + if tools: + formatted_tools = _format_tools_for_span(tools) + if formatted_tools: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + formatted_tools, + unpack=False, + ) + + +def set_span_data_for_response( + span: "Span", integration: "Any", response: "GenerateContentResponse" +) -> None: + """Set span data for the response.""" + if not response: + return + + if should_send_default_pii() and integration.include_prompts: + response_texts = _extract_response_text(response) + if response_texts: + # Format as JSON string array as per documentation + span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts)) + + tool_calls = extract_tool_calls(response) + if tool_calls: + # Tool calls should be JSON serialized + span.set_data(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)) + + finish_reasons = extract_finish_reasons(response) + if finish_reasons: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons + ) + + if getattr(response, "response_id", None): + span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, response.response_id) + + if getattr(response, "model_version", None): + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_version) + + usage_data = extract_usage_data(response) + + if usage_data["input_tokens"]: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage_data["input_tokens"]) + + if usage_data["input_tokens_cached"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, + usage_data["input_tokens_cached"], + ) + + if usage_data["output_tokens"]: + span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage_data["output_tokens"]) + + if usage_data["output_tokens_reasoning"]: + span.set_data( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, + usage_data["output_tokens_reasoning"], + ) + + if usage_data["total_tokens"]: + span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"]) + + +def prepare_generate_content_args( + args: "tuple[Any, ...]", kwargs: "dict[str, Any]" +) -> "tuple[Any, Any, str]": + """Extract and prepare common arguments for generate_content methods.""" + model = args[0] if args else kwargs.get("model", "unknown") + contents = args[1] if len(args) > 1 else kwargs.get("contents") + model_name = get_model_name(model) + + config = kwargs.get("config") + wrapped_config = wrapped_config_with_tools(config) + if wrapped_config is not config: + kwargs["config"] = wrapped_config + + return model, contents, model_name + + +def prepare_embed_content_args( + args: "tuple[Any, ...]", kwargs: "dict[str, Any]" +) -> "tuple[str, Any]": + """Extract and prepare common arguments for embed_content methods. + + Returns: + tuple: (model_name, contents) + """ + model = kwargs.get("model", "unknown") + contents = kwargs.get("contents") + model_name = get_model_name(model) + + return model_name, contents + + +def set_span_data_for_embed_request( + span: "Span", integration: "Any", contents: "Any", kwargs: "dict[str, Any]" +) -> None: + """Set span data for embedding request.""" + # Include input contents if PII is allowed + if should_send_default_pii() and integration.include_prompts: + if contents: + # For embeddings, contents is typically a list of strings/texts + input_texts = [] + + # Handle various content formats + if isinstance(contents, str): + input_texts = [contents] + elif isinstance(contents, list): + for item in contents: + text = extract_contents_text(item) + if text: + input_texts.append(text) + else: + text = extract_contents_text(contents) + if text: + input_texts = [text] + + if input_texts: + set_data_normalized( + span, + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + input_texts, + unpack=False, + ) + + +def set_span_data_for_embed_response( + span: "Span", integration: "Any", response: "EmbedContentResponse" +) -> None: + """Set span data for embedding response.""" + if not response: + return + + # Extract token counts from embeddings statistics (Vertex AI only) + # Each embedding has its own statistics with token_count + if hasattr(response, "embeddings") and response.embeddings: + total_tokens = 0 + + for embedding in response.embeddings: + if hasattr(embedding, "statistics") and embedding.statistics: + token_count = getattr(embedding.statistics, "token_count", None) + if token_count is not None: + total_tokens += int(token_count) + + # Set token count if we found any + if total_tokens > 0: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/gql.py b/lib/python3.12/site-packages/sentry_sdk/integrations/gql.py new file mode 100644 index 0000000000000000000000000000000000000000..c19a518f46cf968ed1c615706219e4a7e982aa80 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/gql.py @@ -0,0 +1,168 @@ +import sentry_sdk +from sentry_sdk.utils import ( + event_from_exception, + ensure_integration_enabled, + parse_version, +) + +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii + +try: + import gql # type: ignore[import-not-found] + from graphql import ( + print_ast, + get_operation_ast, + DocumentNode, + VariableDefinitionNode, + ) + from gql.transport import Transport, AsyncTransport # type: ignore[import-not-found] + from gql.transport.exceptions import TransportQueryError # type: ignore[import-not-found] + + try: + # gql 4.0+ + from gql import GraphQLRequest + except ImportError: + GraphQLRequest = None + +except ImportError: + raise DidNotEnable("gql is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Dict, Tuple, Union + from sentry_sdk._types import Event, EventProcessor + + EventDataType = Dict[str, Union[str, Tuple[VariableDefinitionNode, ...]]] + + +class GQLIntegration(Integration): + identifier = "gql" + + @staticmethod + def setup_once() -> None: + gql_version = parse_version(gql.__version__) + _check_minimum_version(GQLIntegration, gql_version) + + _patch_execute() + + +def _data_from_document(document: "DocumentNode") -> "EventDataType": + try: + operation_ast = get_operation_ast(document) + data: "EventDataType" = {"query": print_ast(document)} + + if operation_ast is not None: + data["variables"] = operation_ast.variable_definitions + if operation_ast.name is not None: + data["operationName"] = operation_ast.name.value + + return data + except (AttributeError, TypeError): + return dict() + + +def _transport_method(transport: "Union[Transport, AsyncTransport]") -> str: + """ + The RequestsHTTPTransport allows defining the HTTP method; all + other transports use POST. + """ + try: + return transport.method + except AttributeError: + return "POST" + + +def _request_info_from_transport( + transport: "Union[Transport, AsyncTransport, None]", +) -> "Dict[str, str]": + if transport is None: + return {} + + request_info = { + "method": _transport_method(transport), + } + + try: + request_info["url"] = transport.url + except AttributeError: + pass + + return request_info + + +def _patch_execute() -> None: + real_execute = gql.Client.execute + + # Maintain signature for backwards compatibility. + # gql.Client.execute() accepts a positional-only "request" + # parameter with version 4.0.0. + @ensure_integration_enabled(GQLIntegration, real_execute) + def sentry_patched_execute( + self: "gql.Client", + document: "DocumentNode", + *args: "Any", + **kwargs: "Any", + ) -> "Any": + scope = sentry_sdk.get_isolation_scope() + # document is a gql.GraphQLRequest with gql v4.0.0. + scope.add_event_processor(_make_gql_event_processor(self, document)) + + try: + # document is a gql.GraphQLRequest with gql v4.0.0. + return real_execute(self, document, *args, **kwargs) + except TransportQueryError as e: + event, hint = event_from_exception( + e, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "gql", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint) + raise e + + gql.Client.execute = sentry_patched_execute + + +def _make_gql_event_processor( + client: "gql.Client", document_or_request: "Union[DocumentNode, gql.GraphQLRequest]" +) -> "EventProcessor": + def processor(event: "Event", hint: "dict[str, Any]") -> "Event": + try: + errors = hint["exc_info"][1].errors + except (AttributeError, KeyError): + errors = None + + request = event.setdefault("request", {}) + request.update( + { + "api_target": "graphql", + **_request_info_from_transport(client.transport), + } + ) + + if should_send_default_pii(): + if GraphQLRequest is not None and isinstance( + document_or_request, GraphQLRequest + ): + # In v4.0.0, gql moved to using GraphQLRequest instead of + # DocumentNode in execute + # https://github.com/graphql-python/gql/pull/556 + document = document_or_request.document + else: + document = document_or_request + + request["data"] = _data_from_document(document) + contexts = event.setdefault("contexts", {}) + response = contexts.setdefault("response", {}) + response.update( + { + "data": {"errors": errors}, + "type": response, + } + ) + + return event + + return processor diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/graphene.py b/lib/python3.12/site-packages/sentry_sdk/integrations/graphene.py new file mode 100644 index 0000000000000000000000000000000000000000..5a61ca5c78b950a547f7bc3b7ecd10e935ad03af --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/graphene.py @@ -0,0 +1,157 @@ +from contextlib import contextmanager + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + package_version, +) + +try: + from graphene.types import schema as graphene_schema # type: ignore +except ImportError: + raise DidNotEnable("graphene is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Generator + from typing import Any, Dict, Union + from graphene.language.source import Source # type: ignore + from graphql.execution import ExecutionResult + from graphql.type import GraphQLSchema + from sentry_sdk._types import Event + + +class GrapheneIntegration(Integration): + identifier = "graphene" + + @staticmethod + def setup_once() -> None: + version = package_version("graphene") + _check_minimum_version(GrapheneIntegration, version) + + _patch_graphql() + + +def _patch_graphql() -> None: + old_graphql_sync = graphene_schema.graphql_sync + old_graphql_async = graphene_schema.graphql + + @ensure_integration_enabled(GrapheneIntegration, old_graphql_sync) + def _sentry_patched_graphql_sync( + schema: "GraphQLSchema", + source: "Union[str, Source]", + *args: "Any", + **kwargs: "Any", + ) -> "ExecutionResult": + scope = sentry_sdk.get_isolation_scope() + scope.add_event_processor(_event_processor) + + with graphql_span(schema, source, kwargs): + result = old_graphql_sync(schema, source, *args, **kwargs) + + with capture_internal_exceptions(): + client = sentry_sdk.get_client() + for error in result.errors or []: + event, hint = event_from_exception( + error, + client_options=client.options, + mechanism={ + "type": GrapheneIntegration.identifier, + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + + return result + + async def _sentry_patched_graphql_async( + schema: "GraphQLSchema", + source: "Union[str, Source]", + *args: "Any", + **kwargs: "Any", + ) -> "ExecutionResult": + integration = sentry_sdk.get_client().get_integration(GrapheneIntegration) + if integration is None: + return await old_graphql_async(schema, source, *args, **kwargs) + + scope = sentry_sdk.get_isolation_scope() + scope.add_event_processor(_event_processor) + + with graphql_span(schema, source, kwargs): + result = await old_graphql_async(schema, source, *args, **kwargs) + + with capture_internal_exceptions(): + client = sentry_sdk.get_client() + for error in result.errors or []: + event, hint = event_from_exception( + error, + client_options=client.options, + mechanism={ + "type": GrapheneIntegration.identifier, + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + + return result + + graphene_schema.graphql_sync = _sentry_patched_graphql_sync + graphene_schema.graphql = _sentry_patched_graphql_async + + +def _event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event": + if should_send_default_pii(): + request_info = event.setdefault("request", {}) + request_info["api_target"] = "graphql" + + elif event.get("request", {}).get("data"): + del event["request"]["data"] + + return event + + +@contextmanager +def graphql_span( + schema: "GraphQLSchema", source: "Union[str, Source]", kwargs: "Dict[str, Any]" +) -> "Generator[None, None, None]": + operation_name = kwargs.get("operation_name") + + operation_type = "query" + op = OP.GRAPHQL_QUERY + if source.strip().startswith("mutation"): + operation_type = "mutation" + op = OP.GRAPHQL_MUTATION + elif source.strip().startswith("subscription"): + operation_type = "subscription" + op = OP.GRAPHQL_SUBSCRIPTION + + sentry_sdk.add_breadcrumb( + crumb={ + "data": { + "operation_name": operation_name, + "operation_type": operation_type, + }, + "category": "graphql.operation", + }, + ) + + scope = sentry_sdk.get_current_scope() + if scope.span: + _graphql_span = scope.span.start_child(op=op, name=operation_name) + else: + _graphql_span = sentry_sdk.start_span(op=op, name=operation_name) + + _graphql_span.set_data("graphql.document", source) + _graphql_span.set_data("graphql.operation.name", operation_name) + _graphql_span.set_data("graphql.operation.type", operation_type) + + try: + yield + finally: + _graphql_span.finish() diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/httpx.py b/lib/python3.12/site-packages/sentry_sdk/integrations/httpx.py new file mode 100644 index 0000000000000000000000000000000000000000..38c4f437bc2f3a63ecd0aee504244fe396b2d3e4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/httpx.py @@ -0,0 +1,156 @@ +import sentry_sdk +from sentry_sdk import start_span +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.tracing import BAGGAGE_HEADER_NAME +from sentry_sdk.tracing_utils import ( + should_propagate_trace, + add_http_request_source, + add_sentry_baggage_to_headers, +) +from sentry_sdk.utils import ( + SENSITIVE_DATA_SUBSTITUTE, + capture_internal_exceptions, + ensure_integration_enabled, + logger, + parse_url, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + + +try: + from httpx import AsyncClient, Client, Request, Response # type: ignore +except ImportError: + raise DidNotEnable("httpx is not installed") + +__all__ = ["HttpxIntegration"] + + +class HttpxIntegration(Integration): + identifier = "httpx" + origin = f"auto.http.{identifier}" + + @staticmethod + def setup_once() -> None: + """ + httpx has its own transport layer and can be customized when needed, + so patch Client.send and AsyncClient.send to support both synchronous and async interfaces. + """ + _install_httpx_client() + _install_httpx_async_client() + + +def _install_httpx_client() -> None: + real_send = Client.send + + @ensure_integration_enabled(HttpxIntegration, real_send) + def send(self: "Client", request: "Request", **kwargs: "Any") -> "Response": + parsed_url = None + with capture_internal_exceptions(): + parsed_url = parse_url(str(request.url), sanitize=False) + + with start_span( + op=OP.HTTP_CLIENT, + name="%s %s" + % ( + request.method, + parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE, + ), + origin=HttpxIntegration.origin, + ) as span: + span.set_data(SPANDATA.HTTP_METHOD, request.method) + if parsed_url is not None: + span.set_data("url", parsed_url.url) + span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query) + span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment) + + if should_propagate_trace(sentry_sdk.get_client(), str(request.url)): + for ( + key, + value, + ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(): + logger.debug( + "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format( + key=key, value=value, url=request.url + ) + ) + + if key == BAGGAGE_HEADER_NAME: + add_sentry_baggage_to_headers(request.headers, value) + else: + request.headers[key] = value + + rv = real_send(self, request, **kwargs) + + span.set_http_status(rv.status_code) + span.set_data("reason", rv.reason_phrase) + + with capture_internal_exceptions(): + add_http_request_source(span) + + return rv + + Client.send = send + + +def _install_httpx_async_client() -> None: + real_send = AsyncClient.send + + async def send( + self: "AsyncClient", request: "Request", **kwargs: "Any" + ) -> "Response": + if sentry_sdk.get_client().get_integration(HttpxIntegration) is None: + return await real_send(self, request, **kwargs) + + parsed_url = None + with capture_internal_exceptions(): + parsed_url = parse_url(str(request.url), sanitize=False) + + with start_span( + op=OP.HTTP_CLIENT, + name="%s %s" + % ( + request.method, + parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE, + ), + origin=HttpxIntegration.origin, + ) as span: + span.set_data(SPANDATA.HTTP_METHOD, request.method) + if parsed_url is not None: + span.set_data("url", parsed_url.url) + span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query) + span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment) + + if should_propagate_trace(sentry_sdk.get_client(), str(request.url)): + for ( + key, + value, + ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(): + logger.debug( + "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format( + key=key, value=value, url=request.url + ) + ) + if key == BAGGAGE_HEADER_NAME and request.headers.get( + BAGGAGE_HEADER_NAME + ): + # do not overwrite any existing baggage, just append to it + request.headers[key] += "," + value + else: + request.headers[key] = value + + rv = await real_send(self, request, **kwargs) + + span.set_http_status(rv.status_code) + span.set_data("reason", rv.reason_phrase) + + with capture_internal_exceptions(): + add_http_request_source(span) + + return rv + + AsyncClient.send = send diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/huey.py b/lib/python3.12/site-packages/sentry_sdk/integrations/huey.py new file mode 100644 index 0000000000000000000000000000000000000000..1c7626f3fa35d34d67120b60cc68a52304b09823 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/huey.py @@ -0,0 +1,166 @@ +import sys +from datetime import datetime + +import sentry_sdk +from sentry_sdk.api import continue_trace, get_baggage, get_traceparent +from sentry_sdk.consts import OP, SPANSTATUS +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SENTRY_TRACE_HEADER_NAME, + TransactionSource, +) +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + SENSITIVE_DATA_SUBSTITUTE, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Optional, Union, TypeVar + + from sentry_sdk._types import EventProcessor, Event, Hint + from sentry_sdk.utils import ExcInfo + + F = TypeVar("F", bound=Callable[..., Any]) + +try: + from huey.api import Huey, Result, ResultGroup, Task, PeriodicTask + from huey.exceptions import CancelExecution, RetryTask, TaskLockedException +except ImportError: + raise DidNotEnable("Huey is not installed") + + +HUEY_CONTROL_FLOW_EXCEPTIONS = (CancelExecution, RetryTask, TaskLockedException) + + +class HueyIntegration(Integration): + identifier = "huey" + origin = f"auto.queue.{identifier}" + + @staticmethod + def setup_once() -> None: + patch_enqueue() + patch_execute() + + +def patch_enqueue() -> None: + old_enqueue = Huey.enqueue + + @ensure_integration_enabled(HueyIntegration, old_enqueue) + def _sentry_enqueue( + self: "Huey", task: "Task" + ) -> "Optional[Union[Result, ResultGroup]]": + with sentry_sdk.start_span( + op=OP.QUEUE_SUBMIT_HUEY, + name=task.name, + origin=HueyIntegration.origin, + ): + if not isinstance(task, PeriodicTask): + # Attach trace propagation data to task kwargs. We do + # not do this for periodic tasks, as these don't + # really have an originating transaction. + task.kwargs["sentry_headers"] = { + BAGGAGE_HEADER_NAME: get_baggage(), + SENTRY_TRACE_HEADER_NAME: get_traceparent(), + } + return old_enqueue(self, task) + + Huey.enqueue = _sentry_enqueue + + +def _make_event_processor(task: "Any") -> "EventProcessor": + def event_processor(event: "Event", hint: "Hint") -> "Optional[Event]": + with capture_internal_exceptions(): + tags = event.setdefault("tags", {}) + tags["huey_task_id"] = task.id + tags["huey_task_retry"] = task.default_retries > task.retries + extra = event.setdefault("extra", {}) + extra["huey-job"] = { + "task": task.name, + "args": ( + task.args + if should_send_default_pii() + else SENSITIVE_DATA_SUBSTITUTE + ), + "kwargs": ( + task.kwargs + if should_send_default_pii() + else SENSITIVE_DATA_SUBSTITUTE + ), + "retry": (task.default_retries or 0) - task.retries, + } + + return event + + return event_processor + + +def _capture_exception(exc_info: "ExcInfo") -> None: + scope = sentry_sdk.get_current_scope() + + if exc_info[0] in HUEY_CONTROL_FLOW_EXCEPTIONS: + scope.transaction.set_status(SPANSTATUS.ABORTED) + return + + scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR) + event, hint = event_from_exception( + exc_info, + client_options=sentry_sdk.get_client().options, + mechanism={"type": HueyIntegration.identifier, "handled": False}, + ) + scope.capture_event(event, hint=hint) + + +def _wrap_task_execute(func: "F") -> "F": + @ensure_integration_enabled(HueyIntegration, func) + def _sentry_execute(*args: "Any", **kwargs: "Any") -> "Any": + try: + result = func(*args, **kwargs) + except Exception: + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result + + return _sentry_execute # type: ignore + + +def patch_execute() -> None: + old_execute = Huey._execute + + @ensure_integration_enabled(HueyIntegration, old_execute) + def _sentry_execute( + self: "Huey", task: "Task", timestamp: "Optional[datetime]" = None + ) -> "Any": + with sentry_sdk.isolation_scope() as scope: + with capture_internal_exceptions(): + scope._name = "huey" + scope.clear_breadcrumbs() + scope.add_event_processor(_make_event_processor(task)) + + sentry_headers = task.kwargs.pop("sentry_headers", None) + + transaction = continue_trace( + sentry_headers or {}, + name=task.name, + op=OP.QUEUE_TASK_HUEY, + source=TransactionSource.TASK, + origin=HueyIntegration.origin, + ) + transaction.set_status(SPANSTATUS.OK) + + if not getattr(task, "_sentry_is_patched", False): + task.execute = _wrap_task_execute(task.execute) + task._sentry_is_patched = True + + with sentry_sdk.start_transaction(transaction): + return old_execute(self, task, timestamp) + + Huey._execute = _sentry_execute diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/huggingface_hub.py b/lib/python3.12/site-packages/sentry_sdk/integrations/huggingface_hub.py new file mode 100644 index 0000000000000000000000000000000000000000..8509cadefae4b450cb46b650bcdb104ebe8ed7d4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/huggingface_hub.py @@ -0,0 +1,376 @@ +import sys +import inspect +from functools import wraps + +import sentry_sdk +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import set_span_errored +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Iterable + +try: + import huggingface_hub.inference._client +except ImportError: + raise DidNotEnable("Huggingface not installed") + + +class HuggingfaceHubIntegration(Integration): + identifier = "huggingface_hub" + origin = f"auto.ai.{identifier}" + + def __init__( + self: "HuggingfaceHubIntegration", include_prompts: bool = True + ) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + # Other tasks that can be called: https://huggingface.co/docs/huggingface_hub/guides/inference#supported-providers-and-tasks + huggingface_hub.inference._client.InferenceClient.text_generation = ( + _wrap_huggingface_task( + huggingface_hub.inference._client.InferenceClient.text_generation, + OP.GEN_AI_GENERATE_TEXT, + ) + ) + huggingface_hub.inference._client.InferenceClient.chat_completion = ( + _wrap_huggingface_task( + huggingface_hub.inference._client.InferenceClient.chat_completion, + OP.GEN_AI_CHAT, + ) + ) + + +def _capture_exception(exc: "Any") -> None: + set_span_errored() + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "huggingface_hub", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _wrap_huggingface_task(f: "Callable[..., Any]", op: str) -> "Callable[..., Any]": + @wraps(f) + def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration) + if integration is None: + return f(*args, **kwargs) + + prompt = None + if "prompt" in kwargs: + prompt = kwargs["prompt"] + elif "messages" in kwargs: + prompt = kwargs["messages"] + elif len(args) >= 2: + if isinstance(args[1], str) or isinstance(args[1], list): + prompt = args[1] + + if prompt is None: + # invalid call, dont instrument, let it return error + return f(*args, **kwargs) + + client = args[0] + model = client.model or kwargs.get("model") or "" + operation_name = op.split(".")[-1] + + span = sentry_sdk.start_span( + op=op, + name=f"{operation_name} {model}", + origin=HuggingfaceHubIntegration.origin, + ) + span.__enter__() + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, operation_name) + + if model: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model) + + # Input attributes + if should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False + ) + + attribute_mapping = { + "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K, + "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING, + } + + for attribute, span_attribute in attribute_mapping.items(): + value = kwargs.get(attribute, None) + if value is not None: + if isinstance(value, (int, float, bool, str)): + span.set_data(span_attribute, value) + else: + set_data_normalized(span, span_attribute, value, unpack=False) + + # LLM Execution + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + span.__exit__(None, None, None) + reraise(*exc_info) + + # Output attributes + finish_reason = None + response_model = None + response_text_buffer: "list[str]" = [] + tokens_used = 0 + tool_calls = None + usage = None + + with capture_internal_exceptions(): + if isinstance(res, str) and res is not None: + response_text_buffer.append(res) + + if hasattr(res, "generated_text") and res.generated_text is not None: + response_text_buffer.append(res.generated_text) + + if hasattr(res, "model") and res.model is not None: + response_model = res.model + + if hasattr(res, "details") and hasattr(res.details, "finish_reason"): + finish_reason = res.details.finish_reason + + if ( + hasattr(res, "details") + and hasattr(res.details, "generated_tokens") + and res.details.generated_tokens is not None + ): + tokens_used = res.details.generated_tokens + + if hasattr(res, "usage") and res.usage is not None: + usage = res.usage + + if hasattr(res, "choices") and res.choices is not None: + for choice in res.choices: + if hasattr(choice, "finish_reason"): + finish_reason = choice.finish_reason + if hasattr(choice, "message") and hasattr( + choice.message, "tool_calls" + ): + tool_calls = choice.message.tool_calls + if ( + hasattr(choice, "message") + and hasattr(choice.message, "content") + and choice.message.content is not None + ): + response_text_buffer.append(choice.message.content) + + if response_model is not None: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) + + if finish_reason is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, + finish_reason, + ) + + if should_send_default_pii() and integration.include_prompts: + if tool_calls is not None and len(tool_calls) > 0: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) + + if len(response_text_buffer) > 0: + text_response = "".join(response_text_buffer) + if text_response: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + text_response, + ) + + if usage is not None: + record_token_usage( + span, + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens, + ) + elif tokens_used > 0: + record_token_usage( + span, + total_tokens=tokens_used, + ) + + # If the response is not a generator (meaning a streaming response) + # we are done and can return the response + if not inspect.isgenerator(res): + span.__exit__(None, None, None) + return res + + if kwargs.get("details", False): + # text-generation stream output + def new_details_iterator() -> "Iterable[Any]": + finish_reason = None + response_text_buffer: "list[str]" = [] + tokens_used = 0 + + with capture_internal_exceptions(): + for chunk in res: + if ( + hasattr(chunk, "token") + and hasattr(chunk.token, "text") + and chunk.token.text is not None + ): + response_text_buffer.append(chunk.token.text) + + if hasattr(chunk, "details") and hasattr( + chunk.details, "finish_reason" + ): + finish_reason = chunk.details.finish_reason + + if ( + hasattr(chunk, "details") + and hasattr(chunk.details, "generated_tokens") + and chunk.details.generated_tokens is not None + ): + tokens_used = chunk.details.generated_tokens + + yield chunk + + if finish_reason is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, + finish_reason, + ) + + if should_send_default_pii() and integration.include_prompts: + if len(response_text_buffer) > 0: + text_response = "".join(response_text_buffer) + if text_response: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + text_response, + ) + + if tokens_used > 0: + record_token_usage( + span, + total_tokens=tokens_used, + ) + + span.__exit__(None, None, None) + + return new_details_iterator() + + else: + # chat-completion stream output + def new_iterator() -> "Iterable[str]": + finish_reason = None + response_model = None + response_text_buffer: "list[str]" = [] + tool_calls = None + usage = None + + with capture_internal_exceptions(): + for chunk in res: + if hasattr(chunk, "model") and chunk.model is not None: + response_model = chunk.model + + if hasattr(chunk, "usage") and chunk.usage is not None: + usage = chunk.usage + + if isinstance(chunk, str): + if chunk is not None: + response_text_buffer.append(chunk) + + if hasattr(chunk, "choices") and chunk.choices is not None: + for choice in chunk.choices: + if ( + hasattr(choice, "delta") + and hasattr(choice.delta, "content") + and choice.delta.content is not None + ): + response_text_buffer.append( + choice.delta.content + ) + + if ( + hasattr(choice, "finish_reason") + and choice.finish_reason is not None + ): + finish_reason = choice.finish_reason + + if ( + hasattr(choice, "delta") + and hasattr(choice.delta, "tool_calls") + and choice.delta.tool_calls is not None + ): + tool_calls = choice.delta.tool_calls + + yield chunk + + if response_model is not None: + span.set_data( + SPANDATA.GEN_AI_RESPONSE_MODEL, response_model + ) + + if finish_reason is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, + finish_reason, + ) + + if should_send_default_pii() and integration.include_prompts: + if tool_calls is not None and len(tool_calls) > 0: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) + + if len(response_text_buffer) > 0: + text_response = "".join(response_text_buffer) + if text_response: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + text_response, + ) + + if usage is not None: + record_token_usage( + span, + input_tokens=usage.prompt_tokens, + output_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens, + ) + + span.__exit__(None, None, None) + + return new_iterator() + + return new_huggingface_task diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/langchain.py b/lib/python3.12/site-packages/sentry_sdk/integrations/langchain.py new file mode 100644 index 0000000000000000000000000000000000000000..7a5e863e4d6b466455a927ff3cd427d4a5c13c0d --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/langchain.py @@ -0,0 +1,1244 @@ +import contextvars +import itertools +import sys +import json +import warnings +from collections import OrderedDict +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.ai.monitoring import set_ai_pipeline_name +from sentry_sdk.ai.utils import ( + GEN_AI_ALLOWED_MESSAGE_ROLES, + get_start_span_function, + normalize_message_roles, + set_data_normalized, + truncate_and_annotate_messages, + transform_content_part, +) +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import _get_value, set_span_errored +from sentry_sdk.utils import capture_internal_exceptions, logger + +if TYPE_CHECKING: + from typing import ( + Any, + AsyncIterator, + Callable, + Dict, + Iterator, + List, + Optional, + Union, + ) + from uuid import UUID + + from sentry_sdk.tracing import Span + from sentry_sdk._types import TextPart + + +try: + from langchain_core.agents import AgentFinish + from langchain_core.callbacks import ( + BaseCallbackHandler, + BaseCallbackManager, + Callbacks, + manager, + ) + from langchain_core.messages import BaseMessage + from langchain_core.outputs import LLMResult + +except ImportError: + raise DidNotEnable("langchain not installed") + + +try: + # >=v1 + from langchain_classic.agents import AgentExecutor # type: ignore[import-not-found] +except ImportError: + try: + # "Dict[str, Any]": + """ + Transform a LangChain content block using the shared transform_content_part function. + + Returns the original content block if transformation is not applicable + (e.g., for text blocks or unrecognized formats). + """ + result = transform_content_part(content_block) + return result if result is not None else content_block + + +def _transform_langchain_message_content(content: "Any") -> "Any": + """ + Transform LangChain message content, handling both string content and + list of content blocks. + """ + if isinstance(content, str): + return content + + if isinstance(content, (list, tuple)): + transformed = [] + for block in content: + if isinstance(block, dict): + transformed.append(_transform_langchain_content_block(block)) + else: + transformed.append(block) + return transformed + + return content + + +# Contextvar to track agent names in a stack for re-entrant agent support +_agent_stack: "contextvars.ContextVar[Optional[List[Optional[str]]]]" = ( + contextvars.ContextVar("langchain_agent_stack", default=None) +) + + +def _push_agent(agent_name: "Optional[str]") -> None: + """Push an agent name onto the stack.""" + stack = _agent_stack.get() + if stack is None: + stack = [] + else: + # Copy the list to maintain contextvar isolation across async contexts + stack = stack.copy() + stack.append(agent_name) + _agent_stack.set(stack) + + +def _pop_agent() -> "Optional[str]": + """Pop an agent name from the stack and return it.""" + stack = _agent_stack.get() + if stack: + # Copy the list to maintain contextvar isolation across async contexts + stack = stack.copy() + agent_name = stack.pop() + _agent_stack.set(stack) + return agent_name + return None + + +def _get_current_agent() -> "Optional[str]": + """Get the current agent name (top of stack) without removing it.""" + stack = _agent_stack.get() + if stack: + return stack[-1] + return None + + +def _get_system_instructions(messages: "List[List[BaseMessage]]") -> "List[str]": + system_instructions = [] + + for list_ in messages: + for message in list_: + # type of content: str | list[str | dict] | None + if message.type == "system" and isinstance(message.content, str): + system_instructions.append(message.content) + + elif message.type == "system" and isinstance(message.content, list): + for item in message.content: + if isinstance(item, str): + system_instructions.append(item) + + elif isinstance(item, dict) and item.get("type") == "text": + instruction = item.get("text") + if isinstance(instruction, str): + system_instructions.append(instruction) + + return system_instructions + + +def _transform_system_instructions( + system_instructions: "List[str]", +) -> "List[TextPart]": + return [ + { + "type": "text", + "content": instruction, + } + for instruction in system_instructions + ] + + +class LangchainIntegration(Integration): + identifier = "langchain" + origin = f"auto.ai.{identifier}" + + def __init__( + self: "LangchainIntegration", + include_prompts: bool = True, + max_spans: "Optional[int]" = None, + ) -> None: + self.include_prompts = include_prompts + self.max_spans = max_spans + + if max_spans is not None: + warnings.warn( + "The `max_spans` parameter of `LangchainIntegration` is " + "deprecated and will be removed in version 3.0 of sentry-sdk.", + DeprecationWarning, + stacklevel=2, + ) + + @staticmethod + def setup_once() -> None: + manager._configure = _wrap_configure(manager._configure) + + if AgentExecutor is not None: + AgentExecutor.invoke = _wrap_agent_executor_invoke(AgentExecutor.invoke) + AgentExecutor.stream = _wrap_agent_executor_stream(AgentExecutor.stream) + + # Patch embeddings providers + _patch_embeddings_provider(OpenAIEmbeddings) + _patch_embeddings_provider(AzureOpenAIEmbeddings) + _patch_embeddings_provider(VertexAIEmbeddings) + _patch_embeddings_provider(BedrockEmbeddings) + _patch_embeddings_provider(CohereEmbeddings) + _patch_embeddings_provider(MistralAIEmbeddings) + _patch_embeddings_provider(HuggingFaceEmbeddings) + _patch_embeddings_provider(OllamaEmbeddings) + + +class WatchedSpan: + span: "Span" = None # type: ignore[assignment] + children: "List[WatchedSpan]" = [] + is_pipeline: bool = False + + def __init__(self, span: "Span") -> None: + self.span = span + + +class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc] + """Callback handler that creates Sentry spans.""" + + def __init__( + self, max_span_map_size: "Optional[int]", include_prompts: bool + ) -> None: + self.span_map: "OrderedDict[UUID, WatchedSpan]" = OrderedDict() + self.max_span_map_size = max_span_map_size + self.include_prompts = include_prompts + + def gc_span_map(self) -> None: + if self.max_span_map_size is not None: + while len(self.span_map) > self.max_span_map_size: + run_id, watched_span = self.span_map.popitem(last=False) + self._exit_span(watched_span, run_id) + + def _handle_error(self, run_id: "UUID", error: "Any") -> None: + with capture_internal_exceptions(): + if not run_id or run_id not in self.span_map: + return + + span_data = self.span_map[run_id] + span = span_data.span + set_span_errored(span) + + sentry_sdk.capture_exception(error, span.scope) + + span.__exit__(None, None, None) + del self.span_map[run_id] + + def _normalize_langchain_message(self, message: "BaseMessage") -> "Any": + # Transform content to handle multimodal data (images, audio, video, files) + transformed_content = _transform_langchain_message_content(message.content) + parsed = {"role": message.type, "content": transformed_content} + parsed.update(message.additional_kwargs) + return parsed + + def _create_span( + self: "SentryLangchainCallback", + run_id: "UUID", + parent_id: "Optional[Any]", + **kwargs: "Any", + ) -> "WatchedSpan": + watched_span: "Optional[WatchedSpan]" = None + if parent_id: + parent_span: "Optional[WatchedSpan]" = self.span_map.get(parent_id) + if parent_span: + watched_span = WatchedSpan(parent_span.span.start_child(**kwargs)) + parent_span.children.append(watched_span) + + if watched_span is None: + watched_span = WatchedSpan(sentry_sdk.start_span(**kwargs)) + + watched_span.span.__enter__() + self.span_map[run_id] = watched_span + self.gc_span_map() + return watched_span + + def _exit_span( + self: "SentryLangchainCallback", span_data: "WatchedSpan", run_id: "UUID" + ) -> None: + if span_data.is_pipeline: + set_ai_pipeline_name(None) + + span_data.span.__exit__(None, None, None) + del self.span_map[run_id] + + def on_llm_start( + self: "SentryLangchainCallback", + serialized: "Dict[str, Any]", + prompts: "List[str]", + *, + run_id: "UUID", + tags: "Optional[List[str]]" = None, + parent_run_id: "Optional[UUID]" = None, + metadata: "Optional[Dict[str, Any]]" = None, + **kwargs: "Any", + ) -> "Any": + """Run when LLM starts running.""" + with capture_internal_exceptions(): + if not run_id: + return + + all_params = kwargs.get("invocation_params", {}) + all_params.update(serialized.get("kwargs", {})) + + model = ( + all_params.get("model") + or all_params.get("model_name") + or all_params.get("model_id") + or "" + ) + + watched_span = self._create_span( + run_id, + parent_run_id, + op=OP.GEN_AI_PIPELINE, + name=kwargs.get("name") or "Langchain LLM call", + origin=LangchainIntegration.origin, + ) + span = watched_span.span + + if model: + span.set_data( + SPANDATA.GEN_AI_REQUEST_MODEL, + model, + ) + + ai_type = all_params.get("_type", "") + if "anthropic" in ai_type: + span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic") + elif "openai" in ai_type: + span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai") + + for key, attribute in DATA_FIELDS.items(): + if key in all_params and all_params[key] is not None: + set_data_normalized(span, attribute, all_params[key], unpack=False) + + _set_tools_on_span(span, all_params.get("tools")) + + if should_send_default_pii() and self.include_prompts: + normalized_messages = [ + { + "role": GEN_AI_ALLOWED_MESSAGE_ROLES.USER, + "content": {"type": "text", "text": prompt}, + } + for prompt in prompts + ] + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + def on_chat_model_start( + self: "SentryLangchainCallback", + serialized: "Dict[str, Any]", + messages: "List[List[BaseMessage]]", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when Chat Model starts running.""" + with capture_internal_exceptions(): + if not run_id: + return + + all_params = kwargs.get("invocation_params", {}) + all_params.update(serialized.get("kwargs", {})) + + model = ( + all_params.get("model") + or all_params.get("model_name") + or all_params.get("model_id") + or "" + ) + + watched_span = self._create_span( + run_id, + kwargs.get("parent_run_id"), + op=OP.GEN_AI_CHAT, + name=f"chat {model}".strip(), + origin=LangchainIntegration.origin, + ) + span = watched_span.span + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + if model: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model) + + ai_type = all_params.get("_type", "") + if "anthropic" in ai_type: + span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic") + elif "openai" in ai_type: + span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai") + + agent_name = _get_current_agent() + if agent_name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name) + + for key, attribute in DATA_FIELDS.items(): + if key in all_params and all_params[key] is not None: + set_data_normalized(span, attribute, all_params[key], unpack=False) + + _set_tools_on_span(span, all_params.get("tools")) + + if should_send_default_pii() and self.include_prompts: + system_instructions = _get_system_instructions(messages) + if len(system_instructions) > 0: + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(_transform_system_instructions(system_instructions)), + ) + + normalized_messages = [] + for list_ in messages: + for message in list_: + if message.type == "system": + continue + + normalized_messages.append( + self._normalize_langchain_message(message) + ) + normalized_messages = normalize_message_roles(normalized_messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + def on_chat_model_end( + self: "SentryLangchainCallback", + response: "LLMResult", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when Chat Model ends running.""" + with capture_internal_exceptions(): + if not run_id or run_id not in self.span_map: + return + + span_data = self.span_map[run_id] + span = span_data.span + + if should_send_default_pii() and self.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + [[x.text for x in list_] for list_ in response.generations], + ) + + _record_token_usage(span, response) + self._exit_span(span_data, run_id) + + def on_llm_end( + self: "SentryLangchainCallback", + response: "LLMResult", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when LLM ends running.""" + with capture_internal_exceptions(): + if not run_id or run_id not in self.span_map: + return + + span_data = self.span_map[run_id] + span = span_data.span + + try: + generation = response.generations[0][0] + except IndexError: + generation = None + + if generation is not None: + try: + response_model = generation.message.response_metadata.get( + "model_name" + ) + if response_model is not None: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) + except AttributeError: + pass + + try: + finish_reason = generation.generation_info.get("finish_reason") + if finish_reason is not None: + span.set_data( + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reason + ) + except AttributeError: + pass + + try: + if should_send_default_pii() and self.include_prompts: + tool_calls = getattr(generation.message, "tool_calls", None) + if tool_calls is not None and tool_calls != []: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) + except AttributeError: + pass + + if should_send_default_pii() and self.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + [[x.text for x in list_] for list_ in response.generations], + ) + + _record_token_usage(span, response) + self._exit_span(span_data, run_id) + + def on_llm_error( + self: "SentryLangchainCallback", + error: "Union[Exception, KeyboardInterrupt]", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when LLM errors.""" + self._handle_error(run_id, error) + + def on_chat_model_error( + self: "SentryLangchainCallback", + error: "Union[Exception, KeyboardInterrupt]", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when Chat Model errors.""" + self._handle_error(run_id, error) + + def on_agent_finish( + self: "SentryLangchainCallback", + finish: "AgentFinish", + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + with capture_internal_exceptions(): + if not run_id or run_id not in self.span_map: + return + + span_data = self.span_map[run_id] + span = span_data.span + + if should_send_default_pii() and self.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, finish.return_values.items() + ) + + self._exit_span(span_data, run_id) + + def on_tool_start( + self: "SentryLangchainCallback", + serialized: "Dict[str, Any]", + input_str: str, + *, + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when tool starts running.""" + with capture_internal_exceptions(): + if not run_id: + return + + tool_name = serialized.get("name") or kwargs.get("name") or "" + + watched_span = self._create_span( + run_id, + kwargs.get("parent_run_id"), + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {tool_name}".strip(), + origin=LangchainIntegration.origin, + ) + span = watched_span.span + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") + span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + + tool_description = serialized.get("description") + if tool_description is not None: + span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description) + + agent_name = _get_current_agent() + if agent_name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name) + + if should_send_default_pii() and self.include_prompts: + set_data_normalized( + span, + SPANDATA.GEN_AI_TOOL_INPUT, + kwargs.get("inputs", [input_str]), + ) + + def on_tool_end( + self: "SentryLangchainCallback", output: str, *, run_id: "UUID", **kwargs: "Any" + ) -> "Any": + """Run when tool ends running.""" + with capture_internal_exceptions(): + if not run_id or run_id not in self.span_map: + return + + span_data = self.span_map[run_id] + span = span_data.span + + if should_send_default_pii() and self.include_prompts: + set_data_normalized(span, SPANDATA.GEN_AI_TOOL_OUTPUT, output) + + self._exit_span(span_data, run_id) + + def on_tool_error( + self, + error: "SentryLangchainCallback", + *args: "Union[Exception, KeyboardInterrupt]", + run_id: "UUID", + **kwargs: "Any", + ) -> "Any": + """Run when tool errors.""" + self._handle_error(run_id, error) + + +def _extract_tokens( + token_usage: "Any", +) -> "tuple[Optional[int], Optional[int], Optional[int]]": + if not token_usage: + return None, None, None + + input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value( + token_usage, "input_tokens" + ) + output_tokens = _get_value(token_usage, "completion_tokens") or _get_value( + token_usage, "output_tokens" + ) + total_tokens = _get_value(token_usage, "total_tokens") + + return input_tokens, output_tokens, total_tokens + + +def _extract_tokens_from_generations( + generations: "Any", +) -> "tuple[Optional[int], Optional[int], Optional[int]]": + """Extract token usage from response.generations structure.""" + if not generations: + return None, None, None + + total_input = 0 + total_output = 0 + total_total = 0 + + for gen_list in generations: + for gen in gen_list: + token_usage = _get_token_usage(gen) + input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage) + total_input += input_tokens if input_tokens is not None else 0 + total_output += output_tokens if output_tokens is not None else 0 + total_total += total_tokens if total_tokens is not None else 0 + + return ( + total_input if total_input > 0 else None, + total_output if total_output > 0 else None, + total_total if total_total > 0 else None, + ) + + +def _get_token_usage(obj: "Any") -> "Optional[Dict[str, Any]]": + """ + Check multiple paths to extract token usage from different objects. + """ + possible_names = ("usage", "token_usage", "usage_metadata") + + message = _get_value(obj, "message") + if message is not None: + for name in possible_names: + usage = _get_value(message, name) + if usage is not None: + return usage + + llm_output = _get_value(obj, "llm_output") + if llm_output is not None: + for name in possible_names: + usage = _get_value(llm_output, name) + if usage is not None: + return usage + + for name in possible_names: + usage = _get_value(obj, name) + if usage is not None: + return usage + + return None + + +def _record_token_usage(span: "Span", response: "Any") -> None: + token_usage = _get_token_usage(response) + if token_usage: + input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage) + else: + input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations( + response.generations + ) + + if input_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + + if output_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + + if total_tokens is not None: + span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) + + +def _get_request_data( + obj: "Any", args: "Any", kwargs: "Any" +) -> "tuple[Optional[str], Optional[List[Any]]]": + """ + Get the agent name and available tools for the agent. + """ + agent = getattr(obj, "agent", None) + runnable = getattr(agent, "runnable", None) + runnable_config = getattr(runnable, "config", {}) + tools = ( + getattr(obj, "tools", None) + or getattr(agent, "tools", None) + or runnable_config.get("tools") + or runnable_config.get("available_tools") + ) + tools = tools if tools and len(tools) > 0 else None + + try: + agent_name = None + if len(args) > 1: + agent_name = args[1].get("run_name") + if agent_name is None: + agent_name = runnable_config.get("run_name") + except Exception: + pass + + return (agent_name, tools) + + +def _simplify_langchain_tools(tools: "Any") -> "Optional[List[Any]]": + """Parse and simplify tools into a cleaner format.""" + if not tools: + return None + + if not isinstance(tools, (list, tuple)): + return None + + simplified_tools = [] + for tool in tools: + try: + if isinstance(tool, dict): + if "function" in tool and isinstance(tool["function"], dict): + func = tool["function"] + simplified_tool = { + "name": func.get("name"), + "description": func.get("description"), + } + if simplified_tool["name"]: + simplified_tools.append(simplified_tool) + elif "name" in tool: + simplified_tool = { + "name": tool.get("name"), + "description": tool.get("description"), + } + simplified_tools.append(simplified_tool) + else: + name = ( + tool.get("name") + or tool.get("tool_name") + or tool.get("function_name") + ) + if name: + simplified_tools.append( + { + "name": name, + "description": tool.get("description") + or tool.get("desc"), + } + ) + elif hasattr(tool, "name"): + simplified_tool = { + "name": getattr(tool, "name", None), + "description": getattr(tool, "description", None) + or getattr(tool, "desc", None), + } + if simplified_tool["name"]: + simplified_tools.append(simplified_tool) + elif hasattr(tool, "__name__"): + simplified_tools.append( + { + "name": tool.__name__, + "description": getattr(tool, "__doc__", None), + } + ) + else: + tool_str = str(tool) + if tool_str and tool_str != "": + simplified_tools.append({"name": tool_str, "description": None}) + except Exception: + continue + + return simplified_tools if simplified_tools else None + + +def _set_tools_on_span(span: "Span", tools: "Any") -> None: + """Set available tools data on a span if tools are provided.""" + if tools is not None: + simplified_tools = _simplify_langchain_tools(tools) + if simplified_tools: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + simplified_tools, + unpack=False, + ) + + +def _wrap_configure(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_configure( + callback_manager_cls: type, + inheritable_callbacks: "Callbacks" = None, + local_callbacks: "Callbacks" = None, + *args: "Any", + **kwargs: "Any", + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(LangchainIntegration) + if integration is None: + return f( + callback_manager_cls, + inheritable_callbacks, + local_callbacks, + *args, + **kwargs, + ) + + local_callbacks = local_callbacks or [] + + # Handle each possible type of local_callbacks. For each type, we + # extract the list of callbacks to check for SentryLangchainCallback, + # and define a function that would add the SentryLangchainCallback + # to the existing callbacks list. + if isinstance(local_callbacks, BaseCallbackManager): + callbacks_list = local_callbacks.handlers + elif isinstance(local_callbacks, BaseCallbackHandler): + callbacks_list = [local_callbacks] + elif isinstance(local_callbacks, list): + callbacks_list = local_callbacks + else: + logger.debug("Unknown callback type: %s", local_callbacks) + # Just proceed with original function call + return f( + callback_manager_cls, + inheritable_callbacks, + local_callbacks, + *args, + **kwargs, + ) + + # Handle each possible type of inheritable_callbacks. + if isinstance(inheritable_callbacks, BaseCallbackManager): + inheritable_callbacks_list = inheritable_callbacks.handlers + elif isinstance(inheritable_callbacks, list): + inheritable_callbacks_list = inheritable_callbacks + else: + inheritable_callbacks_list = [] + + if not any( + isinstance(cb, SentryLangchainCallback) + for cb in itertools.chain(callbacks_list, inheritable_callbacks_list) + ): + sentry_handler = SentryLangchainCallback( + integration.max_spans, + integration.include_prompts, + ) + if isinstance(local_callbacks, BaseCallbackManager): + local_callbacks = local_callbacks.copy() + local_callbacks.handlers = [ + *local_callbacks.handlers, + sentry_handler, + ] + elif isinstance(local_callbacks, BaseCallbackHandler): + local_callbacks = [local_callbacks, sentry_handler] + else: + local_callbacks = [*local_callbacks, sentry_handler] + + return f( + callback_manager_cls, + inheritable_callbacks, + local_callbacks, + *args, + **kwargs, + ) + + return new_configure + + +def _wrap_agent_executor_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LangchainIntegration) + if integration is None: + return f(self, *args, **kwargs) + + agent_name, tools = _get_request_data(self, args, kwargs) + start_span_function = get_start_span_function() + + with start_span_function( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent", + origin=LangchainIntegration.origin, + ) as span: + _push_agent(agent_name) + try: + if agent_name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) + + _set_tools_on_span(span, tools) + + # Run the agent + result = f(self, *args, **kwargs) + + input = result.get("input") + if ( + input is not None + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles([input]) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + output = result.get("output") + if ( + output is not None + and should_send_default_pii() + and integration.include_prompts + ): + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output) + + return result + finally: + # Ensure agent is popped even if an exception occurs + _pop_agent() + + return new_invoke + + +def _wrap_agent_executor_stream(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LangchainIntegration) + if integration is None: + return f(self, *args, **kwargs) + + agent_name, tools = _get_request_data(self, args, kwargs) + start_span_function = get_start_span_function() + + span = start_span_function( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent", + origin=LangchainIntegration.origin, + ) + span.__enter__() + + _push_agent(agent_name) + + if agent_name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + + _set_tools_on_span(span, tools) + + input = args[0].get("input") if len(args) >= 1 else None + if ( + input is not None + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles([input]) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + # Run the agent + result = f(self, *args, **kwargs) + + old_iterator = result + + def new_iterator() -> "Iterator[Any]": + exc_info: "tuple[Any, Any, Any]" = (None, None, None) + try: + for event in old_iterator: + yield event + + try: + output = event.get("output") + except Exception: + output = None + + if ( + output is not None + and should_send_default_pii() + and integration.include_prompts + ): + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output) + except Exception: + exc_info = sys.exc_info() + set_span_errored(span) + raise + finally: + # Ensure cleanup happens even if iterator is abandoned or fails + _pop_agent() + span.__exit__(*exc_info) + + async def new_iterator_async() -> "AsyncIterator[Any]": + exc_info: "tuple[Any, Any, Any]" = (None, None, None) + try: + async for event in old_iterator: + yield event + + try: + output = event.get("output") + except Exception: + output = None + + if ( + output is not None + and should_send_default_pii() + and integration.include_prompts + ): + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output) + except Exception: + exc_info = sys.exc_info() + set_span_errored(span) + raise + finally: + # Ensure cleanup happens even if iterator is abandoned or fails + _pop_agent() + span.__exit__(*exc_info) + + if str(type(result)) == "": + result = new_iterator_async() + else: + result = new_iterator() + + return result + + return new_stream + + +def _patch_embeddings_provider(provider_class: "Any") -> None: + """Patch an embeddings provider class with monitoring wrappers.""" + if provider_class is None: + return + + if hasattr(provider_class, "embed_documents"): + provider_class.embed_documents = _wrap_embedding_method( + provider_class.embed_documents + ) + if hasattr(provider_class, "embed_query"): + provider_class.embed_query = _wrap_embedding_method(provider_class.embed_query) + if hasattr(provider_class, "aembed_documents"): + provider_class.aembed_documents = _wrap_async_embedding_method( + provider_class.aembed_documents + ) + if hasattr(provider_class, "aembed_query"): + provider_class.aembed_query = _wrap_async_embedding_method( + provider_class.aembed_query + ) + + +def _wrap_embedding_method(f: "Callable[..., Any]") -> "Callable[..., Any]": + """Wrap sync embedding methods (embed_documents and embed_query).""" + + @wraps(f) + def new_embedding_method(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LangchainIntegration) + if integration is None: + return f(self, *args, **kwargs) + + model_name = getattr(self, "model", None) or getattr(self, "model_name", None) + with sentry_sdk.start_span( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}" if model_name else "embeddings", + origin=LangchainIntegration.origin, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + if model_name: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + + # Capture input if PII is allowed + if ( + should_send_default_pii() + and integration.include_prompts + and len(args) > 0 + ): + input_data = args[0] + # Normalize to list format + texts = input_data if isinstance(input_data, list) else [input_data] + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, texts, unpack=False + ) + + result = f(self, *args, **kwargs) + return result + + return new_embedding_method + + +def _wrap_async_embedding_method(f: "Callable[..., Any]") -> "Callable[..., Any]": + """Wrap async embedding methods (aembed_documents and aembed_query).""" + + @wraps(f) + async def new_async_embedding_method( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + integration = sentry_sdk.get_client().get_integration(LangchainIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + model_name = getattr(self, "model", None) or getattr(self, "model_name", None) + with sentry_sdk.start_span( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}" if model_name else "embeddings", + origin=LangchainIntegration.origin, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + if model_name: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + + # Capture input if PII is allowed + if ( + should_send_default_pii() + and integration.include_prompts + and len(args) > 0 + ): + input_data = args[0] + # Normalize to list format + texts = input_data if isinstance(input_data, list) else [input_data] + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, texts, unpack=False + ) + + result = await f(self, *args, **kwargs) + return result + + return new_async_embedding_method diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/langgraph.py b/lib/python3.12/site-packages/sentry_sdk/integrations/langgraph.py new file mode 100644 index 0000000000000000000000000000000000000000..e5ea12b90ab244198b9ab74cc1119b711d73d6ef --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/langgraph.py @@ -0,0 +1,386 @@ +from functools import wraps +from typing import Any, Callable, List, Optional + +import sentry_sdk +from sentry_sdk.ai.utils import ( + set_data_normalized, + normalize_message_roles, + truncate_and_annotate_messages, +) +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import safe_serialize + + +try: + from langgraph.graph import StateGraph + from langgraph.pregel import Pregel +except ImportError: + raise DidNotEnable("langgraph not installed") + + +class LanggraphIntegration(Integration): + identifier = "langgraph" + origin = f"auto.ai.{identifier}" + + def __init__(self: "LanggraphIntegration", include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + # LangGraph lets users create agents using a StateGraph or the Functional API. + # StateGraphs are then compiled to a CompiledStateGraph. Both CompiledStateGraph and + # the functional API execute on a Pregel instance. Pregel is the runtime for the graph + # and the invocation happens on Pregel, so patching the invoke methods takes care of both. + # The streaming methods are not patched, because due to some internal reasons, LangGraph + # will automatically patch the streaming methods to run through invoke, and by doing this + # we prevent duplicate spans for invocations. + StateGraph.compile = _wrap_state_graph_compile(StateGraph.compile) + if hasattr(Pregel, "invoke"): + Pregel.invoke = _wrap_pregel_invoke(Pregel.invoke) + if hasattr(Pregel, "ainvoke"): + Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke) + + +def _get_graph_name(graph_obj: "Any") -> "Optional[str]": + for attr in ["name", "graph_name", "__name__", "_name"]: + if hasattr(graph_obj, attr): + name = getattr(graph_obj, attr) + if name and isinstance(name, str): + return name + return None + + +def _normalize_langgraph_message(message: "Any") -> "Any": + if not hasattr(message, "content"): + return None + + parsed = {"role": getattr(message, "type", None), "content": message.content} + + for attr in [ + "name", + "tool_calls", + "function_call", + "tool_call_id", + "response_metadata", + ]: + if hasattr(message, attr): + value = getattr(message, attr) + if value is not None: + parsed[attr] = value + + return parsed + + +def _parse_langgraph_messages(state: "Any") -> "Optional[List[Any]]": + if not state: + return None + + messages = None + + if isinstance(state, dict): + messages = state.get("messages") + elif hasattr(state, "messages"): + messages = state.messages + elif hasattr(state, "get") and callable(state.get): + try: + messages = state.get("messages") + except Exception: + pass + + if not messages or not isinstance(messages, (list, tuple)): + return None + + normalized_messages = [] + for message in messages: + try: + normalized = _normalize_langgraph_message(message) + if normalized: + normalized_messages.append(normalized) + except Exception: + continue + + return normalized_messages if normalized_messages else None + + +def _wrap_state_graph_compile(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_compile(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return f(self, *args, **kwargs) + with sentry_sdk.start_span( + op=OP.GEN_AI_CREATE_AGENT, + origin=LanggraphIntegration.origin, + ) as span: + compiled_graph = f(self, *args, **kwargs) + + compiled_graph_name = getattr(compiled_graph, "name", None) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "create_agent") + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, compiled_graph_name) + + if compiled_graph_name: + span.description = f"create_agent {compiled_graph_name}" + else: + span.description = "create_agent" + + if kwargs.get("model", None) is not None: + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, kwargs.get("model")) + + tools = None + get_graph = getattr(compiled_graph, "get_graph", None) + if get_graph and callable(get_graph): + graph_obj = compiled_graph.get_graph() + nodes = getattr(graph_obj, "nodes", None) + if nodes and isinstance(nodes, dict): + tools_node = nodes.get("tools") + if tools_node: + data = getattr(tools_node, "data", None) + if data and hasattr(data, "tools_by_name"): + tools = list(data.tools_by_name.keys()) + + if tools is not None: + span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools) + + return compiled_graph + + return new_compile + + +def _wrap_pregel_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return f(self, *args, **kwargs) + + graph_name = _get_graph_name(self) + span_name = ( + f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" + ) + + with sentry_sdk.start_span( + op=OP.GEN_AI_INVOKE_AGENT, + name=span_name, + origin=LanggraphIntegration.origin, + ) as span: + if graph_name: + span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + + # Store input messages to later compare with output + input_messages = None + if ( + len(args) > 0 + and should_send_default_pii() + and integration.include_prompts + ): + input_messages = _parse_langgraph_messages(args[0]) + if input_messages: + normalized_input_messages = normalize_message_roles(input_messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_input_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + result = f(self, *args, **kwargs) + + _set_response_attributes(span, input_messages, result, integration) + + return result + + return new_invoke + + +def _wrap_pregel_ainvoke(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return await f(self, *args, **kwargs) + + graph_name = _get_graph_name(self) + span_name = ( + f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" + ) + + with sentry_sdk.start_span( + op=OP.GEN_AI_INVOKE_AGENT, + name=span_name, + origin=LanggraphIntegration.origin, + ) as span: + if graph_name: + span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) + + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + + input_messages = None + if ( + len(args) > 0 + and should_send_default_pii() + and integration.include_prompts + ): + input_messages = _parse_langgraph_messages(args[0]) + if input_messages: + normalized_input_messages = normalize_message_roles(input_messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages( + normalized_input_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + result = await f(self, *args, **kwargs) + + _set_response_attributes(span, input_messages, result, integration) + + return result + + return new_ainvoke + + +def _get_new_messages( + input_messages: "Optional[List[Any]]", output_messages: "Optional[List[Any]]" +) -> "Optional[List[Any]]": + """Extract only the new messages added during this invocation.""" + if not output_messages: + return None + + if not input_messages: + return output_messages + + # only return the new messages, aka the output messages that are not in the input messages + input_count = len(input_messages) + new_messages = ( + output_messages[input_count:] if len(output_messages) > input_count else [] + ) + + return new_messages if new_messages else None + + +def _extract_llm_response_text(messages: "Optional[List[Any]]") -> "Optional[str]": + if not messages: + return None + + for message in reversed(messages): + if isinstance(message, dict): + role = message.get("role") + if role in ["assistant", "ai"]: + content = message.get("content") + if content and isinstance(content, str): + return content + + return None + + +def _extract_tool_calls(messages: "Optional[List[Any]]") -> "Optional[List[Any]]": + if not messages: + return None + + tool_calls = [] + for message in messages: + if isinstance(message, dict): + msg_tool_calls = message.get("tool_calls") + if msg_tool_calls and isinstance(msg_tool_calls, list): + tool_calls.extend(msg_tool_calls) + + return tool_calls if tool_calls else None + + +def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: + input_tokens = 0 + output_tokens = 0 + total_tokens = 0 + + for message in messages: + response_metadata = message.get("response_metadata") + if response_metadata is None: + continue + + token_usage = response_metadata.get("token_usage") + if not token_usage: + continue + + input_tokens += int(token_usage.get("prompt_tokens", 0)) + output_tokens += int(token_usage.get("completion_tokens", 0)) + total_tokens += int(token_usage.get("total_tokens", 0)) + + if input_tokens > 0: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + + if output_tokens > 0: + span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + + if total_tokens > 0: + span.set_data( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + total_tokens, + ) + + +def _set_response_model_name(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: + if len(messages) == 0: + return + + last_message = messages[-1] + response_metadata = last_message.get("response_metadata") + if response_metadata is None: + return + + model_name = response_metadata.get("model_name") + if model_name is None: + return + + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model_name) + + +def _set_response_attributes( + span: "Any", + input_messages: "Optional[List[Any]]", + result: "Any", + integration: "LanggraphIntegration", +) -> None: + parsed_response_messages = _parse_langgraph_messages(result) + new_messages = _get_new_messages(input_messages, parsed_response_messages) + + if new_messages is None: + return + + _set_usage_data(span, new_messages) + _set_response_model_name(span, new_messages) + + if not (should_send_default_pii() and integration.include_prompts): + return + + llm_response_text = _extract_llm_response_text(new_messages) + if llm_response_text: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text) + elif new_messages: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages) + else: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result) + + tool_calls = _extract_tool_calls(new_messages) + if tool_calls: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + safe_serialize(tool_calls), + unpack=False, + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/launchdarkly.py b/lib/python3.12/site-packages/sentry_sdk/integrations/launchdarkly.py new file mode 100644 index 0000000000000000000000000000000000000000..2d86fc5ca49af6ea162beaa8d1b80736d45225b4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/launchdarkly.py @@ -0,0 +1,63 @@ +from typing import TYPE_CHECKING + +from sentry_sdk.feature_flags import add_feature_flag +from sentry_sdk.integrations import DidNotEnable, Integration + +try: + import ldclient + from ldclient.hook import Hook, Metadata + + if TYPE_CHECKING: + from ldclient import LDClient + from ldclient.hook import EvaluationSeriesContext + from ldclient.evaluation import EvaluationDetail + + from typing import Any +except ImportError: + raise DidNotEnable("LaunchDarkly is not installed") + + +class LaunchDarklyIntegration(Integration): + identifier = "launchdarkly" + + def __init__(self, ld_client: "LDClient | None" = None) -> None: + """ + :param client: An initialized LDClient instance. If a client is not provided, this + integration will attempt to use the shared global instance. + """ + try: + client = ld_client or ldclient.get() + except Exception as exc: + raise DidNotEnable("Error getting LaunchDarkly client. " + repr(exc)) + + if not client.is_initialized(): + raise DidNotEnable("LaunchDarkly client is not initialized.") + + # Register the flag collection hook with the LD client. + client.add_hook(LaunchDarklyHook()) + + @staticmethod + def setup_once() -> None: + pass + + +class LaunchDarklyHook(Hook): + @property + def metadata(self) -> "Metadata": + return Metadata(name="sentry-flag-auditor") + + def after_evaluation( + self, + series_context: "EvaluationSeriesContext", + data: "dict[Any, Any]", + detail: "EvaluationDetail", + ) -> "dict[Any, Any]": + if isinstance(detail.value, bool): + add_feature_flag(series_context.key, detail.value) + + return data + + def before_evaluation( + self, series_context: "EvaluationSeriesContext", data: "dict[Any, Any]" + ) -> "dict[Any, Any]": + return data # No-op. diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/litellm.py b/lib/python3.12/site-packages/sentry_sdk/integrations/litellm.py new file mode 100644 index 0000000000000000000000000000000000000000..28bcc34d3e270b0ffb02469ed38daf55b5445f4f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/litellm.py @@ -0,0 +1,325 @@ +import copy +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk import consts +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import ( + get_start_span_function, + set_data_normalized, + truncate_and_annotate_messages, + transform_openai_content_part, + truncate_and_annotate_embedding_inputs, +) +from sentry_sdk.consts import SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.utils import event_from_exception + +if TYPE_CHECKING: + from typing import Any, Dict, List + from datetime import datetime + +try: + import litellm # type: ignore[import-not-found] + from litellm import input_callback, success_callback, failure_callback +except ImportError: + raise DidNotEnable("LiteLLM not installed") + + +def _get_metadata_dict(kwargs: "Dict[str, Any]") -> "Dict[str, Any]": + """Get the metadata dictionary from the kwargs.""" + litellm_params = kwargs.setdefault("litellm_params", {}) + + # we need this weird little dance, as metadata might be set but may be None initially + metadata = litellm_params.get("metadata") + if metadata is None: + metadata = {} + litellm_params["metadata"] = metadata + return metadata + + +def _convert_message_parts(messages: "List[Dict[str, Any]]") -> "List[Dict[str, Any]]": + """ + Convert the message parts from OpenAI format to the `gen_ai.request.messages` format + using the OpenAI-specific transformer (LiteLLM uses OpenAI's message format). + + Deep copies messages to avoid mutating original kwargs. + """ + # Deep copy to avoid mutating original messages from kwargs + messages = copy.deepcopy(messages) + + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, (list, tuple)): + transformed = [] + for item in content: + if isinstance(item, dict): + result = transform_openai_content_part(item) + # If transformation succeeded, use the result; otherwise keep original + transformed.append(result if result is not None else item) + else: + transformed.append(item) + message["content"] = transformed + return messages + + +def _input_callback(kwargs: "Dict[str, Any]") -> None: + """Handle the start of a request.""" + integration = sentry_sdk.get_client().get_integration(LiteLLMIntegration) + + if integration is None: + return + + # Get key parameters + full_model = kwargs.get("model", "") + try: + model, provider, _, _ = litellm.get_llm_provider(full_model) + except Exception: + model = full_model + provider = "unknown" + + call_type = kwargs.get("call_type", None) + if call_type == "embedding": + operation = "embeddings" + else: + operation = "chat" + + # Start a new span/transaction + span = get_start_span_function()( + op=( + consts.OP.GEN_AI_CHAT + if operation == "chat" + else consts.OP.GEN_AI_EMBEDDINGS + ), + name=f"{operation} {model}", + origin=LiteLLMIntegration.origin, + ) + span.__enter__() + + # Store span for later + _get_metadata_dict(kwargs)["_sentry_span"] = span + + # Set basic data + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, provider) + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) + + # Record input/messages if allowed + if should_send_default_pii() and integration.include_prompts: + if operation == "embeddings": + # For embeddings, look for the 'input' parameter + embedding_input = kwargs.get("input") + if embedding_input: + scope = sentry_sdk.get_current_scope() + # Normalize to list format + input_list = ( + embedding_input + if isinstance(embedding_input, list) + else [embedding_input] + ) + messages_data = truncate_and_annotate_embedding_inputs( + input_list, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + messages_data, + unpack=False, + ) + else: + # For chat, look for the 'messages' parameter + messages = kwargs.get("messages", []) + if messages: + scope = sentry_sdk.get_current_scope() + messages = _convert_message_parts(messages) + messages_data = truncate_and_annotate_messages(messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + messages_data, + unpack=False, + ) + + # Record other parameters + params = { + "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING, + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + } + for key, attribute in params.items(): + value = kwargs.get(key) + if value is not None: + set_data_normalized(span, attribute, value) + + # Record LiteLLM-specific parameters + litellm_params = { + "api_base": kwargs.get("api_base"), + "api_version": kwargs.get("api_version"), + "custom_llm_provider": kwargs.get("custom_llm_provider"), + } + for key, value in litellm_params.items(): + if value is not None: + set_data_normalized(span, f"gen_ai.litellm.{key}", value) + + +def _success_callback( + kwargs: "Dict[str, Any]", + completion_response: "Any", + start_time: "datetime", + end_time: "datetime", +) -> None: + """Handle successful completion.""" + + span = _get_metadata_dict(kwargs).get("_sentry_span") + if span is None: + return + + integration = sentry_sdk.get_client().get_integration(LiteLLMIntegration) + if integration is None: + return + + try: + # Record model information + if hasattr(completion_response, "model"): + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, completion_response.model + ) + + # Record response content if allowed + if should_send_default_pii() and integration.include_prompts: + if hasattr(completion_response, "choices"): + response_messages = [] + for choice in completion_response.choices: + if hasattr(choice, "message"): + if hasattr(choice.message, "model_dump"): + response_messages.append(choice.message.model_dump()) + elif hasattr(choice.message, "dict"): + response_messages.append(choice.message.dict()) + else: + # Fallback for basic message objects + msg = {} + if hasattr(choice.message, "role"): + msg["role"] = choice.message.role + if hasattr(choice.message, "content"): + msg["content"] = choice.message.content + if hasattr(choice.message, "tool_calls"): + msg["tool_calls"] = choice.message.tool_calls + response_messages.append(msg) + + if response_messages: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_messages + ) + + # Record token usage + if hasattr(completion_response, "usage"): + usage = completion_response.usage + record_token_usage( + span, + input_tokens=getattr(usage, "prompt_tokens", None), + output_tokens=getattr(usage, "completion_tokens", None), + total_tokens=getattr(usage, "total_tokens", None), + ) + + finally: + # Always finish the span and clean up + span.__exit__(None, None, None) + + +def _failure_callback( + kwargs: "Dict[str, Any]", + exception: Exception, + start_time: "datetime", + end_time: "datetime", +) -> None: + """Handle request failure.""" + span = _get_metadata_dict(kwargs).get("_sentry_span") + if span is None: + return + + try: + # Capture the exception + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "litellm", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + finally: + # Always finish the span and clean up + span.__exit__(type(exception), exception, None) + + +class LiteLLMIntegration(Integration): + """ + LiteLLM integration for Sentry. + + This integration automatically captures LiteLLM API calls and sends them to Sentry + for monitoring and error tracking. It supports all 100+ LLM providers that LiteLLM + supports, including OpenAI, Anthropic, Google, Cohere, and many others. + + Features: + - Automatic exception capture for all LiteLLM calls + - Token usage tracking across all providers + - Provider detection and attribution + - Input/output message capture (configurable) + - Streaming response support + - Cost tracking integration + + Usage: + + ```python + import litellm + import sentry_sdk + + # Initialize Sentry with the LiteLLM integration + sentry_sdk.init( + dsn="your-dsn", + send_default_pii=True + integrations=[ + sentry_sdk.integrations.LiteLLMIntegration( + include_prompts=True # Set to False to exclude message content + ) + ] + ) + + # All LiteLLM calls will now be monitored + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}] + ) + ``` + + Configuration: + - include_prompts (bool): Whether to include prompts and responses in spans. + Defaults to True. Set to False to exclude potentially sensitive data. + """ + + identifier = "litellm" + origin = f"auto.ai.{identifier}" + + def __init__(self: "LiteLLMIntegration", include_prompts: bool = True) -> None: + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + """Set up LiteLLM callbacks for monitoring.""" + litellm.input_callback = input_callback or [] + if _input_callback not in litellm.input_callback: + litellm.input_callback.append(_input_callback) + + litellm.success_callback = success_callback or [] + if _success_callback not in litellm.success_callback: + litellm.success_callback.append(_success_callback) + + litellm.failure_callback = failure_callback or [] + if _failure_callback not in litellm.failure_callback: + litellm.failure_callback.append(_failure_callback) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/litestar.py b/lib/python3.12/site-packages/sentry_sdk/integrations/litestar.py new file mode 100644 index 0000000000000000000000000000000000000000..e0baf7f591c47389db8943aad2f287b396eb9de7 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/litestar.py @@ -0,0 +1,311 @@ +from collections.abc import Set +from copy import deepcopy + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import ( + _DEFAULT_FAILED_REQUEST_STATUS_CODES, + DidNotEnable, + Integration, +) +from sentry_sdk.integrations.asgi import SentryAsgiMiddleware +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import TransactionSource, SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + ensure_integration_enabled, + event_from_exception, + transaction_from_function, +) + +try: + from litestar import Request, Litestar # type: ignore + from litestar.handlers.base import BaseRouteHandler # type: ignore + from litestar.middleware import DefineMiddleware # type: ignore + from litestar.routes.http import HTTPRoute # type: ignore + from litestar.data_extractors import ConnectionDataExtractor # type: ignore + from litestar.exceptions import HTTPException # type: ignore +except ImportError: + raise DidNotEnable("Litestar is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Optional, Union + from litestar.types.asgi_types import ASGIApp # type: ignore + from litestar.types import ( # type: ignore + HTTPReceiveMessage, + HTTPScope, + Message, + Middleware, + Receive, + Scope as LitestarScope, + Send, + WebSocketReceiveMessage, + ) + from litestar.middleware import MiddlewareProtocol + from sentry_sdk._types import Event, Hint + +_DEFAULT_TRANSACTION_NAME = "generic Litestar request" + + +class LitestarIntegration(Integration): + identifier = "litestar" + origin = f"auto.http.{identifier}" + + def __init__( + self, + failed_request_status_codes: "Set[int]" = _DEFAULT_FAILED_REQUEST_STATUS_CODES, + ) -> None: + self.failed_request_status_codes = failed_request_status_codes + + @staticmethod + def setup_once() -> None: + patch_app_init() + patch_middlewares() + patch_http_route_handle() + + # The following line follows the pattern found in other integrations such as `DjangoIntegration.setup_once`. + # The Litestar `ExceptionHandlerMiddleware.__call__` catches exceptions and does the following + # (among other things): + # 1. Logs them, some at least (such as 500s) as errors + # 2. Calls after_exception hooks + # The `LitestarIntegration`` provides an after_exception hook (see `patch_app_init` below) to create a Sentry event + # from an exception, which ends up being called during step 2 above. However, the Sentry `LoggingIntegration` will + # by default create a Sentry event from error logs made in step 1 if we do not prevent it from doing so. + ignore_logger("litestar") + + +class SentryLitestarASGIMiddleware(SentryAsgiMiddleware): + def __init__( + self, app: "ASGIApp", span_origin: str = LitestarIntegration.origin + ) -> None: + super().__init__( + app=app, + unsafe_context_data=False, + transaction_style="endpoint", + mechanism_type="asgi", + span_origin=span_origin, + asgi_version=3, + ) + + def _capture_request_exception(self, exc: Exception) -> None: + """Avoid catching exceptions from request handlers. + + Those exceptions are already handled in Litestar.after_exception handler. + We still catch exceptions from application lifespan handlers. + """ + pass + + +def patch_app_init() -> None: + """ + Replaces the Litestar class's `__init__` function in order to inject `after_exception` handlers and set the + `SentryLitestarASGIMiddleware` as the outmost middleware in the stack. + See: + - https://docs.litestar.dev/2/usage/applications.html#after-exception + - https://docs.litestar.dev/2/usage/middleware/using-middleware.html + """ + old__init__ = Litestar.__init__ + + @ensure_integration_enabled(LitestarIntegration, old__init__) + def injection_wrapper(self: "Litestar", *args: "Any", **kwargs: "Any") -> None: + kwargs["after_exception"] = [ + exception_handler, + *(kwargs.get("after_exception") or []), + ] + + middleware = kwargs.get("middleware") or [] + kwargs["middleware"] = [SentryLitestarASGIMiddleware, *middleware] + old__init__(self, *args, **kwargs) + + Litestar.__init__ = injection_wrapper + + +def patch_middlewares() -> None: + old_resolve_middleware_stack = BaseRouteHandler.resolve_middleware + + @ensure_integration_enabled(LitestarIntegration, old_resolve_middleware_stack) + def resolve_middleware_wrapper(self: "BaseRouteHandler") -> "list[Middleware]": + return [ + enable_span_for_middleware(middleware) + for middleware in old_resolve_middleware_stack(self) + ] + + BaseRouteHandler.resolve_middleware = resolve_middleware_wrapper + + +def enable_span_for_middleware(middleware: "Middleware") -> "Middleware": + if ( + not hasattr(middleware, "__call__") # noqa: B004 + or middleware is SentryLitestarASGIMiddleware + ): + return middleware + + if isinstance(middleware, DefineMiddleware): + old_call: "ASGIApp" = middleware.middleware.__call__ + else: + old_call = middleware.__call__ + + async def _create_span_call( + self: "MiddlewareProtocol", + scope: "LitestarScope", + receive: "Receive", + send: "Send", + ) -> None: + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: + return await old_call(self, scope, receive, send) + + middleware_name = self.__class__.__name__ + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_LITESTAR, + name=middleware_name, + origin=LitestarIntegration.origin, + ) as middleware_span: + middleware_span.set_tag("litestar.middleware_name", middleware_name) + + # Creating spans for the "receive" callback + async def _sentry_receive( + *args: "Any", **kwargs: "Any" + ) -> "Union[HTTPReceiveMessage, WebSocketReceiveMessage]": + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: + return await receive(*args, **kwargs) + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_LITESTAR_RECEIVE, + name=getattr(receive, "__qualname__", str(receive)), + origin=LitestarIntegration.origin, + ) as span: + span.set_tag("litestar.middleware_name", middleware_name) + return await receive(*args, **kwargs) + + receive_name = getattr(receive, "__name__", str(receive)) + receive_patched = receive_name == "_sentry_receive" + new_receive = _sentry_receive if not receive_patched else receive + + # Creating spans for the "send" callback + async def _sentry_send(message: "Message") -> None: + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: + return await send(message) + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_LITESTAR_SEND, + name=getattr(send, "__qualname__", str(send)), + origin=LitestarIntegration.origin, + ) as span: + span.set_tag("litestar.middleware_name", middleware_name) + return await send(message) + + send_name = getattr(send, "__name__", str(send)) + send_patched = send_name == "_sentry_send" + new_send = _sentry_send if not send_patched else send + + return await old_call(self, scope, new_receive, new_send) + + not_yet_patched = old_call.__name__ not in ["_create_span_call"] + + if not_yet_patched: + if isinstance(middleware, DefineMiddleware): + middleware.middleware.__call__ = _create_span_call + else: + middleware.__call__ = _create_span_call + + return middleware + + +def patch_http_route_handle() -> None: + old_handle = HTTPRoute.handle + + async def handle_wrapper( + self: "HTTPRoute", scope: "HTTPScope", receive: "Receive", send: "Send" + ) -> None: + if sentry_sdk.get_client().get_integration(LitestarIntegration) is None: + return await old_handle(self, scope, receive, send) + + sentry_scope = sentry_sdk.get_isolation_scope() + request: "Request[Any, Any]" = scope["app"].request_class( + scope=scope, receive=receive, send=send + ) + extracted_request_data = ConnectionDataExtractor( + parse_body=True, parse_query=True + )(request) + body = extracted_request_data.pop("body") + + request_data = await body + + def event_processor(event: "Event", _: "Hint") -> "Event": + route_handler = scope.get("route_handler") + + request_info = event.get("request", {}) + request_info["content_length"] = len(scope.get("_body", b"")) + if should_send_default_pii(): + request_info["cookies"] = extracted_request_data["cookies"] + if request_data is not None: + request_info["data"] = request_data + + func = None + if route_handler.name is not None: + tx_name = route_handler.name + # Accounts for use of type `Ref` in earlier versions of litestar without the need to reference it as a type + elif hasattr(route_handler.fn, "value"): + func = route_handler.fn.value + else: + func = route_handler.fn + if func is not None: + tx_name = transaction_from_function(func) + + tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]} + + if not tx_name: + tx_name = _DEFAULT_TRANSACTION_NAME + tx_info = {"source": TransactionSource.ROUTE} + + event.update( + { + "request": deepcopy(request_info), + "transaction": tx_name, + "transaction_info": tx_info, + } + ) + return event + + sentry_scope._name = LitestarIntegration.identifier + sentry_scope.add_event_processor(event_processor) + + return await old_handle(self, scope, receive, send) + + HTTPRoute.handle = handle_wrapper + + +def retrieve_user_from_scope(scope: "LitestarScope") -> "Optional[dict[str, Any]]": + scope_user = scope.get("user") + if isinstance(scope_user, dict): + return scope_user + if hasattr(scope_user, "asdict"): # dataclasses + return scope_user.asdict() + + return None + + +@ensure_integration_enabled(LitestarIntegration) +def exception_handler(exc: Exception, scope: "LitestarScope") -> None: + user_info: "Optional[dict[str, Any]]" = None + if should_send_default_pii(): + user_info = retrieve_user_from_scope(scope) + if user_info and isinstance(user_info, dict): + sentry_scope = sentry_sdk.get_isolation_scope() + sentry_scope.set_user(user_info) + + if isinstance(exc, HTTPException): + integration = sentry_sdk.get_client().get_integration(LitestarIntegration) + if ( + integration is not None + and exc.status_code not in integration.failed_request_status_codes + ): + return + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": LitestarIntegration.identifier, "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/logging.py b/lib/python3.12/site-packages/sentry_sdk/integrations/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..42029c5a7a5d50e0a9307a982bd2a4af906133e8 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/logging.py @@ -0,0 +1,409 @@ +import logging +import sys +from datetime import datetime, timezone +from fnmatch import fnmatch + +import sentry_sdk +from sentry_sdk.client import BaseClient +from sentry_sdk.logger import _log_level_to_otel +from sentry_sdk.utils import ( + safe_repr, + to_string, + event_from_exception, + current_stacktrace, + capture_internal_exceptions, + has_logs_enabled, +) +from sentry_sdk.integrations import Integration + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import MutableMapping + from logging import LogRecord + from typing import Any + from typing import Dict + from typing import Optional + +DEFAULT_LEVEL = logging.INFO +DEFAULT_EVENT_LEVEL = logging.ERROR +LOGGING_TO_EVENT_LEVEL = { + logging.NOTSET: "notset", + logging.DEBUG: "debug", + logging.INFO: "info", + logging.WARN: "warning", # WARN is same a WARNING + logging.WARNING: "warning", + logging.ERROR: "error", + logging.FATAL: "fatal", + logging.CRITICAL: "fatal", # CRITICAL is same as FATAL +} + +# Map logging level numbers to corresponding OTel level numbers +SEVERITY_TO_OTEL_SEVERITY = { + logging.CRITICAL: 21, # fatal + logging.ERROR: 17, # error + logging.WARNING: 13, # warn + logging.INFO: 9, # info + logging.DEBUG: 5, # debug +} + + +# Capturing events from those loggers causes recursion errors. We cannot allow +# the user to unconditionally create events from those loggers under any +# circumstances. +# +# Note: Ignoring by logger name here is better than mucking with thread-locals. +# We do not necessarily know whether thread-locals work 100% correctly in the user's environment. +_IGNORED_LOGGERS = set( + ["sentry_sdk.errors", "urllib3.connectionpool", "urllib3.connection"] +) + + +def ignore_logger( + name: str, +) -> None: + """This disables recording (both in breadcrumbs and as events) calls to + a logger of a specific name. Among other uses, many of our integrations + use this to prevent their actions being recorded as breadcrumbs. Exposed + to users as a way to quiet spammy loggers. + + :param name: The name of the logger to ignore (same string you would pass to ``logging.getLogger``). + """ + _IGNORED_LOGGERS.add(name) + + +class LoggingIntegration(Integration): + identifier = "logging" + + def __init__( + self, + level: "Optional[int]" = DEFAULT_LEVEL, + event_level: "Optional[int]" = DEFAULT_EVENT_LEVEL, + sentry_logs_level: "Optional[int]" = DEFAULT_LEVEL, + ) -> None: + self._handler = None + self._breadcrumb_handler = None + self._sentry_logs_handler = None + + if level is not None: + self._breadcrumb_handler = BreadcrumbHandler(level=level) + + if sentry_logs_level is not None: + self._sentry_logs_handler = SentryLogsHandler(level=sentry_logs_level) + + if event_level is not None: + self._handler = EventHandler(level=event_level) + + def _handle_record(self, record: "LogRecord") -> None: + if self._handler is not None and record.levelno >= self._handler.level: + self._handler.handle(record) + + if ( + self._breadcrumb_handler is not None + and record.levelno >= self._breadcrumb_handler.level + ): + self._breadcrumb_handler.handle(record) + + if ( + self._sentry_logs_handler is not None + and record.levelno >= self._sentry_logs_handler.level + ): + self._sentry_logs_handler.handle(record) + + @staticmethod + def setup_once() -> None: + old_callhandlers = logging.Logger.callHandlers + + def sentry_patched_callhandlers(self: "Any", record: "LogRecord") -> "Any": + # keeping a local reference because the + # global might be discarded on shutdown + ignored_loggers = _IGNORED_LOGGERS + + try: + return old_callhandlers(self, record) + finally: + # This check is done twice, once also here before we even get + # the integration. Otherwise we have a high chance of getting + # into a recursion error when the integration is resolved + # (this also is slower). + if ( + ignored_loggers is not None + and record.name.strip() not in ignored_loggers + ): + integration = sentry_sdk.get_client().get_integration( + LoggingIntegration + ) + if integration is not None: + integration._handle_record(record) + + logging.Logger.callHandlers = sentry_patched_callhandlers # type: ignore + + +class _BaseHandler(logging.Handler): + COMMON_RECORD_ATTRS = frozenset( + ( + "args", + "created", + "exc_info", + "exc_text", + "filename", + "funcName", + "levelname", + "levelno", + "linenno", + "lineno", + "message", + "module", + "msecs", + "msg", + "name", + "pathname", + "process", + "processName", + "relativeCreated", + "stack", + "tags", + "taskName", + "thread", + "threadName", + "stack_info", + ) + ) + + def _can_record(self, record: "LogRecord") -> bool: + """Prevents ignored loggers from recording""" + for logger in _IGNORED_LOGGERS: + if fnmatch(record.name.strip(), logger): + return False + return True + + def _logging_to_event_level(self, record: "LogRecord") -> str: + return LOGGING_TO_EVENT_LEVEL.get( + record.levelno, record.levelname.lower() if record.levelname else "" + ) + + def _extra_from_record(self, record: "LogRecord") -> "MutableMapping[str, object]": + return { + k: v + for k, v in vars(record).items() + if k not in self.COMMON_RECORD_ATTRS + and (not isinstance(k, str) or not k.startswith("_")) + } + + +class EventHandler(_BaseHandler): + """ + A logging handler that emits Sentry events for each log record + + Note that you do not have to use this class if the logging integration is enabled, which it is by default. + """ + + def emit(self, record: "LogRecord") -> "Any": + with capture_internal_exceptions(): + self.format(record) + return self._emit(record) + + def _emit(self, record: "LogRecord") -> None: + if not self._can_record(record): + return + + client = sentry_sdk.get_client() + if not client.is_active(): + return + + client_options = client.options + + # exc_info might be None or (None, None, None) + # + # exc_info may also be any falsy value due to Python stdlib being + # liberal with what it receives and Celery's billiard being "liberal" + # with what it sends. See + # https://github.com/getsentry/sentry-python/issues/904 + if record.exc_info and record.exc_info[0] is not None: + event, hint = event_from_exception( + record.exc_info, + client_options=client_options, + mechanism={"type": "logging", "handled": True}, + ) + elif (record.exc_info and record.exc_info[0] is None) or record.stack_info: + event = {} + hint = {} + with capture_internal_exceptions(): + event["threads"] = { + "values": [ + { + "stacktrace": current_stacktrace( + include_local_variables=client_options[ + "include_local_variables" + ], + max_value_length=client_options["max_value_length"], + ), + "crashed": False, + "current": True, + } + ] + } + else: + event = {} + hint = {} + + hint["log_record"] = record + + level = self._logging_to_event_level(record) + if level in {"debug", "info", "warning", "error", "critical", "fatal"}: + event["level"] = level # type: ignore[typeddict-item] + event["logger"] = record.name + + if ( + sys.version_info < (3, 11) + and record.name == "py.warnings" + and record.msg == "%s" + ): + # warnings module on Python 3.10 and below sets record.msg to "%s" + # and record.args[0] to the actual warning message. + # This was fixed in https://github.com/python/cpython/pull/30975. + message = record.args[0] + params = () + else: + message = record.msg + params = record.args + + event["logentry"] = { + "message": to_string(message), + "formatted": record.getMessage(), + "params": params, + } + + event["extra"] = self._extra_from_record(record) + + sentry_sdk.capture_event(event, hint=hint) + + +# Legacy name +SentryHandler = EventHandler + + +class BreadcrumbHandler(_BaseHandler): + """ + A logging handler that records breadcrumbs for each log record. + + Note that you do not have to use this class if the logging integration is enabled, which it is by default. + """ + + def emit(self, record: "LogRecord") -> "Any": + with capture_internal_exceptions(): + self.format(record) + return self._emit(record) + + def _emit(self, record: "LogRecord") -> None: + if not self._can_record(record): + return + + sentry_sdk.add_breadcrumb( + self._breadcrumb_from_record(record), hint={"log_record": record} + ) + + def _breadcrumb_from_record(self, record: "LogRecord") -> "Dict[str, Any]": + return { + "type": "log", + "level": self._logging_to_event_level(record), + "category": record.name, + "message": record.message, + "timestamp": datetime.fromtimestamp(record.created, timezone.utc), + "data": self._extra_from_record(record), + } + + +class SentryLogsHandler(_BaseHandler): + """ + A logging handler that records Sentry logs for each Python log record. + + Note that you do not have to use this class if the logging integration is enabled, which it is by default. + """ + + def emit(self, record: "LogRecord") -> "Any": + with capture_internal_exceptions(): + self.format(record) + if not self._can_record(record): + return + + client = sentry_sdk.get_client() + if not client.is_active(): + return + + if not has_logs_enabled(client.options): + return + + self._capture_log_from_record(client, record) + + def _capture_log_from_record( + self, client: "BaseClient", record: "LogRecord" + ) -> None: + otel_severity_number, otel_severity_text = _log_level_to_otel( + record.levelno, SEVERITY_TO_OTEL_SEVERITY + ) + project_root = client.options["project_root"] + + attrs: "Any" = self._extra_from_record(record) + attrs["sentry.origin"] = "auto.log.stdlib" + + parameters_set = False + if record.args is not None: + if isinstance(record.args, tuple): + parameters_set = bool(record.args) + for i, arg in enumerate(record.args): + attrs[f"sentry.message.parameter.{i}"] = ( + arg + if isinstance(arg, (str, float, int, bool)) + else safe_repr(arg) + ) + elif isinstance(record.args, dict): + parameters_set = bool(record.args) + for key, value in record.args.items(): + attrs[f"sentry.message.parameter.{key}"] = ( + value + if isinstance(value, (str, float, int, bool)) + else safe_repr(value) + ) + + if parameters_set and isinstance(record.msg, str): + # only include template if there is at least one + # sentry.message.parameter.X set + attrs["sentry.message.template"] = record.msg + + if record.lineno: + attrs["code.line.number"] = record.lineno + + if record.pathname: + if project_root is not None and record.pathname.startswith(project_root): + attrs["code.file.path"] = record.pathname[len(project_root) + 1 :] + else: + attrs["code.file.path"] = record.pathname + + if record.funcName: + attrs["code.function.name"] = record.funcName + + if record.thread: + attrs["thread.id"] = record.thread + if record.threadName: + attrs["thread.name"] = record.threadName + + if record.process: + attrs["process.pid"] = record.process + if record.processName: + attrs["process.executable.name"] = record.processName + if record.name: + attrs["logger.name"] = record.name + + # noinspection PyProtectedMember + sentry_sdk.get_current_scope()._capture_log( + { + "severity_text": otel_severity_text, + "severity_number": otel_severity_number, + "body": record.message, + "attributes": attrs, + "time_unix_nano": int(record.created * 1e9), + "trace_id": None, + "span_id": None, + }, + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/loguru.py b/lib/python3.12/site-packages/sentry_sdk/integrations/loguru.py new file mode 100644 index 0000000000000000000000000000000000000000..00bd3c022bd8fc41b6dd75ec5546c2bc00f7b7d8 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/loguru.py @@ -0,0 +1,209 @@ +import enum + +import sentry_sdk +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.integrations.logging import ( + BreadcrumbHandler, + EventHandler, + _BaseHandler, +) +from sentry_sdk.logger import _log_level_to_otel +from sentry_sdk.utils import has_logs_enabled, safe_repr + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from logging import LogRecord + from typing import Any, Optional + +try: + import loguru + from loguru import logger + from loguru._defaults import LOGURU_FORMAT as DEFAULT_FORMAT + + if TYPE_CHECKING: + from loguru import Message +except ImportError: + raise DidNotEnable("LOGURU is not installed") + + +class LoggingLevels(enum.IntEnum): + TRACE = 5 + DEBUG = 10 + INFO = 20 + SUCCESS = 25 + WARNING = 30 + ERROR = 40 + CRITICAL = 50 + + +DEFAULT_LEVEL = LoggingLevels.INFO.value +DEFAULT_EVENT_LEVEL = LoggingLevels.ERROR.value + + +SENTRY_LEVEL_FROM_LOGURU_LEVEL = { + "TRACE": "DEBUG", + "DEBUG": "DEBUG", + "INFO": "INFO", + "SUCCESS": "INFO", + "WARNING": "WARNING", + "ERROR": "ERROR", + "CRITICAL": "CRITICAL", +} + +# Map Loguru level numbers to corresponding OTel level numbers +SEVERITY_TO_OTEL_SEVERITY = { + LoggingLevels.CRITICAL: 21, # fatal + LoggingLevels.ERROR: 17, # error + LoggingLevels.WARNING: 13, # warn + LoggingLevels.SUCCESS: 11, # info + LoggingLevels.INFO: 9, # info + LoggingLevels.DEBUG: 5, # debug + LoggingLevels.TRACE: 1, # trace +} + + +class LoguruIntegration(Integration): + identifier = "loguru" + + level: "Optional[int]" = DEFAULT_LEVEL + event_level: "Optional[int]" = DEFAULT_EVENT_LEVEL + breadcrumb_format = DEFAULT_FORMAT + event_format = DEFAULT_FORMAT + sentry_logs_level: "Optional[int]" = DEFAULT_LEVEL + + def __init__( + self, + level: "Optional[int]" = DEFAULT_LEVEL, + event_level: "Optional[int]" = DEFAULT_EVENT_LEVEL, + breadcrumb_format: "str | loguru.FormatFunction" = DEFAULT_FORMAT, + event_format: "str | loguru.FormatFunction" = DEFAULT_FORMAT, + sentry_logs_level: "Optional[int]" = DEFAULT_LEVEL, + ) -> None: + LoguruIntegration.level = level + LoguruIntegration.event_level = event_level + LoguruIntegration.breadcrumb_format = breadcrumb_format + LoguruIntegration.event_format = event_format + LoguruIntegration.sentry_logs_level = sentry_logs_level + + @staticmethod + def setup_once() -> None: + if LoguruIntegration.level is not None: + logger.add( + LoguruBreadcrumbHandler(level=LoguruIntegration.level), + level=LoguruIntegration.level, + format=LoguruIntegration.breadcrumb_format, + ) + + if LoguruIntegration.event_level is not None: + logger.add( + LoguruEventHandler(level=LoguruIntegration.event_level), + level=LoguruIntegration.event_level, + format=LoguruIntegration.event_format, + ) + + if LoguruIntegration.sentry_logs_level is not None: + logger.add( + loguru_sentry_logs_handler, + level=LoguruIntegration.sentry_logs_level, + ) + + +class _LoguruBaseHandler(_BaseHandler): + def __init__(self, *args: "Any", **kwargs: "Any") -> None: + if kwargs.get("level"): + kwargs["level"] = SENTRY_LEVEL_FROM_LOGURU_LEVEL.get( + kwargs.get("level", ""), DEFAULT_LEVEL + ) + + super().__init__(*args, **kwargs) + + def _logging_to_event_level(self, record: "LogRecord") -> str: + try: + return SENTRY_LEVEL_FROM_LOGURU_LEVEL[ + LoggingLevels(record.levelno).name + ].lower() + except (ValueError, KeyError): + return record.levelname.lower() if record.levelname else "" + + +class LoguruEventHandler(_LoguruBaseHandler, EventHandler): + """Modified version of :class:`sentry_sdk.integrations.logging.EventHandler` to use loguru's level names.""" + + pass + + +class LoguruBreadcrumbHandler(_LoguruBaseHandler, BreadcrumbHandler): + """Modified version of :class:`sentry_sdk.integrations.logging.BreadcrumbHandler` to use loguru's level names.""" + + pass + + +def loguru_sentry_logs_handler(message: "Message") -> None: + # This is intentionally a callable sink instead of a standard logging handler + # since otherwise we wouldn't get direct access to message.record + client = sentry_sdk.get_client() + + if not client.is_active(): + return + + if not has_logs_enabled(client.options): + return + + record = message.record + + if ( + LoguruIntegration.sentry_logs_level is None + or record["level"].no < LoguruIntegration.sentry_logs_level + ): + return + + otel_severity_number, otel_severity_text = _log_level_to_otel( + record["level"].no, SEVERITY_TO_OTEL_SEVERITY + ) + + attrs: "dict[str, Any]" = {"sentry.origin": "auto.log.loguru"} + + project_root = client.options["project_root"] + if record.get("file"): + if project_root is not None and record["file"].path.startswith(project_root): + attrs["code.file.path"] = record["file"].path[len(project_root) + 1 :] + else: + attrs["code.file.path"] = record["file"].path + + if record.get("line") is not None: + attrs["code.line.number"] = record["line"] + + if record.get("function"): + attrs["code.function.name"] = record["function"] + + if record.get("thread"): + attrs["thread.name"] = record["thread"].name + attrs["thread.id"] = record["thread"].id + + if record.get("process"): + attrs["process.pid"] = record["process"].id + attrs["process.executable.name"] = record["process"].name + + if record.get("name"): + attrs["logger.name"] = record["name"] + + extra = record.get("extra") + if isinstance(extra, dict): + for key, value in extra.items(): + if isinstance(value, (str, int, float, bool)): + attrs[f"sentry.message.parameter.{key}"] = value + else: + attrs[f"sentry.message.parameter.{key}"] = safe_repr(value) + + sentry_sdk.get_current_scope()._capture_log( + { + "severity_text": otel_severity_text, + "severity_number": otel_severity_number, + "body": record["message"], + "attributes": attrs, + "time_unix_nano": int(record["time"].timestamp() * 1e9), + "trace_id": None, + "span_id": None, + } + ) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/mcp.py b/lib/python3.12/site-packages/sentry_sdk/integrations/mcp.py new file mode 100644 index 0000000000000000000000000000000000000000..58f9cd94e993167433bdd5cfe3993f0092e183cd --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/mcp.py @@ -0,0 +1,636 @@ +""" +Sentry integration for MCP (Model Context Protocol) servers. + +This integration instruments MCP servers to create spans for tool, prompt, +and resource handler execution, and captures errors that occur during execution. + +Supports the low-level `mcp.server.lowlevel.Server` API. +""" + +import inspect +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.ai.utils import get_start_span_function +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.utils import safe_serialize +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.integrations._wsgi_common import nullcontext + +try: + from mcp.server.lowlevel import Server # type: ignore[import-not-found] + from mcp.server.lowlevel.server import request_ctx # type: ignore[import-not-found] + from mcp.server.streamable_http import StreamableHTTPServerTransport # type: ignore[import-not-found] +except ImportError: + raise DidNotEnable("MCP SDK not installed") + +try: + from fastmcp import FastMCP # type: ignore[import-not-found] +except ImportError: + FastMCP = None + + +if TYPE_CHECKING: + from typing import Any, Callable, Optional, Tuple, ContextManager + + from starlette.types import Receive, Scope, Send # type: ignore[import-not-found] + + +class MCPIntegration(Integration): + identifier = "mcp" + origin = "auto.ai.mcp" + + def __init__(self, include_prompts: bool = True) -> None: + """ + Initialize the MCP integration. + + Args: + include_prompts: Whether to include prompts (tool results and prompt content) + in span data. Requires send_default_pii=True. Default is True. + """ + self.include_prompts = include_prompts + + @staticmethod + def setup_once() -> None: + """ + Patches MCP server classes to instrument handler execution. + """ + _patch_lowlevel_server() + _patch_handle_request() + + if FastMCP is not None: + _patch_fastmcp() + + +def _get_active_http_scopes() -> ( + "Optional[Tuple[Optional[sentry_sdk.Scope], Optional[sentry_sdk.Scope]]]" +): + try: + ctx = request_ctx.get() + except LookupError: + return None + + if ( + ctx is None + or not hasattr(ctx, "request") + or ctx.request is None + or "state" not in ctx.request.scope + ): + return None + + return ( + ctx.request.scope["state"].get("sentry_sdk.isolation_scope"), + ctx.request.scope["state"].get("sentry_sdk.current_scope"), + ) + + +def _get_request_context_data() -> "tuple[Optional[str], Optional[str], str]": + """ + Extract request ID, session ID, and MCP transport type from the request context. + + Returns: + Tuple of (request_id, session_id, mcp_transport). + - request_id: May be None if not available + - session_id: May be None if not available + - mcp_transport: "http", "sse", "stdio" + """ + request_id: "Optional[str]" = None + session_id: "Optional[str]" = None + mcp_transport: str = "stdio" + + try: + ctx = request_ctx.get() + + if ctx is not None: + request_id = ctx.request_id + if hasattr(ctx, "request") and ctx.request is not None: + request = ctx.request + # Detect transport type by checking request characteristics + if hasattr(request, "query_params") and request.query_params.get( + "session_id" + ): + # SSE transport uses query parameter + mcp_transport = "sse" + session_id = request.query_params.get("session_id") + elif hasattr(request, "headers") and request.headers.get( + "mcp-session-id" + ): + # StreamableHTTP transport uses header + mcp_transport = "http" + session_id = request.headers.get("mcp-session-id") + + except LookupError: + # No request context available - default to stdio + pass + + return request_id, session_id, mcp_transport + + +def _get_span_config( + handler_type: str, item_name: str +) -> "tuple[str, str, str, Optional[str]]": + """ + Get span configuration based on handler type. + + Returns: + Tuple of (span_data_key, span_name, mcp_method_name, result_data_key) + Note: result_data_key is None for resources + """ + if handler_type == "tool": + span_data_key = SPANDATA.MCP_TOOL_NAME + mcp_method_name = "tools/call" + result_data_key = SPANDATA.MCP_TOOL_RESULT_CONTENT + elif handler_type == "prompt": + span_data_key = SPANDATA.MCP_PROMPT_NAME + mcp_method_name = "prompts/get" + result_data_key = SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT + else: # resource + span_data_key = SPANDATA.MCP_RESOURCE_URI + mcp_method_name = "resources/read" + result_data_key = None # Resources don't capture result content + + span_name = f"{mcp_method_name} {item_name}" + return span_data_key, span_name, mcp_method_name, result_data_key + + +def _set_span_input_data( + span: "Any", + handler_name: str, + span_data_key: str, + mcp_method_name: str, + arguments: "dict[str, Any]", + request_id: "Optional[str]", + session_id: "Optional[str]", + mcp_transport: str, +) -> None: + """Set input span data for MCP handlers.""" + + # Set handler identifier + span.set_data(span_data_key, handler_name) + span.set_data(SPANDATA.MCP_METHOD_NAME, mcp_method_name) + + # Set transport/MCP transport type + span.set_data( + SPANDATA.NETWORK_TRANSPORT, "pipe" if mcp_transport == "stdio" else "tcp" + ) + span.set_data(SPANDATA.MCP_TRANSPORT, mcp_transport) + + # Set request_id if provided + if request_id: + span.set_data(SPANDATA.MCP_REQUEST_ID, request_id) + + # Set session_id if provided + if session_id: + span.set_data(SPANDATA.MCP_SESSION_ID, session_id) + + # Set request arguments (excluding common request context objects) + for k, v in arguments.items(): + span.set_data(f"mcp.request.argument.{k}", safe_serialize(v)) + + +def _extract_tool_result_content(result: "Any") -> "Any": + """ + Extract meaningful content from MCP tool result. + + Tool handlers can return: + - tuple (UnstructuredContent, StructuredContent): Return the structured content (dict) + - dict (StructuredContent): Return as-is + - Iterable (UnstructuredContent): Extract text from content blocks + """ + if result is None: + return None + + # Handle CombinationContent: tuple of (UnstructuredContent, StructuredContent) + if isinstance(result, tuple) and len(result) == 2: + # Return the structured content (2nd element) + return result[1] + + # Handle StructuredContent: dict + if isinstance(result, dict): + return result + + # Handle UnstructuredContent: iterable of ContentBlock objects + # Try to extract text content + if hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict)): + texts = [] + try: + for item in result: + # Try to get text attribute from ContentBlock objects + if hasattr(item, "text"): + texts.append(item.text) + elif isinstance(item, dict) and "text" in item: + texts.append(item["text"]) + except Exception: + # If extraction fails, return the original + return result + return " ".join(texts) if texts else result + + return result + + +def _set_span_output_data( + span: "Any", result: "Any", result_data_key: "Optional[str]", handler_type: str +) -> None: + """Set output span data for MCP handlers.""" + if result is None: + return + + # Get integration to check PII settings + integration = sentry_sdk.get_client().get_integration(MCPIntegration) + if integration is None: + return + + # Check if we should include sensitive data + should_include_data = should_send_default_pii() and integration.include_prompts + + # For tools, extract the meaningful content + if handler_type == "tool": + extracted = _extract_tool_result_content(result) + if extracted is not None and should_include_data: + span.set_data(result_data_key, safe_serialize(extracted)) + # Set content count if result is a dict + if isinstance(extracted, dict): + span.set_data(SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, len(extracted)) + elif handler_type == "prompt": + # For prompts, count messages and set role/content only for single-message prompts + try: + messages: "Optional[list[str]]" = None + message_count = 0 + + # Check if result has messages attribute (GetPromptResult) + if hasattr(result, "messages") and result.messages: + messages = result.messages + message_count = len(messages) + # Also check if result is a dict with messages + elif isinstance(result, dict) and result.get("messages"): + messages = result["messages"] + message_count = len(messages) + + # Always set message count if we found messages + if message_count > 0: + span.set_data(SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count) + + # Only set role and content for single-message prompts if PII is allowed + if message_count == 1 and should_include_data and messages: + first_message = messages[0] + # Extract role + role = None + if hasattr(first_message, "role"): + role = first_message.role + elif isinstance(first_message, dict) and "role" in first_message: + role = first_message["role"] + + if role: + span.set_data(SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role) + + # Extract content text + content_text = None + if hasattr(first_message, "content"): + msg_content = first_message.content + # Content can be a TextContent object or similar + if hasattr(msg_content, "text"): + content_text = msg_content.text + elif isinstance(msg_content, dict) and "text" in msg_content: + content_text = msg_content["text"] + elif isinstance(msg_content, str): + content_text = msg_content + elif isinstance(first_message, dict) and "content" in first_message: + msg_content = first_message["content"] + if isinstance(msg_content, dict) and "text" in msg_content: + content_text = msg_content["text"] + elif isinstance(msg_content, str): + content_text = msg_content + + if content_text: + span.set_data(result_data_key, content_text) + except Exception: + # Silently ignore if we can't extract message info + pass + # Resources don't capture result content (result_data_key is None) + + +# Handler data preparation and wrapping + + +def _prepare_handler_data( + handler_type: str, + original_args: "tuple[Any, ...]", + original_kwargs: "Optional[dict[str, Any]]" = None, +) -> "tuple[str, dict[str, Any], str, str, str, Optional[str]]": + """ + Prepare common handler data for both async and sync wrappers. + + Returns: + Tuple of (handler_name, arguments, span_data_key, span_name, mcp_method_name, result_data_key) + """ + original_kwargs = original_kwargs or {} + + # Extract handler-specific data based on handler type + if handler_type == "tool": + if original_args: + handler_name = original_args[0] + elif original_kwargs.get("name"): + handler_name = original_kwargs["name"] + + arguments = {} + if len(original_args) > 1: + arguments = original_args[1] + elif original_kwargs.get("arguments"): + arguments = original_kwargs["arguments"] + + elif handler_type == "prompt": + if original_args: + handler_name = original_args[0] + elif original_kwargs.get("name"): + handler_name = original_kwargs["name"] + + arguments = {} + if len(original_args) > 1: + arguments = original_args[1] + elif original_kwargs.get("arguments"): + arguments = original_kwargs["arguments"] + + # Include name in arguments dict for span data + arguments = {"name": handler_name, **(arguments or {})} + + else: # resource + handler_name = "unknown" + if original_args: + handler_name = str(original_args[0]) + elif original_kwargs.get("uri"): + handler_name = str(original_kwargs["uri"]) + + arguments = {} + + # Get span configuration + span_data_key, span_name, mcp_method_name, result_data_key = _get_span_config( + handler_type, handler_name + ) + + return ( + handler_name, + arguments, + span_data_key, + span_name, + mcp_method_name, + result_data_key, + ) + + +async def _handler_wrapper( + handler_type: str, + func: "Callable[..., Any]", + original_args: "tuple[Any, ...]", + original_kwargs: "Optional[dict[str, Any]]" = None, + self: "Optional[Any]" = None, + force_await: bool = True, +) -> "Any": + """ + Wrapper for MCP handlers. + + Args: + handler_type: "tool", "prompt", or "resource" + func: The handler function to wrap + original_args: Original arguments passed to the handler + original_kwargs: Original keyword arguments passed to the handler + self: Optional instance for bound methods + """ + if original_kwargs is None: + original_kwargs = {} + + ( + handler_name, + arguments, + span_data_key, + span_name, + mcp_method_name, + result_data_key, + ) = _prepare_handler_data(handler_type, original_args, original_kwargs) + + scopes = _get_active_http_scopes() + + isolation_scope_context: "ContextManager[Any]" + current_scope_context: "ContextManager[Any]" + + if scopes is None: + isolation_scope_context = nullcontext() + current_scope_context = nullcontext() + else: + isolation_scope, current_scope = scopes + + isolation_scope_context = ( + nullcontext() + if isolation_scope is None + else sentry_sdk.scope.use_isolation_scope(isolation_scope) + ) + current_scope_context = ( + nullcontext() + if current_scope is None + else sentry_sdk.scope.use_scope(current_scope) + ) + + # Get request ID, session ID, and transport from context + request_id, session_id, mcp_transport = _get_request_context_data() + + # Start span and execute + with isolation_scope_context: + with current_scope_context: + with get_start_span_function()( + op=OP.MCP_SERVER, + name=span_name, + origin=MCPIntegration.origin, + ) as span: + # Set input span data + _set_span_input_data( + span, + handler_name, + span_data_key, + mcp_method_name, + arguments, + request_id, + session_id, + mcp_transport, + ) + + # For resources, extract and set protocol + if handler_type == "resource": + if original_args: + uri = original_args[0] + else: + uri = original_kwargs.get("uri") + + protocol = None + if hasattr(uri, "scheme"): + protocol = uri.scheme + elif handler_name and "://" in handler_name: + protocol = handler_name.split("://")[0] + if protocol: + span.set_data(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) + + try: + # Execute the async handler + if self is not None: + original_args = (self, *original_args) + + result = func(*original_args, **original_kwargs) + if force_await or inspect.isawaitable(result): + result = await result + + except Exception as e: + # Set error flag for tools + if handler_type == "tool": + span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True) + sentry_sdk.capture_exception(e) + raise + + _set_span_output_data(span, result, result_data_key, handler_type) + + return result + + +def _create_instrumented_decorator( + original_decorator: "Callable[..., Any]", + handler_type: str, + *decorator_args: "Any", + **decorator_kwargs: "Any", +) -> "Callable[..., Any]": + """ + Create an instrumented version of an MCP decorator. + + This function intercepts MCP decorators (like @server.call_tool()) and injects + Sentry instrumentation into the handler registration flow. The returned decorator + will: + 1. Receive the user's handler function + 2. Pass the instrumented version to the original MCP decorator + + This ensures that when the handler is called at runtime, it's already wrapped + with Sentry spans and metrics collection. + + Args: + original_decorator: The original MCP decorator method (e.g., Server.call_tool) + handler_type: "tool", "prompt", or "resource" - determines span configuration + decorator_args: Positional arguments to pass to the original decorator (e.g., self) + decorator_kwargs: Keyword arguments to pass to the original decorator + + Returns: + A decorator function that instruments handlers before registering them + """ + + def instrumented_decorator(func: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(func) + async def wrapper(*args: "Any") -> "Any": + return await _handler_wrapper(handler_type, func, args, force_await=False) + + # Then register it with the original MCP decorator + return original_decorator(*decorator_args, **decorator_kwargs)(wrapper) + + return instrumented_decorator + + +def _patch_lowlevel_server() -> None: + """ + Patches the mcp.server.lowlevel.Server class to instrument handler execution. + """ + # Patch call_tool decorator + original_call_tool = Server.call_tool + + def patched_call_tool( + self: "Server", **kwargs: "Any" + ) -> "Callable[[Callable[..., Any]], Callable[..., Any]]": + """Patched version of Server.call_tool that adds Sentry instrumentation.""" + return lambda func: _create_instrumented_decorator( + original_call_tool, "tool", self, **kwargs + )(func) + + Server.call_tool = patched_call_tool + + # Patch get_prompt decorator + original_get_prompt = Server.get_prompt + + def patched_get_prompt( + self: "Server", + ) -> "Callable[[Callable[..., Any]], Callable[..., Any]]": + """Patched version of Server.get_prompt that adds Sentry instrumentation.""" + return lambda func: _create_instrumented_decorator( + original_get_prompt, "prompt", self + )(func) + + Server.get_prompt = patched_get_prompt + + # Patch read_resource decorator + original_read_resource = Server.read_resource + + def patched_read_resource( + self: "Server", + ) -> "Callable[[Callable[..., Any]], Callable[..., Any]]": + """Patched version of Server.read_resource that adds Sentry instrumentation.""" + return lambda func: _create_instrumented_decorator( + original_read_resource, "resource", self + )(func) + + Server.read_resource = patched_read_resource + + +def _patch_handle_request() -> None: + original_handle_request = StreamableHTTPServerTransport.handle_request + + @wraps(original_handle_request) + async def patched_handle_request( + self: "StreamableHTTPServerTransport", + scope: "Scope", + receive: "Receive", + send: "Send", + ) -> None: + scope.setdefault("state", {})["sentry_sdk.isolation_scope"] = ( + sentry_sdk.get_isolation_scope() + ) + scope["state"]["sentry_sdk.current_scope"] = sentry_sdk.get_current_scope() + await original_handle_request(self, scope, receive, send) + + StreamableHTTPServerTransport.handle_request = patched_handle_request + + +def _patch_fastmcp() -> None: + """ + Patches the standalone fastmcp package's FastMCP class. + + The standalone fastmcp package (v2.14.0+) registers its own handlers for + prompts and resources directly, bypassing the Server decorators we patch. + This function patches the _get_prompt_mcp and _read_resource_mcp methods + to add instrumentation for those handlers. + """ + if hasattr(FastMCP, "_get_prompt_mcp"): + original_get_prompt_mcp = FastMCP._get_prompt_mcp + + @wraps(original_get_prompt_mcp) + async def patched_get_prompt_mcp( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + return await _handler_wrapper( + "prompt", + original_get_prompt_mcp, + args, + kwargs, + self, + ) + + FastMCP._get_prompt_mcp = patched_get_prompt_mcp + + if hasattr(FastMCP, "_read_resource_mcp"): + original_read_resource_mcp = FastMCP._read_resource_mcp + + @wraps(original_read_resource_mcp) + async def patched_read_resource_mcp( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> "Any": + return await _handler_wrapper( + "resource", + original_read_resource_mcp, + args, + kwargs, + self, + ) + + FastMCP._read_resource_mcp = patched_read_resource_mcp diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/modules.py b/lib/python3.12/site-packages/sentry_sdk/integrations/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..086f537030316c8bb96101b10e5c9e33575204a4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/modules.py @@ -0,0 +1,27 @@ +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import add_global_event_processor +from sentry_sdk.utils import _get_installed_modules + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from sentry_sdk._types import Event + + +class ModulesIntegration(Integration): + identifier = "modules" + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def processor(event: "Event", hint: "Any") -> "Event": + if event.get("type") == "transaction": + return event + + if sentry_sdk.get_client().get_integration(ModulesIntegration) is None: + return event + + event["modules"] = _get_installed_modules() + return event diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/openai.py b/lib/python3.12/site-packages/sentry_sdk/integrations/openai.py new file mode 100644 index 0000000000000000000000000000000000000000..863f146a514352d2cb852b93c52a3677ad5d6484 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/openai.py @@ -0,0 +1,940 @@ +import sys +import json +import time +from functools import wraps +from collections.abc import Iterable + +import sentry_sdk +from sentry_sdk import consts +from sentry_sdk.ai.monitoring import record_token_usage +from sentry_sdk.ai.utils import ( + set_data_normalized, + normalize_message_roles, + truncate_and_annotate_messages, + truncate_and_annotate_embedding_inputs, +) +from sentry_sdk.ai._openai_completions_api import ( + _is_system_instruction as _is_system_instruction_completions, + _get_system_instructions as _get_system_instructions_completions, + _transform_system_instructions, + _get_text_items, +) +from sentry_sdk.ai._openai_responses_api import ( + _is_system_instruction as _is_system_instruction_responses, + _get_system_instructions as _get_system_instructions_responses, +) +from sentry_sdk.consts import SPANDATA +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing_utils import set_span_errored +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, + safe_serialize, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import ( + Any, + List, + Optional, + Callable, + AsyncIterator, + Iterator, + Union, + Iterable, + ) + from sentry_sdk.tracing import Span + from sentry_sdk._types import TextPart + + from openai.types.responses import ResponseInputParam + from openai import Omit + +try: + try: + from openai import NotGiven + except ImportError: + NotGiven = None + + try: + from openai import Omit + except ImportError: + Omit = None + + from openai.resources.chat.completions import Completions, AsyncCompletions + from openai.resources import Embeddings, AsyncEmbeddings + + if TYPE_CHECKING: + from openai.types.chat import ( + ChatCompletionMessageParam, + ChatCompletionChunk, + ChatCompletionSystemMessageParam, + ) +except ImportError: + raise DidNotEnable("OpenAI not installed") + +RESPONSES_API_ENABLED = True +try: + # responses API support was introduced in v1.66.0 + from openai.resources.responses import Responses, AsyncResponses + from openai.types.responses.response_completed_event import ResponseCompletedEvent +except ImportError: + RESPONSES_API_ENABLED = False + + +class OpenAIIntegration(Integration): + identifier = "openai" + origin = f"auto.ai.{identifier}" + + def __init__( + self: "OpenAIIntegration", + include_prompts: bool = True, + tiktoken_encoding_name: "Optional[str]" = None, + ) -> None: + self.include_prompts = include_prompts + + self.tiktoken_encoding = None + if tiktoken_encoding_name is not None: + import tiktoken # type: ignore + + self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name) + + @staticmethod + def setup_once() -> None: + Completions.create = _wrap_chat_completion_create(Completions.create) + AsyncCompletions.create = _wrap_async_chat_completion_create( + AsyncCompletions.create + ) + + Embeddings.create = _wrap_embeddings_create(Embeddings.create) + AsyncEmbeddings.create = _wrap_async_embeddings_create(AsyncEmbeddings.create) + + if RESPONSES_API_ENABLED: + Responses.create = _wrap_responses_create(Responses.create) + AsyncResponses.create = _wrap_async_responses_create(AsyncResponses.create) + + def count_tokens(self: "OpenAIIntegration", s: str) -> int: + if self.tiktoken_encoding is None: + return 0 + try: + return len(self.tiktoken_encoding.encode_ordinary(s)) + except Exception: + return 0 + + +def _capture_exception(exc: "Any", manual_span_cleanup: bool = True) -> None: + # Close an eventually open span + # We need to do this by hand because we are not using the start_span context manager + current_span = sentry_sdk.get_current_span() + set_span_errored(current_span) + + if manual_span_cleanup and current_span is not None: + current_span.__exit__(None, None, None) + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "openai", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + +def _get_usage(usage: "Any", names: "List[str]") -> int: + for name in names: + if hasattr(usage, name) and isinstance(getattr(usage, name), int): + return getattr(usage, name) + return 0 + + +def _calculate_token_usage( + messages: "Optional[Iterable[ChatCompletionMessageParam]]", + response: "Any", + span: "Span", + streaming_message_responses: "Optional[List[str]]", + count_tokens: "Callable[..., Any]", +) -> None: + input_tokens: "Optional[int]" = 0 + input_tokens_cached: "Optional[int]" = 0 + output_tokens: "Optional[int]" = 0 + output_tokens_reasoning: "Optional[int]" = 0 + total_tokens: "Optional[int]" = 0 + + if hasattr(response, "usage"): + input_tokens = _get_usage(response.usage, ["input_tokens", "prompt_tokens"]) + if hasattr(response.usage, "input_tokens_details"): + input_tokens_cached = _get_usage( + response.usage.input_tokens_details, ["cached_tokens"] + ) + + output_tokens = _get_usage( + response.usage, ["output_tokens", "completion_tokens"] + ) + if hasattr(response.usage, "output_tokens_details"): + output_tokens_reasoning = _get_usage( + response.usage.output_tokens_details, ["reasoning_tokens"] + ) + + total_tokens = _get_usage(response.usage, ["total_tokens"]) + + # Manually count tokens + if input_tokens == 0: + for message in messages or []: + if isinstance(message, str): + input_tokens += count_tokens(message) + continue + elif isinstance(message, dict): + message_content = message.get("content") + if message_content is None: + continue + # Deliberate use of Completions function for both Completions and Responses input format. + text_items = _get_text_items(message_content) + input_tokens += sum(count_tokens(text) for text in text_items) + continue + + if output_tokens == 0: + if streaming_message_responses is not None: + for message in streaming_message_responses: + output_tokens += count_tokens(message) + elif hasattr(response, "choices"): + for choice in response.choices: + if hasattr(choice, "message") and hasattr(choice.message, "content"): + output_tokens += count_tokens(choice.message.content) + + # Do not set token data if it is 0 + input_tokens = input_tokens or None + input_tokens_cached = input_tokens_cached or None + output_tokens = output_tokens or None + output_tokens_reasoning = output_tokens_reasoning or None + total_tokens = total_tokens or None + + record_token_usage( + span, + input_tokens=input_tokens, + input_tokens_cached=input_tokens_cached, + output_tokens=output_tokens, + output_tokens_reasoning=output_tokens_reasoning, + total_tokens=total_tokens, + ) + + +def _get_input_messages( + kwargs: "dict[str, Any]", +) -> "Optional[Union[Iterable[Any], list[str]]]": + # Input messages (the prompt or data sent to the model) + messages = kwargs.get("messages") + if messages is None: + messages = kwargs.get("input") + + if isinstance(messages, str): + messages = [messages] + + return messages + + +def _commmon_set_input_data( + span: "Span", + kwargs: "dict[str, Any]", +) -> None: + # Input attributes: Common + set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "openai") + + # Input attributes: Optional + kwargs_keys_to_attributes = { + "model": SPANDATA.GEN_AI_REQUEST_MODEL, + "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING, + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + } + for key, attribute in kwargs_keys_to_attributes.items(): + value = kwargs.get(key) + + if value is not None and _is_given(value): + set_data_normalized(span, attribute, value) + + # Input attributes: Tools + tools = kwargs.get("tools") + if tools is not None and _is_given(tools) and len(tools) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + + +def _set_responses_api_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + explicit_instructions: "Union[Optional[str], Omit]" = kwargs.get("instructions") + messages: "Optional[Union[str, ResponseInputParam]]" = kwargs.get("input") + + if not should_send_default_pii() or not integration.include_prompts: + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + return + + if ( + messages is None + and explicit_instructions is not None + and _is_given(explicit_instructions) + ): + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps( + [ + { + "type": "text", + "content": explicit_instructions, + } + ] + ), + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + return + + if messages is None: + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + return + + instructions_text_parts: "list[TextPart]" = [] + if explicit_instructions is not None and _is_given(explicit_instructions): + instructions_text_parts.append( + { + "type": "text", + "content": explicit_instructions, + } + ) + + system_instructions = _get_system_instructions_responses(messages) + # Deliberate use of function accepting completions API type because + # of shared structure FOR THIS PURPOSE ONLY. + instructions_text_parts += _transform_system_instructions(system_instructions) + + if len(instructions_text_parts) > 0: + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(instructions_text_parts), + ) + + if isinstance(messages, str): + normalized_messages = normalize_message_roles([messages]) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + return + + non_system_messages = [ + message for message in messages if not _is_system_instruction_responses(message) + ] + if len(non_system_messages) > 0: + normalized_messages = normalize_message_roles(non_system_messages) + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "responses") + _commmon_set_input_data(span, kwargs) + + +def _set_completions_api_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + messages: "Optional[Union[str, Iterable[ChatCompletionMessageParam]]]" = kwargs.get( + "messages" + ) + + if ( + not should_send_default_pii() + or not integration.include_prompts + or messages is None + ): + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + _commmon_set_input_data(span, kwargs) + return + + system_instructions = _get_system_instructions_completions(messages) + if len(system_instructions) > 0: + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(_transform_system_instructions(system_instructions)), + ) + + if isinstance(messages, str): + normalized_messages = normalize_message_roles([messages]) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + _commmon_set_input_data(span, kwargs) + return + + non_system_messages = [ + message + for message in messages + if not _is_system_instruction_completions(message) + ] + if len(non_system_messages) > 0: + normalized_messages = normalize_message_roles(non_system_messages) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_messages(normalized_messages, span, scope) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat") + _commmon_set_input_data(span, kwargs) + + +def _set_embeddings_input_data( + span: "Span", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", +) -> None: + messages = _get_input_messages(kwargs) + + if ( + messages is not None + and len(messages) > 0 # type: ignore + and should_send_default_pii() + and integration.include_prompts + ): + normalized_messages = normalize_message_roles(messages) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_embedding_inputs( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False + ) + + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + _commmon_set_input_data(span, kwargs) + + +def _set_output_data( + span: "Span", + response: "Any", + kwargs: "dict[str, Any]", + integration: "OpenAIIntegration", + start_time: "Optional[float]" = None, + finish_span: bool = True, +) -> None: + if hasattr(response, "model"): + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model) + + # Input messages (the prompt or data sent to the model) + # used for the token usage calculation + messages = kwargs.get("messages") + if messages is None: + messages = kwargs.get("input") + + if messages is not None and isinstance(messages, str): + messages = [messages] + + ttft: "Optional[float]" = None + + if hasattr(response, "choices"): + if should_send_default_pii() and integration.include_prompts: + response_text = [ + choice.message.model_dump() + for choice in response.choices + if choice.message is not None + ] + if len(response_text) > 0: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_text) + + _calculate_token_usage(messages, response, span, None, integration.count_tokens) + + if finish_span: + span.__exit__(None, None, None) + + elif hasattr(response, "output"): + if should_send_default_pii() and integration.include_prompts: + output_messages: "dict[str, list[Any]]" = { + "response": [], + "tool": [], + } + + for output in response.output: + if output.type == "function_call": + output_messages["tool"].append(output.dict()) + elif output.type == "message": + for output_message in output.content: + try: + output_messages["response"].append(output_message.text) + except AttributeError: + # Unknown output message type, just return the json + output_messages["response"].append(output_message.dict()) + + if len(output_messages["tool"]) > 0: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + output_messages["tool"], + unpack=False, + ) + + if len(output_messages["response"]) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"] + ) + + _calculate_token_usage(messages, response, span, None, integration.count_tokens) + + if finish_span: + span.__exit__(None, None, None) + + elif hasattr(response, "_iterator"): + data_buf: "list[list[str]]" = [] # one for each choice + + old_iterator = response._iterator + + def new_iterator() -> "Iterator[ChatCompletionChunk]": + nonlocal ttft + count_tokens_manually = True + for x in old_iterator: + with capture_internal_exceptions(): + # OpenAI chat completion API + if hasattr(x, "choices"): + choice_index = 0 + for choice in x.choices: + if hasattr(choice, "delta") and hasattr( + choice.delta, "content" + ): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + content = choice.delta.content + if len(data_buf) <= choice_index: + data_buf.append([]) + data_buf[choice_index].append(content or "") + choice_index += 1 + + # OpenAI responses API + elif hasattr(x, "delta"): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + if len(data_buf) == 0: + data_buf.append([]) + data_buf[0].append(x.delta or "") + + # OpenAI responses API end of streaming response + if RESPONSES_API_ENABLED and isinstance(x, ResponseCompletedEvent): + _calculate_token_usage( + messages, + x.response, + span, + None, + integration.count_tokens, + ) + count_tokens_manually = False + + yield x + + with capture_internal_exceptions(): + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + if len(data_buf) > 0: + all_responses = ["".join(chunk) for chunk in data_buf] + if should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) + if count_tokens_manually: + _calculate_token_usage( + messages, + response, + span, + all_responses, + integration.count_tokens, + ) + + if finish_span: + span.__exit__(None, None, None) + + async def new_iterator_async() -> "AsyncIterator[ChatCompletionChunk]": + nonlocal ttft + count_tokens_manually = True + async for x in old_iterator: + with capture_internal_exceptions(): + # OpenAI chat completion API + if hasattr(x, "choices"): + choice_index = 0 + for choice in x.choices: + if hasattr(choice, "delta") and hasattr( + choice.delta, "content" + ): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + content = choice.delta.content + if len(data_buf) <= choice_index: + data_buf.append([]) + data_buf[choice_index].append(content or "") + choice_index += 1 + + # OpenAI responses API + elif hasattr(x, "delta"): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + if len(data_buf) == 0: + data_buf.append([]) + data_buf[0].append(x.delta or "") + + # OpenAI responses API end of streaming response + if RESPONSES_API_ENABLED and isinstance(x, ResponseCompletedEvent): + _calculate_token_usage( + messages, + x.response, + span, + None, + integration.count_tokens, + ) + count_tokens_manually = False + + yield x + + with capture_internal_exceptions(): + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + if len(data_buf) > 0: + all_responses = ["".join(chunk) for chunk in data_buf] + if should_send_default_pii() and integration.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) + if count_tokens_manually: + _calculate_token_usage( + messages, + response, + span, + all_responses, + integration.count_tokens, + ) + if finish_span: + span.__exit__(None, None, None) + + if str(type(response._iterator)) == "": + response._iterator = new_iterator_async() + else: + response._iterator = new_iterator() + else: + _calculate_token_usage(messages, response, span, None, integration.count_tokens) + if finish_span: + span.__exit__(None, None, None) + + +def _new_chat_completion_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return f(*args, **kwargs) + + if "messages" not in kwargs: + # invalid call (in all versions of openai), let it return error + return f(*args, **kwargs) + + try: + iter(kwargs["messages"]) + except TypeError: + # invalid call (in all versions), messages must be iterable + return f(*args, **kwargs) + + model = kwargs.get("model") + + span = sentry_sdk.start_span( + op=consts.OP.GEN_AI_CHAT, + name=f"chat {model}", + origin=OpenAIIntegration.origin, + ) + span.__enter__() + + _set_completions_api_input_data(span, kwargs, integration) + + start_time = time.perf_counter() + response = yield f, args, kwargs + + _set_output_data(span, response, kwargs, integration, start_time, finish_span=True) + + return response + + +def _wrap_chat_completion_create(f: "Callable[..., Any]") -> "Callable[..., Any]": + def _execute_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_chat_completion_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return e.value + + try: + try: + result = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None or "messages" not in kwargs: + # no "messages" means invalid call (in all versions of openai), let it return error + return f(*args, **kwargs) + + return _execute_sync(f, *args, **kwargs) + + return _sentry_patched_create_sync + + +def _wrap_async_chat_completion_create(f: "Callable[..., Any]") -> "Callable[..., Any]": + async def _execute_async(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_chat_completion_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return await e.value + + try: + try: + result = await f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + async def _sentry_patched_create_async(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None or "messages" not in kwargs: + # no "messages" means invalid call (in all versions of openai), let it return error + return await f(*args, **kwargs) + + return await _execute_async(f, *args, **kwargs) + + return _sentry_patched_create_async + + +def _new_embeddings_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return f(*args, **kwargs) + + model = kwargs.get("model") + + with sentry_sdk.start_span( + op=consts.OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model}", + origin=OpenAIIntegration.origin, + ) as span: + _set_embeddings_input_data(span, kwargs, integration) + + response = yield f, args, kwargs + + _set_output_data(span, response, kwargs, integration, finish_span=False) + + return response + + +def _wrap_embeddings_create(f: "Any") -> "Any": + def _execute_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_embeddings_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return e.value + + try: + try: + result = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e, manual_span_cleanup=False) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return f(*args, **kwargs) + + return _execute_sync(f, *args, **kwargs) + + return _sentry_patched_create_sync + + +def _wrap_async_embeddings_create(f: "Any") -> "Any": + async def _execute_async(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_embeddings_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return await e.value + + try: + try: + result = await f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e, manual_span_cleanup=False) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + async def _sentry_patched_create_async(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return await f(*args, **kwargs) + + return await _execute_async(f, *args, **kwargs) + + return _sentry_patched_create_async + + +def _new_responses_create_common(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return f(*args, **kwargs) + + model = kwargs.get("model") + + span = sentry_sdk.start_span( + op=consts.OP.GEN_AI_RESPONSES, + name=f"responses {model}", + origin=OpenAIIntegration.origin, + ) + span.__enter__() + + _set_responses_api_input_data(span, kwargs, integration) + + start_time = time.perf_counter() + response = yield f, args, kwargs + + _set_output_data(span, response, kwargs, integration, start_time, finish_span=True) + + return response + + +def _wrap_responses_create(f: "Any") -> "Any": + def _execute_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_responses_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return e.value + + try: + try: + result = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + def _sentry_patched_create_sync(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return f(*args, **kwargs) + + return _execute_sync(f, *args, **kwargs) + + return _sentry_patched_create_sync + + +def _wrap_async_responses_create(f: "Any") -> "Any": + async def _execute_async(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": + gen = _new_responses_create_common(f, *args, **kwargs) + + try: + f, args, kwargs = next(gen) + except StopIteration as e: + return await e.value + + try: + try: + result = await f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(e) + reraise(*exc_info) + + return gen.send(result) + except StopIteration as e: + return e.value + + @wraps(f) + async def _sentry_patched_responses_async(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(OpenAIIntegration) + if integration is None: + return await f(*args, **kwargs) + + return await _execute_async(f, *args, **kwargs) + + return _sentry_patched_responses_async + + +def _is_given(obj: "Any") -> bool: + """ + Check for givenness safely across different openai versions. + """ + if NotGiven is not None and isinstance(obj, NotGiven): + return False + if Omit is not None and isinstance(obj, Omit): + return False + return True diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/openfeature.py b/lib/python3.12/site-packages/sentry_sdk/integrations/openfeature.py new file mode 100644 index 0000000000000000000000000000000000000000..281604fe3829c6ddb3b903fb21536fedddb84fc2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/openfeature.py @@ -0,0 +1,34 @@ +from typing import TYPE_CHECKING, Any + +from sentry_sdk.feature_flags import add_feature_flag +from sentry_sdk.integrations import DidNotEnable, Integration + +try: + from openfeature import api + from openfeature.hook import Hook + + if TYPE_CHECKING: + from openfeature.hook import HookContext, HookHints +except ImportError: + raise DidNotEnable("OpenFeature is not installed") + + +class OpenFeatureIntegration(Integration): + identifier = "openfeature" + + @staticmethod + def setup_once() -> None: + # Register the hook within the global openfeature hooks list. + api.add_hooks(hooks=[OpenFeatureHook()]) + + +class OpenFeatureHook(Hook): + def after(self, hook_context: "Any", details: "Any", hints: "Any") -> None: + if isinstance(details.value, bool): + add_feature_flag(details.flag_key, details.value) + + def error( + self, hook_context: "HookContext", exception: Exception, hints: "HookHints" + ) -> None: + if isinstance(hook_context.default_value, bool): + add_feature_flag(hook_context.flag_key, hook_context.default_value) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/otlp.py b/lib/python3.12/site-packages/sentry_sdk/integrations/otlp.py new file mode 100644 index 0000000000000000000000000000000000000000..19c6099970f7169d1d522faf418b5dd1d0b6af36 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/otlp.py @@ -0,0 +1,217 @@ +from sentry_sdk import get_client, capture_event +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.scope import register_external_propagation_context +from sentry_sdk.utils import ( + Dsn, + logger, + event_from_exception, + capture_internal_exceptions, +) +from sentry_sdk.consts import VERSION, EndpointType +from sentry_sdk.tracing_utils import Baggage +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SENTRY_TRACE_HEADER_NAME, +) + +try: + from opentelemetry.propagate import set_global_textmap + from opentelemetry.sdk.trace import TracerProvider, Span + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + + from opentelemetry.trace import ( + get_current_span, + get_tracer_provider, + set_tracer_provider, + format_trace_id, + format_span_id, + SpanContext, + INVALID_SPAN_ID, + INVALID_TRACE_ID, + ) + + from opentelemetry.context import ( + Context, + get_current, + get_value, + ) + + from opentelemetry.propagators.textmap import ( + CarrierT, + Setter, + default_setter, + ) + + from sentry_sdk.integrations.opentelemetry.propagator import SentryPropagator + from sentry_sdk.integrations.opentelemetry.consts import SENTRY_BAGGAGE_KEY +except ImportError: + raise DidNotEnable("opentelemetry-distro[otlp] is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional, Dict, Any, Tuple + + +def otel_propagation_context() -> "Optional[Tuple[str, str]]": + """ + Get the (trace_id, span_id) from opentelemetry if exists. + """ + ctx = get_current_span().get_span_context() + + if ctx.trace_id == INVALID_TRACE_ID or ctx.span_id == INVALID_SPAN_ID: + return None + + return (format_trace_id(ctx.trace_id), format_span_id(ctx.span_id)) + + +def setup_otlp_traces_exporter(dsn: "Optional[str]" = None) -> None: + tracer_provider = get_tracer_provider() + + if not isinstance(tracer_provider, TracerProvider): + logger.debug("[OTLP] No TracerProvider configured by user, creating a new one") + tracer_provider = TracerProvider() + set_tracer_provider(tracer_provider) + + endpoint = None + headers = None + if dsn: + auth = Dsn(dsn).to_auth(f"sentry.python/{VERSION}") + endpoint = auth.get_api_url(EndpointType.OTLP_TRACES) + headers = {"X-Sentry-Auth": auth.to_header()} + logger.debug(f"[OTLP] Sending traces to {endpoint}") + + otlp_exporter = OTLPSpanExporter(endpoint=endpoint, headers=headers) + span_processor = BatchSpanProcessor(otlp_exporter) + tracer_provider.add_span_processor(span_processor) + + +_sentry_patched_exception = False + + +def setup_capture_exceptions() -> None: + """ + Intercept otel's Span.record_exception to automatically capture those exceptions in Sentry. + """ + global _sentry_patched_exception + _original_record_exception = Span.record_exception + + if _sentry_patched_exception: + return + + def _sentry_patched_record_exception( + self: "Span", exception: "BaseException", *args: "Any", **kwargs: "Any" + ) -> None: + otlp_integration = get_client().get_integration(OTLPIntegration) + if otlp_integration and otlp_integration.capture_exceptions: + with capture_internal_exceptions(): + event, hint = event_from_exception( + exception, + client_options=get_client().options, + mechanism={"type": OTLPIntegration.identifier, "handled": False}, + ) + capture_event(event, hint=hint) + + _original_record_exception(self, exception, *args, **kwargs) + + Span.record_exception = _sentry_patched_record_exception # type: ignore[method-assign] + _sentry_patched_exception = True + + +class SentryOTLPPropagator(SentryPropagator): + """ + We need to override the inject of the older propagator since that + is SpanProcessor based. + + !!! Note regarding baggage: + We cannot meaningfully populate a new baggage as a head SDK + when we are using OTLP since we don't have any sort of transaction semantic to + track state across a group of spans. + + For incoming baggage, we just pass it on as is so that case is correctly handled. + """ + + def inject( + self, + carrier: "CarrierT", + context: "Optional[Context]" = None, + setter: "Setter[CarrierT]" = default_setter, + ) -> None: + otlp_integration = get_client().get_integration(OTLPIntegration) + if otlp_integration is None: + return + + if context is None: + context = get_current() + + current_span = get_current_span(context) + current_span_context = current_span.get_span_context() + + if not current_span_context.is_valid: + return + + sentry_trace = _to_traceparent(current_span_context) + setter.set(carrier, SENTRY_TRACE_HEADER_NAME, sentry_trace) + + baggage = get_value(SENTRY_BAGGAGE_KEY, context) + if baggage is not None and isinstance(baggage, Baggage): + baggage_data = baggage.serialize() + if baggage_data: + setter.set(carrier, BAGGAGE_HEADER_NAME, baggage_data) + + +def _to_traceparent(span_context: "SpanContext") -> str: + """ + Helper method to generate the sentry-trace header. + """ + span_id = format_span_id(span_context.span_id) + trace_id = format_trace_id(span_context.trace_id) + sampled = span_context.trace_flags.sampled + + return f"{trace_id}-{span_id}-{'1' if sampled else '0'}" + + +class OTLPIntegration(Integration): + """ + Automatically setup OTLP ingestion from the DSN. + + :param setup_otlp_traces_exporter: Automatically configure an Exporter to send OTLP traces from the DSN, defaults to True. + Set to False if using a custom collector or to setup the TracerProvider manually. + :param setup_propagator: Automatically configure the Sentry Propagator for Distributed Tracing, defaults to True. + Set to False to configure propagators manually or to disable propagation. + :param capture_exceptions: Intercept and capture exceptions on the OpenTelemetry Span in Sentry as well, defaults to False. + Set to True to turn on capturing but be aware that since Sentry captures most exceptions, duplicate exceptions might be dropped by DedupeIntegration in many cases. + """ + + identifier = "otlp" + + def __init__( + self, + setup_otlp_traces_exporter: bool = True, + setup_propagator: bool = True, + capture_exceptions: bool = False, + ) -> None: + self.setup_otlp_traces_exporter = setup_otlp_traces_exporter + self.setup_propagator = setup_propagator + self.capture_exceptions = capture_exceptions + + @staticmethod + def setup_once() -> None: + logger.debug("[OTLP] Setting up trace linking for all events") + register_external_propagation_context(otel_propagation_context) + + def setup_once_with_options( + self, options: "Optional[Dict[str, Any]]" = None + ) -> None: + if self.setup_otlp_traces_exporter: + logger.debug("[OTLP] Setting up OTLP exporter") + dsn: "Optional[str]" = options.get("dsn") if options else None + setup_otlp_traces_exporter(dsn) + + if self.setup_propagator: + logger.debug("[OTLP] Setting up propagator for distributed tracing") + # TODO-neel better propagator support, chain with existing ones if possible instead of replacing + set_global_textmap(SentryOTLPPropagator()) + + setup_capture_exceptions() diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/pure_eval.py b/lib/python3.12/site-packages/sentry_sdk/integrations/pure_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..f8c7a9fba6616f6dcd75ebbdf441e41266f8bb1c --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/pure_eval.py @@ -0,0 +1,137 @@ +import ast + +import sentry_sdk +from sentry_sdk import serializer +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.scope import add_global_event_processor +from sentry_sdk.utils import walk_exception_chain, iter_stacks + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional, Dict, Any, Tuple, List + from types import FrameType + + from sentry_sdk._types import Event, Hint + +try: + from executing import Source +except ImportError: + raise DidNotEnable("executing is not installed") + +try: + from pure_eval import Evaluator +except ImportError: + raise DidNotEnable("pure_eval is not installed") + +try: + # Used implicitly, just testing it's available + import asttokens # noqa +except ImportError: + raise DidNotEnable("asttokens is not installed") + + +class PureEvalIntegration(Integration): + identifier = "pure_eval" + + @staticmethod + def setup_once() -> None: + @add_global_event_processor + def add_executing_info( + event: "Event", hint: "Optional[Hint]" + ) -> "Optional[Event]": + if sentry_sdk.get_client().get_integration(PureEvalIntegration) is None: + return event + + if hint is None: + return event + + exc_info = hint.get("exc_info", None) + + if exc_info is None: + return event + + exception = event.get("exception", None) + + if exception is None: + return event + + values = exception.get("values", None) + + if values is None: + return event + + for exception, (_exc_type, _exc_value, exc_tb) in zip( + reversed(values), walk_exception_chain(exc_info) + ): + sentry_frames = [ + frame + for frame in exception.get("stacktrace", {}).get("frames", []) + if frame.get("function") + ] + tbs = list(iter_stacks(exc_tb)) + if len(sentry_frames) != len(tbs): + continue + + for sentry_frame, tb in zip(sentry_frames, tbs): + sentry_frame["vars"] = ( + pure_eval_frame(tb.tb_frame) or sentry_frame["vars"] + ) + return event + + +def pure_eval_frame(frame: "FrameType") -> "Dict[str, Any]": + source = Source.for_frame(frame) + if not source.tree: + return {} + + statements = source.statements_at_line(frame.f_lineno) + if not statements: + return {} + + scope = stmt = list(statements)[0] + while True: + # Get the parent first in case the original statement is already + # a function definition, e.g. if we're calling a decorator + # In that case we still want the surrounding scope, not that function + scope = scope.parent + if isinstance(scope, (ast.FunctionDef, ast.ClassDef, ast.Module)): + break + + evaluator = Evaluator.from_frame(frame) + expressions = evaluator.interesting_expressions_grouped(scope) + + def closeness(expression: "Tuple[List[Any], Any]") -> "Tuple[int, int]": + # Prioritise expressions with a node closer to the statement executed + # without being after that statement + # A higher return value is better - the expression will appear + # earlier in the list of values and is less likely to be trimmed + nodes, _value = expression + + def start(n: "ast.expr") -> "Tuple[int, int]": + return (n.lineno, n.col_offset) + + nodes_before_stmt = [ + node + for node in nodes + if start(node) < stmt.last_token.end # type: ignore + ] + if nodes_before_stmt: + # The position of the last node before or in the statement + return max(start(node) for node in nodes_before_stmt) + else: + # The position of the first node after the statement + # Negative means it's always lower priority than nodes that come before + # Less negative means closer to the statement and higher priority + lineno, col_offset = min(start(node) for node in nodes) + return (-lineno, -col_offset) + + # This adds the first_token and last_token attributes to nodes + atok = source.asttokens() + + expressions.sort(key=closeness, reverse=True) + vars = { + atok.get_text(nodes[0]): value + for nodes, value in expressions[: serializer.MAX_DATABAG_BREADTH] + } + return serializer.serialize(vars, is_vars=True) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/pymongo.py b/lib/python3.12/site-packages/sentry_sdk/integrations/pymongo.py new file mode 100644 index 0000000000000000000000000000000000000000..86399b54d1fd3dc4300d9e9bc161dcef38b4b1db --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/pymongo.py @@ -0,0 +1,209 @@ +import copy +import json + +import sentry_sdk +from sentry_sdk.consts import SPANSTATUS, SPANDATA, OP +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import Span +from sentry_sdk.utils import capture_internal_exceptions + +try: + from pymongo import monitoring +except ImportError: + raise DidNotEnable("Pymongo not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Dict, Union + + from pymongo.monitoring import ( + CommandFailedEvent, + CommandStartedEvent, + CommandSucceededEvent, + ) + + +SAFE_COMMAND_ATTRIBUTES = [ + "insert", + "ordered", + "find", + "limit", + "singleBatch", + "aggregate", + "createIndexes", + "indexes", + "delete", + "findAndModify", + "renameCollection", + "to", + "drop", +] + + +def _strip_pii(command: "Dict[str, Any]") -> "Dict[str, Any]": + for key in command: + is_safe_field = key in SAFE_COMMAND_ATTRIBUTES + if is_safe_field: + # Skip if safe key + continue + + update_db_command = key == "update" and "findAndModify" not in command + if update_db_command: + # Also skip "update" db command because it is save. + # There is also an "update" key in the "findAndModify" command, which is NOT safe! + continue + + # Special stripping for documents + is_document = key == "documents" + if is_document: + for doc in command[key]: + for doc_key in doc: + doc[doc_key] = "%s" + continue + + # Special stripping for dict style fields + is_dict_field = key in ["filter", "query", "update"] + if is_dict_field: + for item_key in command[key]: + command[key][item_key] = "%s" + continue + + # For pipeline fields strip the `$match` dict + is_pipeline_field = key == "pipeline" + if is_pipeline_field: + for pipeline in command[key]: + for match_key in pipeline["$match"] if "$match" in pipeline else []: + pipeline["$match"][match_key] = "%s" + continue + + # Default stripping + command[key] = "%s" + + return command + + +def _get_db_data(event: "Any") -> "Dict[str, Any]": + data = {} + + data[SPANDATA.DB_SYSTEM] = "mongodb" + + db_name = event.database_name + if db_name is not None: + data[SPANDATA.DB_NAME] = db_name + + server_address = event.connection_id[0] + if server_address is not None: + data[SPANDATA.SERVER_ADDRESS] = server_address + + server_port = event.connection_id[1] + if server_port is not None: + data[SPANDATA.SERVER_PORT] = server_port + + return data + + +class CommandTracer(monitoring.CommandListener): + def __init__(self) -> None: + self._ongoing_operations: "Dict[int, Span]" = {} + + def _operation_key( + self, + event: "Union[CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent]", + ) -> int: + return event.request_id + + def started(self, event: "CommandStartedEvent") -> None: + if sentry_sdk.get_client().get_integration(PyMongoIntegration) is None: + return + + with capture_internal_exceptions(): + command = dict(copy.deepcopy(event.command)) + + command.pop("$db", None) + command.pop("$clusterTime", None) + command.pop("$signature", None) + + tags = { + "db.name": event.database_name, + SPANDATA.DB_SYSTEM: "mongodb", + SPANDATA.DB_OPERATION: event.command_name, + SPANDATA.DB_MONGODB_COLLECTION: command.get(event.command_name), + } + + try: + tags["net.peer.name"] = event.connection_id[0] + tags["net.peer.port"] = str(event.connection_id[1]) + except TypeError: + pass + + data: "Dict[str, Any]" = {"operation_ids": {}} + data["operation_ids"]["operation"] = event.operation_id + data["operation_ids"]["request"] = event.request_id + + data.update(_get_db_data(event)) + + try: + lsid = command.pop("lsid")["id"] + data["operation_ids"]["session"] = str(lsid) + except KeyError: + pass + + if not should_send_default_pii(): + command = _strip_pii(command) + + query = json.dumps(command, default=str) + span = sentry_sdk.start_span( + op=OP.DB, + name=query, + origin=PyMongoIntegration.origin, + ) + + for tag, value in tags.items(): + # set the tag for backwards-compatibility. + # TODO: remove the set_tag call in the next major release! + span.set_tag(tag, value) + + span.set_data(tag, value) + + for key, value in data.items(): + span.set_data(key, value) + + with capture_internal_exceptions(): + sentry_sdk.add_breadcrumb( + message=query, category="query", type=OP.DB, data=tags + ) + + self._ongoing_operations[self._operation_key(event)] = span.__enter__() + + def failed(self, event: "CommandFailedEvent") -> None: + if sentry_sdk.get_client().get_integration(PyMongoIntegration) is None: + return + + try: + span = self._ongoing_operations.pop(self._operation_key(event)) + span.set_status(SPANSTATUS.INTERNAL_ERROR) + span.__exit__(None, None, None) + except KeyError: + return + + def succeeded(self, event: "CommandSucceededEvent") -> None: + if sentry_sdk.get_client().get_integration(PyMongoIntegration) is None: + return + + try: + span = self._ongoing_operations.pop(self._operation_key(event)) + span.set_status(SPANSTATUS.OK) + span.__exit__(None, None, None) + except KeyError: + pass + + +class PyMongoIntegration(Integration): + identifier = "pymongo" + origin = f"auto.db.{identifier}" + + @staticmethod + def setup_once() -> None: + monitoring.register(CommandTracer()) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/pyramid.py b/lib/python3.12/site-packages/sentry_sdk/integrations/pyramid.py new file mode 100644 index 0000000000000000000000000000000000000000..82e629c86274472e0c761cf39717f34c8f90a5de --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/pyramid.py @@ -0,0 +1,223 @@ +import functools +import os +import sys +import weakref + +import sentry_sdk +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import RequestExtractor +from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + reraise, +) + +try: + from pyramid.httpexceptions import HTTPException + from pyramid.request import Request +except ImportError: + raise DidNotEnable("Pyramid not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pyramid.response import Response + from typing import Any + from sentry_sdk.integrations.wsgi import _ScopedResponse + from typing import Callable + from typing import Dict + from typing import Optional + from webob.cookies import RequestCookies + from webob.request import _FieldStorageWithFile + + from sentry_sdk.utils import ExcInfo + from sentry_sdk._types import Event, EventProcessor + + +if getattr(Request, "authenticated_userid", None): + + def authenticated_userid(request: "Request") -> "Optional[Any]": + return request.authenticated_userid + +else: + # bw-compat for pyramid < 1.5 + from pyramid.security import authenticated_userid # type: ignore + + +TRANSACTION_STYLE_VALUES = ("route_name", "route_pattern") + + +class PyramidIntegration(Integration): + identifier = "pyramid" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__(self, transaction_style: str = "route_name") -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + + @staticmethod + def setup_once() -> None: + from pyramid import router + + old_call_view = router._call_view + + @functools.wraps(old_call_view) + def sentry_patched_call_view( + registry: "Any", request: "Request", *args: "Any", **kwargs: "Any" + ) -> "Response": + integration = sentry_sdk.get_client().get_integration(PyramidIntegration) + if integration is None: + return old_call_view(registry, request, *args, **kwargs) + + _set_transaction_name_and_source( + sentry_sdk.get_current_scope(), integration.transaction_style, request + ) + scope = sentry_sdk.get_isolation_scope() + scope.add_event_processor( + _make_event_processor(weakref.ref(request), integration) + ) + + return old_call_view(registry, request, *args, **kwargs) + + router._call_view = sentry_patched_call_view + + if hasattr(Request, "invoke_exception_view"): + old_invoke_exception_view = Request.invoke_exception_view + + def sentry_patched_invoke_exception_view( + self: "Request", *args: "Any", **kwargs: "Any" + ) -> "Any": + rv = old_invoke_exception_view(self, *args, **kwargs) + + if ( + self.exc_info + and all(self.exc_info) + and rv.status_int == 500 + and sentry_sdk.get_client().get_integration(PyramidIntegration) + is not None + ): + _capture_exception(self.exc_info) + + return rv + + Request.invoke_exception_view = sentry_patched_invoke_exception_view + + old_wsgi_call = router.Router.__call__ + + @ensure_integration_enabled(PyramidIntegration, old_wsgi_call) + def sentry_patched_wsgi_call( + self: "Any", environ: "Dict[str, str]", start_response: "Callable[..., Any]" + ) -> "_ScopedResponse": + def sentry_patched_inner_wsgi_call( + environ: "Dict[str, Any]", start_response: "Callable[..., Any]" + ) -> "Any": + try: + return old_wsgi_call(self, environ, start_response) + except Exception: + einfo = sys.exc_info() + _capture_exception(einfo) + reraise(*einfo) + + middleware = SentryWsgiMiddleware( + sentry_patched_inner_wsgi_call, + span_origin=PyramidIntegration.origin, + ) + return middleware(environ, start_response) + + router.Router.__call__ = sentry_patched_wsgi_call + + +@ensure_integration_enabled(PyramidIntegration) +def _capture_exception(exc_info: "ExcInfo") -> None: + if exc_info[0] is None or issubclass(exc_info[0], HTTPException): + return + + event, hint = event_from_exception( + exc_info, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "pyramid", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) + + +def _set_transaction_name_and_source( + scope: "sentry_sdk.Scope", transaction_style: str, request: "Request" +) -> None: + try: + name_for_style = { + "route_name": request.matched_route.name, + "route_pattern": request.matched_route.pattern, + } + scope.set_transaction_name( + name_for_style[transaction_style], + source=SOURCE_FOR_STYLE[transaction_style], + ) + except Exception: + pass + + +class PyramidRequestExtractor(RequestExtractor): + def url(self) -> str: + return self.request.path_url + + def env(self) -> "Dict[str, str]": + return self.request.environ + + def cookies(self) -> "RequestCookies": + return self.request.cookies + + def raw_data(self) -> str: + return self.request.text + + def form(self) -> "Dict[str, str]": + return { + key: value + for key, value in self.request.POST.items() + if not getattr(value, "filename", None) + } + + def files(self) -> "Dict[str, _FieldStorageWithFile]": + return { + key: value + for key, value in self.request.POST.items() + if getattr(value, "filename", None) + } + + def size_of_file(self, postdata: "_FieldStorageWithFile") -> int: + file = postdata.file + try: + return os.fstat(file.fileno()).st_size + except Exception: + return 0 + + +def _make_event_processor( + weak_request: "Callable[[], Request]", integration: "PyramidIntegration" +) -> "EventProcessor": + def pyramid_event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event": + request = weak_request() + if request is None: + return event + + with capture_internal_exceptions(): + PyramidRequestExtractor(request).extract_into_event(event) + + if should_send_default_pii(): + with capture_internal_exceptions(): + user_info = event.setdefault("user", {}) + user_info.setdefault("id", authenticated_userid(request)) + + return event + + return pyramid_event_processor diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/quart.py b/lib/python3.12/site-packages/sentry_sdk/integrations/quart.py new file mode 100644 index 0000000000000000000000000000000000000000..c1b8fca717c22f6aabb9011291167e013d5dfbe3 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/quart.py @@ -0,0 +1,228 @@ +import asyncio +import inspect +from functools import wraps + +import sentry_sdk +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.integrations._wsgi_common import _filter_headers +from sentry_sdk.integrations.asgi import SentryAsgiMiddleware +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import SOURCE_FOR_STYLE +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, +) +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Union + + from sentry_sdk._types import Event, EventProcessor + +try: + import quart_auth # type: ignore +except ImportError: + quart_auth = None + +try: + from quart import ( # type: ignore + has_request_context, + has_websocket_context, + Request, + Quart, + request, + websocket, + ) + from quart.signals import ( # type: ignore + got_background_exception, + got_request_exception, + got_websocket_exception, + request_started, + websocket_started, + ) +except ImportError: + raise DidNotEnable("Quart is not installed") +else: + # Quart 0.19 is based on Flask and hence no longer has a Scaffold + try: + from quart.scaffold import Scaffold # type: ignore + except ImportError: + from flask.sansio.scaffold import Scaffold # type: ignore + +TRANSACTION_STYLE_VALUES = ("endpoint", "url") + + +class QuartIntegration(Integration): + identifier = "quart" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__(self, transaction_style: str = "endpoint") -> None: + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + + @staticmethod + def setup_once() -> None: + request_started.connect(_request_websocket_started) + websocket_started.connect(_request_websocket_started) + got_background_exception.connect(_capture_exception) + got_request_exception.connect(_capture_exception) + got_websocket_exception.connect(_capture_exception) + + patch_asgi_app() + patch_scaffold_route() + + +def patch_asgi_app() -> None: + old_app = Quart.__call__ + + async def sentry_patched_asgi_app( + self: "Any", scope: "Any", receive: "Any", send: "Any" + ) -> "Any": + if sentry_sdk.get_client().get_integration(QuartIntegration) is None: + return await old_app(self, scope, receive, send) + + middleware = SentryAsgiMiddleware( + lambda *a, **kw: old_app(self, *a, **kw), + span_origin=QuartIntegration.origin, + asgi_version=3, + ) + return await middleware(scope, receive, send) + + Quart.__call__ = sentry_patched_asgi_app + + +def patch_scaffold_route() -> None: + old_route = Scaffold.route + + def _sentry_route(*args: "Any", **kwargs: "Any") -> "Any": + old_decorator = old_route(*args, **kwargs) + + def decorator(old_func: "Any") -> "Any": + if inspect.isfunction(old_func) and not asyncio.iscoroutinefunction( + old_func + ): + + @wraps(old_func) + @ensure_integration_enabled(QuartIntegration, old_func) + def _sentry_func(*args: "Any", **kwargs: "Any") -> "Any": + current_scope = sentry_sdk.get_current_scope() + if current_scope.transaction is not None: + current_scope.transaction.update_active_thread() + + sentry_scope = sentry_sdk.get_isolation_scope() + if sentry_scope.profile is not None: + sentry_scope.profile.update_active_thread_id() + + return old_func(*args, **kwargs) + + return old_decorator(_sentry_func) + + return old_decorator(old_func) + + return decorator + + Scaffold.route = _sentry_route + + +def _set_transaction_name_and_source( + scope: "sentry_sdk.Scope", transaction_style: str, request: "Request" +) -> None: + try: + name_for_style = { + "url": request.url_rule.rule, + "endpoint": request.url_rule.endpoint, + } + scope.set_transaction_name( + name_for_style[transaction_style], + source=SOURCE_FOR_STYLE[transaction_style], + ) + except Exception: + pass + + +async def _request_websocket_started(app: "Quart", **kwargs: "Any") -> None: + integration = sentry_sdk.get_client().get_integration(QuartIntegration) + if integration is None: + return + + if has_request_context(): + request_websocket = request._get_current_object() + if has_websocket_context(): + request_websocket = websocket._get_current_object() + + # Set the transaction name here, but rely on ASGI middleware + # to actually start the transaction + _set_transaction_name_and_source( + sentry_sdk.get_current_scope(), integration.transaction_style, request_websocket + ) + + scope = sentry_sdk.get_isolation_scope() + evt_processor = _make_request_event_processor(app, request_websocket, integration) + scope.add_event_processor(evt_processor) + + +def _make_request_event_processor( + app: "Quart", request: "Request", integration: "QuartIntegration" +) -> "EventProcessor": + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + # if the request is gone we are fine not logging the data from + # it. This might happen if the processor is pushed away to + # another thread. + if request is None: + return event + + with capture_internal_exceptions(): + # TODO: Figure out what to do with request body. Methods on request + # are async, but event processors are not. + + request_info = event.setdefault("request", {}) + request_info["url"] = request.url + request_info["query_string"] = request.query_string + request_info["method"] = request.method + request_info["headers"] = _filter_headers(dict(request.headers)) + + if should_send_default_pii(): + request_info["env"] = {"REMOTE_ADDR": request.access_route[0]} + _add_user_to_event(event) + + return event + + return inner + + +async def _capture_exception( + sender: "Quart", exception: "Union[ValueError, BaseException]", **kwargs: "Any" +) -> None: + integration = sentry_sdk.get_client().get_integration(QuartIntegration) + if integration is None: + return + + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "quart", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) + + +def _add_user_to_event(event: "Event") -> None: + if quart_auth is None: + return + + user = quart_auth.current_user + if user is None: + return + + with capture_internal_exceptions(): + user_info = event.setdefault("user", {}) + + user_info["id"] = quart_auth.current_user._auth_id diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/ray.py b/lib/python3.12/site-packages/sentry_sdk/integrations/ray.py new file mode 100644 index 0000000000000000000000000000000000000000..d1559d5c197a8718dc6b8474c85b7c41ccfaa71d --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/ray.py @@ -0,0 +1,181 @@ +import inspect +import functools +import sys + +import sentry_sdk +from sentry_sdk.consts import OP, SPANSTATUS +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + event_from_exception, + logger, + package_version, + qualname_from_function, + reraise, +) + +try: + import ray # type: ignore[import-not-found] + from ray import remote +except ImportError: + raise DidNotEnable("Ray not installed.") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + from typing import Any, Optional + from sentry_sdk.utils import ExcInfo + + +def _check_sentry_initialized() -> None: + if sentry_sdk.get_client().is_active(): + return + + logger.debug( + "[Tracing] Sentry not initialized in ray cluster worker, performance data will be discarded." + ) + + +def _insert_sentry_tracing_in_signature(func: "Callable[..., Any]") -> None: + # Patching new_func signature to add the _sentry_tracing parameter to it + # Ray later inspects the signature and finds the unexpected parameter otherwise + signature = inspect.signature(func) + params = list(signature.parameters.values()) + sentry_tracing_param = inspect.Parameter( + "_sentry_tracing", + kind=inspect.Parameter.KEYWORD_ONLY, + default=None, + ) + + # Keyword only arguments are penultimate if function has variadic keyword arguments + if params and params[-1].kind is inspect.Parameter.VAR_KEYWORD: + params.insert(-1, sentry_tracing_param) + else: + params.append(sentry_tracing_param) + + func.__signature__ = signature.replace(parameters=params) # type: ignore[attr-defined] + + +def _patch_ray_remote() -> None: + old_remote = remote + + @functools.wraps(old_remote) + def new_remote( + f: "Optional[Callable[..., Any]]" = None, *args: "Any", **kwargs: "Any" + ) -> "Callable[..., Any]": + if inspect.isclass(f): + # Ray Actors + # (https://docs.ray.io/en/latest/ray-core/actors.html) + # are not supported + # (Only Ray Tasks are supported) + return old_remote(f, *args, **kwargs) + + def wrapper(user_f: "Callable[..., Any]") -> "Any": + if inspect.isclass(user_f): + # Ray Actors + # (https://docs.ray.io/en/latest/ray-core/actors.html) + # are not supported + # (Only Ray Tasks are supported) + return old_remote(*args, **kwargs)(user_f) + + @functools.wraps(user_f) + def new_func( + *f_args: "Any", + _sentry_tracing: "Optional[dict[str, Any]]" = None, + **f_kwargs: "Any", + ) -> "Any": + _check_sentry_initialized() + + transaction = sentry_sdk.continue_trace( + _sentry_tracing or {}, + op=OP.QUEUE_TASK_RAY, + name=qualname_from_function(user_f), + origin=RayIntegration.origin, + source=TransactionSource.TASK, + ) + + with sentry_sdk.start_transaction(transaction) as transaction: + try: + result = user_f(*f_args, **f_kwargs) + transaction.set_status(SPANSTATUS.OK) + except Exception: + transaction.set_status(SPANSTATUS.INTERNAL_ERROR) + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result + + _insert_sentry_tracing_in_signature(new_func) + + if f: + rv = old_remote(new_func) + else: + rv = old_remote(*args, **kwargs)(new_func) + old_remote_method = rv.remote + + def _remote_method_with_header_propagation( + *args: "Any", **kwargs: "Any" + ) -> "Any": + """ + Ray Client + """ + with sentry_sdk.start_span( + op=OP.QUEUE_SUBMIT_RAY, + name=qualname_from_function(user_f), + origin=RayIntegration.origin, + ) as span: + tracing = { + k: v + for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() + } + try: + result = old_remote_method( + *args, **kwargs, _sentry_tracing=tracing + ) + span.set_status(SPANSTATUS.OK) + except Exception: + span.set_status(SPANSTATUS.INTERNAL_ERROR) + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result + + rv.remote = _remote_method_with_header_propagation + + return rv + + if f is not None: + return wrapper(f) + else: + return wrapper + + ray.remote = new_remote + + +def _capture_exception(exc_info: "ExcInfo", **kwargs: "Any") -> None: + client = sentry_sdk.get_client() + + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={ + "handled": False, + "type": RayIntegration.identifier, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + + +class RayIntegration(Integration): + identifier = "ray" + origin = f"auto.queue.{identifier}" + + @staticmethod + def setup_once() -> None: + version = package_version("ray") + _check_minimum_version(RayIntegration, version) + + _patch_ray_remote() diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/rq.py b/lib/python3.12/site-packages/sentry_sdk/integrations/rq.py new file mode 100644 index 0000000000000000000000000000000000000000..8caf46b171d00f4f021edaac9dceebd2f1deac03 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/rq.py @@ -0,0 +1,160 @@ +import weakref + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.api import continue_trace +from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + format_timestamp, + parse_version, +) + +try: + from rq.queue import Queue + from rq.timeouts import JobTimeoutException + from rq.version import VERSION as RQ_VERSION + from rq.worker import Worker + from rq.job import JobStatus +except ImportError: + raise DidNotEnable("RQ not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable + + from sentry_sdk._types import Event, EventProcessor + from sentry_sdk.utils import ExcInfo + + from rq.job import Job + + +class RqIntegration(Integration): + identifier = "rq" + origin = f"auto.queue.{identifier}" + + @staticmethod + def setup_once() -> None: + version = parse_version(RQ_VERSION) + _check_minimum_version(RqIntegration, version) + + old_perform_job = Worker.perform_job + + @ensure_integration_enabled(RqIntegration, old_perform_job) + def sentry_patched_perform_job( + self: "Any", job: "Job", *args: "Queue", **kwargs: "Any" + ) -> bool: + with sentry_sdk.new_scope() as scope: + scope.clear_breadcrumbs() + scope.add_event_processor(_make_event_processor(weakref.ref(job))) + + transaction = continue_trace( + job.meta.get("_sentry_trace_headers") or {}, + op=OP.QUEUE_TASK_RQ, + name="unknown RQ task", + source=TransactionSource.TASK, + origin=RqIntegration.origin, + ) + + with capture_internal_exceptions(): + transaction.name = job.func_name + + with sentry_sdk.start_transaction( + transaction, + custom_sampling_context={"rq_job": job}, + ): + rv = old_perform_job(self, job, *args, **kwargs) + + if self.is_horse: + # We're inside of a forked process and RQ is + # about to call `os._exit`. Make sure that our + # events get sent out. + sentry_sdk.get_client().flush() + + return rv + + Worker.perform_job = sentry_patched_perform_job + + old_handle_exception = Worker.handle_exception + + def sentry_patched_handle_exception( + self: "Worker", job: "Any", *exc_info: "Any", **kwargs: "Any" + ) -> "Any": + retry = ( + hasattr(job, "retries_left") + and job.retries_left + and job.retries_left > 0 + ) + failed = job._status == JobStatus.FAILED or job.is_failed + if failed and not retry: + _capture_exception(exc_info) + + return old_handle_exception(self, job, *exc_info, **kwargs) + + Worker.handle_exception = sentry_patched_handle_exception + + old_enqueue_job = Queue.enqueue_job + + @ensure_integration_enabled(RqIntegration, old_enqueue_job) + def sentry_patched_enqueue_job( + self: "Queue", job: "Any", **kwargs: "Any" + ) -> "Any": + scope = sentry_sdk.get_current_scope() + if scope.span is not None: + job.meta["_sentry_trace_headers"] = dict( + scope.iter_trace_propagation_headers() + ) + + return old_enqueue_job(self, job, **kwargs) + + Queue.enqueue_job = sentry_patched_enqueue_job + + ignore_logger("rq.worker") + + +def _make_event_processor(weak_job: "Callable[[], Job]") -> "EventProcessor": + def event_processor(event: "Event", hint: "dict[str, Any]") -> "Event": + job = weak_job() + if job is not None: + with capture_internal_exceptions(): + extra = event.setdefault("extra", {}) + rq_job = { + "job_id": job.id, + "func": job.func_name, + "args": job.args, + "kwargs": job.kwargs, + "description": job.description, + } + + if job.enqueued_at: + rq_job["enqueued_at"] = format_timestamp(job.enqueued_at) + if job.started_at: + rq_job["started_at"] = format_timestamp(job.started_at) + + extra["rq-job"] = rq_job + + if "exc_info" in hint: + with capture_internal_exceptions(): + if issubclass(hint["exc_info"][0], JobTimeoutException): + event["fingerprint"] = ["rq", "JobTimeoutException", job.func_name] + + return event + + return event_processor + + +def _capture_exception(exc_info: "ExcInfo", **kwargs: "Any") -> None: + client = sentry_sdk.get_client() + + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "rq", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/rust_tracing.py b/lib/python3.12/site-packages/sentry_sdk/integrations/rust_tracing.py new file mode 100644 index 0000000000000000000000000000000000000000..e2b203a2861cb7dffb3e48f75cedd836a2145186 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/rust_tracing.py @@ -0,0 +1,283 @@ +""" +This integration ingests tracing data from native extensions written in Rust. + +Using it requires additional setup on the Rust side to accept a +`RustTracingLayer` Python object and register it with the `tracing-subscriber` +using an adapter from the `pyo3-python-tracing-subscriber` crate. For example: +```rust +#[pyfunction] +pub fn initialize_tracing(py_impl: Bound<'_, PyAny>) { + tracing_subscriber::registry() + .with(pyo3_python_tracing_subscriber::PythonCallbackLayerBridge::new(py_impl)) + .init(); +} +``` + +Usage in Python would then look like: +``` +sentry_sdk.init( + dsn=sentry_dsn, + integrations=[ + RustTracingIntegration( + "demo_rust_extension", + demo_rust_extension.initialize_tracing, + event_type_mapping=event_type_mapping, + ) + ], +) +``` + +Each native extension requires its own integration. +""" + +import json +from enum import Enum, auto +from typing import Any, Callable, Dict, Tuple, Optional + +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import Span as SentrySpan +from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE + +TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]] + + +class RustTracingLevel(Enum): + Trace = "TRACE" + Debug = "DEBUG" + Info = "INFO" + Warn = "WARN" + Error = "ERROR" + + +class EventTypeMapping(Enum): + Ignore = auto() + Exc = auto() + Breadcrumb = auto() + Event = auto() + + +def tracing_level_to_sentry_level(level: str) -> "sentry_sdk._types.LogLevelStr": + level = RustTracingLevel(level) + if level in (RustTracingLevel.Trace, RustTracingLevel.Debug): + return "debug" + elif level == RustTracingLevel.Info: + return "info" + elif level == RustTracingLevel.Warn: + return "warning" + elif level == RustTracingLevel.Error: + return "error" + else: + # Better this than crashing + return "info" + + +def extract_contexts(event: "Dict[str, Any]") -> "Dict[str, Any]": + metadata = event.get("metadata", {}) + contexts = {} + + location = {} + for field in ["module_path", "file", "line"]: + if field in metadata: + location[field] = metadata[field] + if len(location) > 0: + contexts["rust_tracing_location"] = location + + fields = {} + for field in metadata.get("fields", []): + fields[field] = event.get(field) + if len(fields) > 0: + contexts["rust_tracing_fields"] = fields + + return contexts + + +def process_event(event: "Dict[str, Any]") -> None: + metadata = event.get("metadata", {}) + + logger = metadata.get("target") + level = tracing_level_to_sentry_level(metadata.get("level")) + message: "sentry_sdk._types.Any" = event.get("message") + contexts = extract_contexts(event) + + sentry_event: "sentry_sdk._types.Event" = { + "logger": logger, + "level": level, + "message": message, + "contexts": contexts, + } + + sentry_sdk.capture_event(sentry_event) + + +def process_exception(event: "Dict[str, Any]") -> None: + process_event(event) + + +def process_breadcrumb(event: "Dict[str, Any]") -> None: + level = tracing_level_to_sentry_level(event.get("metadata", {}).get("level")) + message = event.get("message") + + sentry_sdk.add_breadcrumb(level=level, message=message) + + +def default_span_filter(metadata: "Dict[str, Any]") -> bool: + return RustTracingLevel(metadata.get("level")) in ( + RustTracingLevel.Error, + RustTracingLevel.Warn, + RustTracingLevel.Info, + ) + + +def default_event_type_mapping(metadata: "Dict[str, Any]") -> "EventTypeMapping": + level = RustTracingLevel(metadata.get("level")) + if level == RustTracingLevel.Error: + return EventTypeMapping.Exc + elif level in (RustTracingLevel.Warn, RustTracingLevel.Info): + return EventTypeMapping.Breadcrumb + elif level in (RustTracingLevel.Debug, RustTracingLevel.Trace): + return EventTypeMapping.Ignore + else: + return EventTypeMapping.Ignore + + +class RustTracingLayer: + def __init__( + self, + origin: str, + event_type_mapping: """Callable[ + [Dict[str, Any]], EventTypeMapping + ]""" = default_event_type_mapping, + span_filter: "Callable[[Dict[str, Any]], bool]" = default_span_filter, + include_tracing_fields: "Optional[bool]" = None, + ): + self.origin = origin + self.event_type_mapping = event_type_mapping + self.span_filter = span_filter + self.include_tracing_fields = include_tracing_fields + + def _include_tracing_fields(self) -> bool: + """ + By default, the values of tracing fields are not included in case they + contain PII. A user may override that by passing `True` for the + `include_tracing_fields` keyword argument of this integration or by + setting `send_default_pii` to `True` in their Sentry client options. + """ + return ( + should_send_default_pii() + if self.include_tracing_fields is None + else self.include_tracing_fields + ) + + def on_event(self, event: str, _span_state: "TraceState") -> None: + deserialized_event = json.loads(event) + metadata = deserialized_event.get("metadata", {}) + + event_type = self.event_type_mapping(metadata) + if event_type == EventTypeMapping.Ignore: + return + elif event_type == EventTypeMapping.Exc: + process_exception(deserialized_event) + elif event_type == EventTypeMapping.Breadcrumb: + process_breadcrumb(deserialized_event) + elif event_type == EventTypeMapping.Event: + process_event(deserialized_event) + + def on_new_span(self, attrs: str, span_id: str) -> "TraceState": + attrs = json.loads(attrs) + metadata = attrs.get("metadata", {}) + + if not self.span_filter(metadata): + return None + + module_path = metadata.get("module_path") + name = metadata.get("name") + message = attrs.get("message") + + if message is not None: + sentry_span_name = message + elif module_path is not None and name is not None: + sentry_span_name = f"{module_path}::{name}" # noqa: E231 + elif name is not None: + sentry_span_name = name + else: + sentry_span_name = "" + + kwargs = { + "op": "function", + "name": sentry_span_name, + "origin": self.origin, + } + + scope = sentry_sdk.get_current_scope() + parent_sentry_span = scope.span + if parent_sentry_span: + sentry_span = parent_sentry_span.start_child(**kwargs) + else: + sentry_span = scope.start_span(**kwargs) + + fields = metadata.get("fields", []) + for field in fields: + if self._include_tracing_fields(): + sentry_span.set_data(field, attrs.get(field)) + else: + sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE) + + scope.span = sentry_span + return (parent_sentry_span, sentry_span) + + def on_close(self, span_id: str, span_state: "TraceState") -> None: + if span_state is None: + return + + parent_sentry_span, sentry_span = span_state + sentry_span.finish() + sentry_sdk.get_current_scope().span = parent_sentry_span + + def on_record(self, span_id: str, values: str, span_state: "TraceState") -> None: + if span_state is None: + return + _parent_sentry_span, sentry_span = span_state + + deserialized_values = json.loads(values) + for key, value in deserialized_values.items(): + if self._include_tracing_fields(): + sentry_span.set_data(key, value) + else: + sentry_span.set_data(key, SENSITIVE_DATA_SUBSTITUTE) + + +class RustTracingIntegration(Integration): + """ + Ingests tracing data from a Rust native extension's `tracing` instrumentation. + + If a project uses more than one Rust native extension, each one will need + its own instance of `RustTracingIntegration` with an initializer function + specific to that extension. + + Since all of the setup for this integration requires instance-specific state + which is not available in `setup_once()`, setup instead happens in `__init__()`. + """ + + def __init__( + self, + identifier: str, + initializer: "Callable[[RustTracingLayer], None]", + event_type_mapping: """Callable[ + [Dict[str, Any]], EventTypeMapping + ]""" = default_event_type_mapping, + span_filter: "Callable[[Dict[str, Any]], bool]" = default_span_filter, + include_tracing_fields: "Optional[bool]" = None, + ): + self.identifier = identifier + origin = f"auto.function.rust_tracing.{identifier}" + self.tracing_layer = RustTracingLayer( + origin, event_type_mapping, span_filter, include_tracing_fields + ) + + initializer(self.tracing_layer) + + @staticmethod + def setup_once() -> None: + pass diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/sanic.py b/lib/python3.12/site-packages/sentry_sdk/integrations/sanic.py new file mode 100644 index 0000000000000000000000000000000000000000..9199b76ebaf27e558fc7dc9000fd8a1fb3c69a58 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/sanic.py @@ -0,0 +1,354 @@ +import sys +import weakref +from inspect import isawaitable +from urllib.parse import urlsplit + +import sentry_sdk +from sentry_sdk import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import RequestExtractor, _filter_headers +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + HAS_REAL_CONTEXTVARS, + CONTEXTVARS_ERROR_MESSAGE, + parse_version, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Container + from typing import Any + from typing import Callable + from typing import Optional + from typing import Union + from typing import Dict + + from sanic.request import Request, RequestParameters + from sanic.response import BaseHTTPResponse + + from sentry_sdk._types import Event, EventProcessor, ExcInfo, Hint + from sanic.router import Route + +try: + from sanic import Sanic, __version__ as SANIC_VERSION + from sanic.exceptions import SanicException + from sanic.router import Router + from sanic.handlers import ErrorHandler +except ImportError: + raise DidNotEnable("Sanic not installed") + +old_error_handler_lookup = ErrorHandler.lookup +old_handle_request = Sanic.handle_request +old_router_get = Router.get + +try: + # This method was introduced in Sanic v21.9 + old_startup = Sanic._startup +except AttributeError: + pass + + +class SanicIntegration(Integration): + identifier = "sanic" + origin = f"auto.http.{identifier}" + version = None + + def __init__( + self, unsampled_statuses: "Optional[Container[int]]" = frozenset({404}) + ) -> None: + """ + The unsampled_statuses parameter can be used to specify for which HTTP statuses the + transactions should not be sent to Sentry. By default, transactions are sent for all + HTTP statuses, except 404. Set unsampled_statuses to None to send transactions for all + HTTP statuses, including 404. + """ + self._unsampled_statuses = unsampled_statuses or set() + + @staticmethod + def setup_once() -> None: + SanicIntegration.version = parse_version(SANIC_VERSION) + _check_minimum_version(SanicIntegration, SanicIntegration.version) + + if not HAS_REAL_CONTEXTVARS: + # We better have contextvars or we're going to leak state between + # requests. + raise DidNotEnable( + "The sanic integration for Sentry requires Python 3.7+ " + " or the aiocontextvars package." + CONTEXTVARS_ERROR_MESSAGE + ) + + if SANIC_VERSION.startswith("0.8."): + # Sanic 0.8 and older creates a logger named "root" and puts a + # stringified version of every exception in there (without exc_info), + # which our error deduplication can't detect. + # + # We explicitly check the version here because it is a very + # invasive step to ignore this logger and not necessary in newer + # versions at all. + # + # https://github.com/huge-success/sanic/issues/1332 + ignore_logger("root") + + if SanicIntegration.version is not None and SanicIntegration.version < (21, 9): + _setup_legacy_sanic() + return + + _setup_sanic() + + +class SanicRequestExtractor(RequestExtractor): + def content_length(self) -> int: + if self.request.body is None: + return 0 + return len(self.request.body) + + def cookies(self) -> "Dict[str, str]": + return dict(self.request.cookies) + + def raw_data(self) -> bytes: + return self.request.body + + def form(self) -> "RequestParameters": + return self.request.form + + def is_json(self) -> bool: + raise NotImplementedError() + + def json(self) -> "Optional[Any]": + return self.request.json + + def files(self) -> "RequestParameters": + return self.request.files + + def size_of_file(self, file: "Any") -> int: + return len(file.body or ()) + + +def _setup_sanic() -> None: + Sanic._startup = _startup + ErrorHandler.lookup = _sentry_error_handler_lookup + + +def _setup_legacy_sanic() -> None: + Sanic.handle_request = _legacy_handle_request + Router.get = _legacy_router_get + ErrorHandler.lookup = _sentry_error_handler_lookup + + +async def _startup(self: "Sanic") -> None: + # This happens about as early in the lifecycle as possible, just after the + # Request object is created. The body has not yet been consumed. + self.signal("http.lifecycle.request")(_context_enter) + + # This happens after the handler is complete. In v21.9 this signal is not + # dispatched when there is an exception. Therefore we need to close out + # and call _context_exit from the custom exception handler as well. + # See https://github.com/sanic-org/sanic/issues/2297 + self.signal("http.lifecycle.response")(_context_exit) + + # This happens inside of request handling immediately after the route + # has been identified by the router. + self.signal("http.routing.after")(_set_transaction) + + # The above signals need to be declared before this can be called. + await old_startup(self) + + +async def _context_enter(request: "Request") -> None: + request.ctx._sentry_do_integration = ( + sentry_sdk.get_client().get_integration(SanicIntegration) is not None + ) + + if not request.ctx._sentry_do_integration: + return + + weak_request = weakref.ref(request) + request.ctx._sentry_scope = sentry_sdk.isolation_scope() + scope = request.ctx._sentry_scope.__enter__() + scope.clear_breadcrumbs() + scope.add_event_processor(_make_request_processor(weak_request)) + + transaction = continue_trace( + dict(request.headers), + op=OP.HTTP_SERVER, + # Unless the request results in a 404 error, the name and source will get overwritten in _set_transaction + name=request.path, + source=TransactionSource.URL, + origin=SanicIntegration.origin, + ) + request.ctx._sentry_transaction = sentry_sdk.start_transaction( + transaction + ).__enter__() + + +async def _context_exit( + request: "Request", response: "Optional[BaseHTTPResponse]" = None +) -> None: + with capture_internal_exceptions(): + if not request.ctx._sentry_do_integration: + return + + integration = sentry_sdk.get_client().get_integration(SanicIntegration) + + response_status = None if response is None else response.status + + # This capture_internal_exceptions block has been intentionally nested here, so that in case an exception + # happens while trying to end the transaction, we still attempt to exit the hub. + with capture_internal_exceptions(): + request.ctx._sentry_transaction.set_http_status(response_status) + request.ctx._sentry_transaction.sampled &= ( + isinstance(integration, SanicIntegration) + and response_status not in integration._unsampled_statuses + ) + request.ctx._sentry_transaction.__exit__(None, None, None) + + request.ctx._sentry_scope.__exit__(None, None, None) + + +async def _set_transaction(request: "Request", route: "Route", **_: "Any") -> None: + if request.ctx._sentry_do_integration: + with capture_internal_exceptions(): + scope = sentry_sdk.get_current_scope() + route_name = route.name.replace(request.app.name, "").strip(".") + scope.set_transaction_name(route_name, source=TransactionSource.COMPONENT) + + +def _sentry_error_handler_lookup( + self: "Any", exception: Exception, *args: "Any", **kwargs: "Any" +) -> "Optional[object]": + _capture_exception(exception) + old_error_handler = old_error_handler_lookup(self, exception, *args, **kwargs) + + if old_error_handler is None: + return None + + if sentry_sdk.get_client().get_integration(SanicIntegration) is None: + return old_error_handler + + async def sentry_wrapped_error_handler( + request: "Request", exception: Exception + ) -> "Any": + try: + response = old_error_handler(request, exception) + if isawaitable(response): + response = await response + return response + except Exception: + # Report errors that occur in Sanic error handler. These + # exceptions will not even show up in Sanic's + # `sanic.exceptions` logger. + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + finally: + # As mentioned in previous comment in _startup, this can be removed + # after https://github.com/sanic-org/sanic/issues/2297 is resolved + if SanicIntegration.version and SanicIntegration.version == (21, 9): + await _context_exit(request) + + return sentry_wrapped_error_handler + + +async def _legacy_handle_request( + self: "Any", request: "Request", *args: "Any", **kwargs: "Any" +) -> "Any": + if sentry_sdk.get_client().get_integration(SanicIntegration) is None: + return await old_handle_request(self, request, *args, **kwargs) + + weak_request = weakref.ref(request) + + with sentry_sdk.isolation_scope() as scope: + scope.clear_breadcrumbs() + scope.add_event_processor(_make_request_processor(weak_request)) + + response = old_handle_request(self, request, *args, **kwargs) + if isawaitable(response): + response = await response + + return response + + +def _legacy_router_get(self: "Any", *args: "Union[Any, Request]") -> "Any": + rv = old_router_get(self, *args) + if sentry_sdk.get_client().get_integration(SanicIntegration) is not None: + with capture_internal_exceptions(): + scope = sentry_sdk.get_isolation_scope() + if SanicIntegration.version and SanicIntegration.version >= (21, 3): + # Sanic versions above and including 21.3 append the app name to the + # route name, and so we need to remove it from Route name so the + # transaction name is consistent across all versions + sanic_app_name = self.ctx.app.name + sanic_route = rv[0].name + + if sanic_route.startswith("%s." % sanic_app_name): + # We add a 1 to the len of the sanic_app_name because there is a dot + # that joins app name and the route name + # Format: app_name.route_name + sanic_route = sanic_route[len(sanic_app_name) + 1 :] + + scope.set_transaction_name( + sanic_route, source=TransactionSource.COMPONENT + ) + else: + scope.set_transaction_name( + rv[0].__name__, source=TransactionSource.COMPONENT + ) + + return rv + + +@ensure_integration_enabled(SanicIntegration) +def _capture_exception(exception: "Union[ExcInfo, BaseException]") -> None: + with capture_internal_exceptions(): + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "sanic", "handled": False}, + ) + + if hint and hasattr(hint["exc_info"][0], "quiet") and hint["exc_info"][0].quiet: + return + + sentry_sdk.capture_event(event, hint=hint) + + +def _make_request_processor(weak_request: "Callable[[], Request]") -> "EventProcessor": + def sanic_processor(event: "Event", hint: "Optional[Hint]") -> "Optional[Event]": + try: + if hint and issubclass(hint["exc_info"][0], SanicException): + return None + except KeyError: + pass + + request = weak_request() + if request is None: + return event + + with capture_internal_exceptions(): + extractor = SanicRequestExtractor(request) + extractor.extract_into_event(event) + + request_info = event["request"] + urlparts = urlsplit(request.url) + + request_info["url"] = "%s://%s%s" % ( + urlparts.scheme, + urlparts.netloc, + urlparts.path, + ) + + request_info["query_string"] = urlparts.query + request_info["method"] = request.method + request_info["env"] = {"REMOTE_ADDR": request.remote_addr} + request_info["headers"] = _filter_headers(dict(request.headers)) + + return event + + return sanic_processor diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/serverless.py b/lib/python3.12/site-packages/sentry_sdk/integrations/serverless.py new file mode 100644 index 0000000000000000000000000000000000000000..16f91b28aeb9206f950eece356b0e24fec071f41 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/serverless.py @@ -0,0 +1,65 @@ +import sys +from functools import wraps +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.utils import event_from_exception, reraise + +if TYPE_CHECKING: + from typing import Any, Callable, Optional, TypeVar, Union, overload + + F = TypeVar("F", bound=Callable[..., Any]) + +else: + + def overload(x: "F") -> "F": + return x + + +@overload +def serverless_function(f: "F", flush: bool = True) -> "F": + pass + + +@overload +def serverless_function(f: None = None, flush: bool = True) -> "Callable[[F], F]": # noqa: F811 + pass + + +def serverless_function( # noqa + f: "Optional[F]" = None, flush: bool = True +) -> "Union[F, Callable[[F], F]]": + def wrapper(f: "F") -> "F": + @wraps(f) + def inner(*args: "Any", **kwargs: "Any") -> "Any": + with sentry_sdk.isolation_scope() as scope: + scope.clear_breadcrumbs() + + try: + return f(*args, **kwargs) + except Exception: + _capture_and_reraise() + finally: + if flush: + sentry_sdk.flush() + + return inner # type: ignore + + if f is None: + return wrapper + else: + return wrapper(f) + + +def _capture_and_reraise() -> None: + exc_info = sys.exc_info() + client = sentry_sdk.get_client() + if client.is_active(): + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "serverless", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + reraise(*exc_info) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/socket.py b/lib/python3.12/site-packages/sentry_sdk/integrations/socket.py new file mode 100644 index 0000000000000000000000000000000000000000..472b909d28590636de51827e49f6b9d656bb493b --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/socket.py @@ -0,0 +1,98 @@ +import socket + +import sentry_sdk +from sentry_sdk._types import MYPY +from sentry_sdk.consts import OP +from sentry_sdk.integrations import Integration + +if MYPY: + from socket import AddressFamily, SocketKind + from typing import Tuple, Optional, Union, List + +__all__ = ["SocketIntegration"] + + +class SocketIntegration(Integration): + identifier = "socket" + origin = f"auto.socket.{identifier}" + + @staticmethod + def setup_once() -> None: + """ + patches two of the most used functions of socket: create_connection and getaddrinfo(dns resolver) + """ + _patch_create_connection() + _patch_getaddrinfo() + + +def _get_span_description( + host: "Union[bytes, str, None]", port: "Union[bytes, str, int, None]" +) -> str: + try: + host = host.decode() # type: ignore + except (UnicodeDecodeError, AttributeError): + pass + + try: + port = port.decode() # type: ignore + except (UnicodeDecodeError, AttributeError): + pass + + description = "%s:%s" % (host, port) # type: ignore + return description + + +def _patch_create_connection() -> None: + real_create_connection = socket.create_connection + + def create_connection( + address: "Tuple[Optional[str], int]", + timeout: "Optional[float]" = socket._GLOBAL_DEFAULT_TIMEOUT, # type: ignore + source_address: "Optional[Tuple[Union[bytearray, bytes, str], int]]" = None, + ) -> "socket.socket": + integration = sentry_sdk.get_client().get_integration(SocketIntegration) + if integration is None: + return real_create_connection(address, timeout, source_address) + + with sentry_sdk.start_span( + op=OP.SOCKET_CONNECTION, + name=_get_span_description(address[0], address[1]), + origin=SocketIntegration.origin, + ) as span: + span.set_data("address", address) + span.set_data("timeout", timeout) + span.set_data("source_address", source_address) + + return real_create_connection( + address=address, timeout=timeout, source_address=source_address + ) + + socket.create_connection = create_connection # type: ignore + + +def _patch_getaddrinfo() -> None: + real_getaddrinfo = socket.getaddrinfo + + def getaddrinfo( + host: "Union[bytes, str, None]", + port: "Union[bytes, str, int, None]", + family: int = 0, + type: int = 0, + proto: int = 0, + flags: int = 0, + ) -> "List[Tuple[AddressFamily, SocketKind, int, str, Union[Tuple[str, int], Tuple[str, int, int, int], Tuple[int, bytes]]]]": + integration = sentry_sdk.get_client().get_integration(SocketIntegration) + if integration is None: + return real_getaddrinfo(host, port, family, type, proto, flags) + + with sentry_sdk.start_span( + op=OP.SOCKET_DNS, + name=_get_span_description(host, port), + origin=SocketIntegration.origin, + ) as span: + span.set_data("host", host) + span.set_data("port", port) + + return real_getaddrinfo(host, port, family, type, proto, flags) + + socket.getaddrinfo = getaddrinfo diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/sqlalchemy.py b/lib/python3.12/site-packages/sentry_sdk/integrations/sqlalchemy.py new file mode 100644 index 0000000000000000000000000000000000000000..7d3ed9537388e925228d478d08e8d275553248b7 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/sqlalchemy.py @@ -0,0 +1,153 @@ +from sentry_sdk.consts import SPANSTATUS, SPANDATA +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.tracing_utils import add_query_source, record_sql_queries +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + parse_version, +) + +try: + from sqlalchemy.engine import Engine # type: ignore + from sqlalchemy.event import listen # type: ignore + from sqlalchemy import __version__ as SQLALCHEMY_VERSION # type: ignore +except ImportError: + raise DidNotEnable("SQLAlchemy not installed.") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import ContextManager + from typing import Optional + + from sentry_sdk.tracing import Span + + +class SqlalchemyIntegration(Integration): + identifier = "sqlalchemy" + origin = f"auto.db.{identifier}" + + @staticmethod + def setup_once() -> None: + version = parse_version(SQLALCHEMY_VERSION) + _check_minimum_version(SqlalchemyIntegration, version) + + listen(Engine, "before_cursor_execute", _before_cursor_execute) + listen(Engine, "after_cursor_execute", _after_cursor_execute) + listen(Engine, "handle_error", _handle_error) + + +@ensure_integration_enabled(SqlalchemyIntegration) +def _before_cursor_execute( + conn: "Any", + cursor: "Any", + statement: "Any", + parameters: "Any", + context: "Any", + executemany: bool, + *args: "Any", +) -> None: + ctx_mgr = record_sql_queries( + cursor, + statement, + parameters, + paramstyle=context and context.dialect and context.dialect.paramstyle or None, + executemany=executemany, + span_origin=SqlalchemyIntegration.origin, + ) + context._sentry_sql_span_manager = ctx_mgr + + span = ctx_mgr.__enter__() + + if span is not None: + _set_db_data(span, conn) + context._sentry_sql_span = span + + +@ensure_integration_enabled(SqlalchemyIntegration) +def _after_cursor_execute( + conn: "Any", + cursor: "Any", + statement: "Any", + parameters: "Any", + context: "Any", + *args: "Any", +) -> None: + ctx_mgr: "Optional[ContextManager[Any]]" = getattr( + context, "_sentry_sql_span_manager", None + ) + + if ctx_mgr is not None: + context._sentry_sql_span_manager = None + ctx_mgr.__exit__(None, None, None) + + span: "Optional[Span]" = getattr(context, "_sentry_sql_span", None) + if span is not None: + with capture_internal_exceptions(): + add_query_source(span) + + +def _handle_error(context: "Any", *args: "Any") -> None: + execution_context = context.execution_context + if execution_context is None: + return + + span: "Optional[Span]" = getattr(execution_context, "_sentry_sql_span", None) + + if span is not None: + span.set_status(SPANSTATUS.INTERNAL_ERROR) + + # _after_cursor_execute does not get called for crashing SQL stmts. Judging + # from SQLAlchemy codebase it does seem like any error coming into this + # handler is going to be fatal. + ctx_mgr: "Optional[ContextManager[Any]]" = getattr( + execution_context, "_sentry_sql_span_manager", None + ) + + if ctx_mgr is not None: + execution_context._sentry_sql_span_manager = None + ctx_mgr.__exit__(None, None, None) + + +# See: https://docs.sqlalchemy.org/en/20/dialects/index.html +def _get_db_system(name: str) -> "Optional[str]": + name = str(name) + + if "sqlite" in name: + return "sqlite" + + if "postgres" in name: + return "postgresql" + + if "mariadb" in name: + return "mariadb" + + if "mysql" in name: + return "mysql" + + if "oracle" in name: + return "oracle" + + return None + + +def _set_db_data(span: "Span", conn: "Any") -> None: + db_system = _get_db_system(conn.engine.name) + if db_system is not None: + span.set_data(SPANDATA.DB_SYSTEM, db_system) + + if conn.engine.url is None: + return + + db_name = conn.engine.url.database + if db_name is not None: + span.set_data(SPANDATA.DB_NAME, db_name) + + server_address = conn.engine.url.host + if server_address is not None: + span.set_data(SPANDATA.SERVER_ADDRESS, server_address) + + server_port = conn.engine.url.port + if server_port is not None: + span.set_data(SPANDATA.SERVER_PORT, server_port) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/starlette.py b/lib/python3.12/site-packages/sentry_sdk/integrations/starlette.py new file mode 100644 index 0000000000000000000000000000000000000000..0b797ebcde9c6ad0bafeb48e11961a5173e5410b --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/starlette.py @@ -0,0 +1,739 @@ +import asyncio +import functools +import warnings +from collections.abc import Set +from copy import deepcopy +from json import JSONDecodeError + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import ( + DidNotEnable, + Integration, + _DEFAULT_FAILED_REQUEST_STATUS_CODES, +) +from sentry_sdk.integrations._wsgi_common import ( + DEFAULT_HTTP_METHODS_TO_CAPTURE, + HttpCodeRangeContainer, + _is_json_content_type, + request_body_within_bounds, +) +from sentry_sdk.integrations.asgi import SentryAsgiMiddleware +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import ( + SOURCE_FOR_STYLE, + TransactionSource, +) +from sentry_sdk.utils import ( + AnnotatedValue, + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + parse_version, + transaction_from_function, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Awaitable, Callable, Container, Dict, Optional, Tuple, Union + + from sentry_sdk._types import Event, HttpStatusCodeRange + +try: + import starlette # type: ignore + from starlette import __version__ as STARLETTE_VERSION + from starlette.applications import Starlette # type: ignore + from starlette.datastructures import UploadFile # type: ignore + from starlette.middleware import Middleware # type: ignore + from starlette.middleware.authentication import ( # type: ignore + AuthenticationMiddleware, + ) + from starlette.requests import Request # type: ignore + from starlette.routing import Match # type: ignore + from starlette.types import ASGIApp, Receive, Scope as StarletteScope, Send # type: ignore +except ImportError: + raise DidNotEnable("Starlette is not installed") + +try: + # Starlette 0.20 + from starlette.middleware.exceptions import ExceptionMiddleware # type: ignore +except ImportError: + # Startlette 0.19.1 + from starlette.exceptions import ExceptionMiddleware # type: ignore + +try: + # Optional dependency of Starlette to parse form data. + try: + # python-multipart 0.0.13 and later + import python_multipart as multipart # type: ignore + except ImportError: + # python-multipart 0.0.12 and earlier + import multipart # type: ignore +except ImportError: + multipart = None + + +_DEFAULT_TRANSACTION_NAME = "generic Starlette request" + +TRANSACTION_STYLE_VALUES = ("endpoint", "url") + + +class StarletteIntegration(Integration): + identifier = "starlette" + origin = f"auto.http.{identifier}" + + transaction_style = "" + + def __init__( + self, + transaction_style: str = "url", + failed_request_status_codes: "Union[Set[int], list[HttpStatusCodeRange], None]" = _DEFAULT_FAILED_REQUEST_STATUS_CODES, + middleware_spans: bool = False, + http_methods_to_capture: "tuple[str, ...]" = DEFAULT_HTTP_METHODS_TO_CAPTURE, + ): + if transaction_style not in TRANSACTION_STYLE_VALUES: + raise ValueError( + "Invalid value for transaction_style: %s (must be in %s)" + % (transaction_style, TRANSACTION_STYLE_VALUES) + ) + self.transaction_style = transaction_style + self.middleware_spans = middleware_spans + self.http_methods_to_capture = tuple(map(str.upper, http_methods_to_capture)) + + if isinstance(failed_request_status_codes, Set): + self.failed_request_status_codes: "Container[int]" = ( + failed_request_status_codes + ) + else: + warnings.warn( + "Passing a list or None for failed_request_status_codes is deprecated. " + "Please pass a set of int instead.", + DeprecationWarning, + stacklevel=2, + ) + + if failed_request_status_codes is None: + self.failed_request_status_codes = _DEFAULT_FAILED_REQUEST_STATUS_CODES + else: + self.failed_request_status_codes = HttpCodeRangeContainer( + failed_request_status_codes + ) + + @staticmethod + def setup_once() -> None: + version = parse_version(STARLETTE_VERSION) + + if version is None: + raise DidNotEnable( + "Unparsable Starlette version: {}".format(STARLETTE_VERSION) + ) + + patch_middlewares() + patch_asgi_app() + patch_request_response() + + if version >= (0, 24): + patch_templates() + + +def _enable_span_for_middleware(middleware_class: "Any") -> type: + old_call = middleware_class.__call__ + + async def _create_span_call( + app: "Any", + scope: "Dict[str, Any]", + receive: "Callable[[], Awaitable[Dict[str, Any]]]", + send: "Callable[[Dict[str, Any]], Awaitable[None]]", + **kwargs: "Any", + ) -> None: + integration = sentry_sdk.get_client().get_integration(StarletteIntegration) + if integration is None: + return await old_call(app, scope, receive, send, **kwargs) + + # Update transaction name with middleware name + name, source = _get_transaction_from_middleware(app, scope, integration) + + if name is not None: + sentry_sdk.get_current_scope().set_transaction_name( + name, + source=source, + ) + + if not integration.middleware_spans: + return await old_call(app, scope, receive, send, **kwargs) + + middleware_name = app.__class__.__name__ + + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLETTE, + name=middleware_name, + origin=StarletteIntegration.origin, + ) as middleware_span: + middleware_span.set_tag("starlette.middleware_name", middleware_name) + + # Creating spans for the "receive" callback + async def _sentry_receive(*args: "Any", **kwargs: "Any") -> "Any": + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLETTE_RECEIVE, + name=getattr(receive, "__qualname__", str(receive)), + origin=StarletteIntegration.origin, + ) as span: + span.set_tag("starlette.middleware_name", middleware_name) + return await receive(*args, **kwargs) + + receive_name = getattr(receive, "__name__", str(receive)) + receive_patched = receive_name == "_sentry_receive" + new_receive = _sentry_receive if not receive_patched else receive + + # Creating spans for the "send" callback + async def _sentry_send(*args: "Any", **kwargs: "Any") -> "Any": + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLETTE_SEND, + name=getattr(send, "__qualname__", str(send)), + origin=StarletteIntegration.origin, + ) as span: + span.set_tag("starlette.middleware_name", middleware_name) + return await send(*args, **kwargs) + + send_name = getattr(send, "__name__", str(send)) + send_patched = send_name == "_sentry_send" + new_send = _sentry_send if not send_patched else send + + return await old_call(app, scope, new_receive, new_send, **kwargs) + + not_yet_patched = old_call.__name__ not in [ + "_create_span_call", + "_sentry_authenticationmiddleware_call", + "_sentry_exceptionmiddleware_call", + ] + + if not_yet_patched: + middleware_class.__call__ = _create_span_call + + return middleware_class + + +@ensure_integration_enabled(StarletteIntegration) +def _capture_exception(exception: BaseException, handled: "Any" = False) -> None: + event, hint = event_from_exception( + exception, + client_options=sentry_sdk.get_client().options, + mechanism={"type": StarletteIntegration.identifier, "handled": handled}, + ) + + sentry_sdk.capture_event(event, hint=hint) + + +def patch_exception_middleware(middleware_class: "Any") -> None: + """ + Capture all exceptions in Starlette app and + also extract user information. + """ + old_middleware_init = middleware_class.__init__ + + not_yet_patched = "_sentry_middleware_init" not in str(old_middleware_init) + + if not_yet_patched: + + def _sentry_middleware_init(self: "Any", *args: "Any", **kwargs: "Any") -> None: + old_middleware_init(self, *args, **kwargs) + + # Patch existing exception handlers + old_handlers = self._exception_handlers.copy() + + async def _sentry_patched_exception_handler( + self: "Any", *args: "Any", **kwargs: "Any" + ) -> None: + integration = sentry_sdk.get_client().get_integration( + StarletteIntegration + ) + + exp = args[0] + + if integration is not None: + is_http_server_error = ( + hasattr(exp, "status_code") + and isinstance(exp.status_code, int) + and exp.status_code in integration.failed_request_status_codes + ) + if is_http_server_error: + _capture_exception(exp, handled=True) + + # Find a matching handler + old_handler = None + for cls in type(exp).__mro__: + if cls in old_handlers: + old_handler = old_handlers[cls] + break + + if old_handler is None: + return + + if _is_async_callable(old_handler): + return await old_handler(self, *args, **kwargs) + else: + return old_handler(self, *args, **kwargs) + + for key in self._exception_handlers.keys(): + self._exception_handlers[key] = _sentry_patched_exception_handler + + middleware_class.__init__ = _sentry_middleware_init + + old_call = middleware_class.__call__ + + async def _sentry_exceptionmiddleware_call( + self: "Dict[str, Any]", + scope: "Dict[str, Any]", + receive: "Callable[[], Awaitable[Dict[str, Any]]]", + send: "Callable[[Dict[str, Any]], Awaitable[None]]", + ) -> None: + # Also add the user (that was eventually set by be Authentication middle + # that was called before this middleware). This is done because the authentication + # middleware sets the user in the scope and then (in the same function) + # calls this exception middelware. In case there is no exception (or no handler + # for the type of exception occuring) then the exception bubbles up and setting the + # user information into the sentry scope is done in auth middleware and the + # ASGI middleware will then send everything to Sentry and this is fine. + # But if there is an exception happening that the exception middleware here + # has a handler for, it will send the exception directly to Sentry, so we need + # the user information right now. + # This is why we do it here. + _add_user_to_sentry_scope(scope) + await old_call(self, scope, receive, send) + + middleware_class.__call__ = _sentry_exceptionmiddleware_call + + +@ensure_integration_enabled(StarletteIntegration) +def _add_user_to_sentry_scope(scope: "Dict[str, Any]") -> None: + """ + Extracts user information from the ASGI scope and + adds it to Sentry's scope. + """ + if "user" not in scope: + return + + if not should_send_default_pii(): + return + + user_info: "Dict[str, Any]" = {} + starlette_user = scope["user"] + + username = getattr(starlette_user, "username", None) + if username: + user_info.setdefault("username", starlette_user.username) + + user_id = getattr(starlette_user, "id", None) + if user_id: + user_info.setdefault("id", starlette_user.id) + + email = getattr(starlette_user, "email", None) + if email: + user_info.setdefault("email", starlette_user.email) + + sentry_scope = sentry_sdk.get_isolation_scope() + sentry_scope.set_user(user_info) + + +def patch_authentication_middleware(middleware_class: "Any") -> None: + """ + Add user information to Sentry scope. + """ + old_call = middleware_class.__call__ + + not_yet_patched = "_sentry_authenticationmiddleware_call" not in str(old_call) + + if not_yet_patched: + + async def _sentry_authenticationmiddleware_call( + self: "Dict[str, Any]", + scope: "Dict[str, Any]", + receive: "Callable[[], Awaitable[Dict[str, Any]]]", + send: "Callable[[Dict[str, Any]], Awaitable[None]]", + ) -> None: + await old_call(self, scope, receive, send) + _add_user_to_sentry_scope(scope) + + middleware_class.__call__ = _sentry_authenticationmiddleware_call + + +def patch_middlewares() -> None: + """ + Patches Starlettes `Middleware` class to record + spans for every middleware invoked. + """ + old_middleware_init = Middleware.__init__ + + not_yet_patched = "_sentry_middleware_init" not in str(old_middleware_init) + + if not_yet_patched: + + def _sentry_middleware_init( + self: "Any", cls: "Any", *args: "Any", **kwargs: "Any" + ) -> None: + if cls == SentryAsgiMiddleware: + return old_middleware_init(self, cls, *args, **kwargs) + + span_enabled_cls = _enable_span_for_middleware(cls) + old_middleware_init(self, span_enabled_cls, *args, **kwargs) + + if cls == AuthenticationMiddleware: + patch_authentication_middleware(cls) + + if cls == ExceptionMiddleware: + patch_exception_middleware(cls) + + Middleware.__init__ = _sentry_middleware_init + + +def patch_asgi_app() -> None: + """ + Instrument Starlette ASGI app using the SentryAsgiMiddleware. + """ + old_app = Starlette.__call__ + + async def _sentry_patched_asgi_app( + self: "Starlette", scope: "StarletteScope", receive: "Receive", send: "Send" + ) -> None: + integration = sentry_sdk.get_client().get_integration(StarletteIntegration) + if integration is None: + return await old_app(self, scope, receive, send) + + middleware = SentryAsgiMiddleware( + lambda *a, **kw: old_app(self, *a, **kw), + mechanism_type=StarletteIntegration.identifier, + transaction_style=integration.transaction_style, + span_origin=StarletteIntegration.origin, + http_methods_to_capture=( + integration.http_methods_to_capture + if integration + else DEFAULT_HTTP_METHODS_TO_CAPTURE + ), + asgi_version=3, + ) + + return await middleware(scope, receive, send) + + Starlette.__call__ = _sentry_patched_asgi_app + + +# This was vendored in from Starlette to support Starlette 0.19.1 because +# this function was only introduced in 0.20.x +def _is_async_callable(obj: "Any") -> bool: + while isinstance(obj, functools.partial): + obj = obj.func + + return asyncio.iscoroutinefunction(obj) or ( + callable(obj) and asyncio.iscoroutinefunction(obj.__call__) + ) + + +def patch_request_response() -> None: + old_request_response = starlette.routing.request_response + + def _sentry_request_response(func: "Callable[[Any], Any]") -> "ASGIApp": + old_func = func + + is_coroutine = _is_async_callable(old_func) + if is_coroutine: + + async def _sentry_async_func(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration( + StarletteIntegration + ) + if integration is None: + return await old_func(*args, **kwargs) + + request = args[0] + + _set_transaction_name_and_source( + sentry_sdk.get_current_scope(), + integration.transaction_style, + request, + ) + + sentry_scope = sentry_sdk.get_isolation_scope() + extractor = StarletteRequestExtractor(request) + info = await extractor.extract_request_info() + + def _make_request_event_processor( + req: "Any", integration: "Any" + ) -> "Callable[[Event, dict[str, Any]], Event]": + def event_processor( + event: "Event", hint: "Dict[str, Any]" + ) -> "Event": + # Add info from request to event + request_info = event.get("request", {}) + if info: + if "cookies" in info: + request_info["cookies"] = info["cookies"] + if "data" in info: + request_info["data"] = info["data"] + event["request"] = deepcopy(request_info) + + return event + + return event_processor + + sentry_scope._name = StarletteIntegration.identifier + sentry_scope.add_event_processor( + _make_request_event_processor(request, integration) + ) + + return await old_func(*args, **kwargs) + + func = _sentry_async_func + + else: + + @functools.wraps(old_func) + def _sentry_sync_func(*args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration( + StarletteIntegration + ) + if integration is None: + return old_func(*args, **kwargs) + + current_scope = sentry_sdk.get_current_scope() + if current_scope.transaction is not None: + current_scope.transaction.update_active_thread() + + sentry_scope = sentry_sdk.get_isolation_scope() + if sentry_scope.profile is not None: + sentry_scope.profile.update_active_thread_id() + + request = args[0] + + _set_transaction_name_and_source( + sentry_scope, integration.transaction_style, request + ) + + extractor = StarletteRequestExtractor(request) + cookies = extractor.extract_cookies_from_request() + + def _make_request_event_processor( + req: "Any", integration: "Any" + ) -> "Callable[[Event, dict[str, Any]], Event]": + def event_processor( + event: "Event", hint: "dict[str, Any]" + ) -> "Event": + # Extract information from request + request_info = event.get("request", {}) + if cookies: + request_info["cookies"] = cookies + + event["request"] = deepcopy(request_info) + + return event + + return event_processor + + sentry_scope._name = StarletteIntegration.identifier + sentry_scope.add_event_processor( + _make_request_event_processor(request, integration) + ) + + return old_func(*args, **kwargs) + + func = _sentry_sync_func + + return old_request_response(func) + + starlette.routing.request_response = _sentry_request_response + + +def patch_templates() -> None: + # If markupsafe is not installed, then Jinja2 is not installed + # (markupsafe is a dependency of Jinja2) + # In this case we do not need to patch the Jinja2Templates class + try: + from markupsafe import Markup + except ImportError: + return # Nothing to do + + from starlette.templating import Jinja2Templates # type: ignore + + old_jinja2templates_init = Jinja2Templates.__init__ + + not_yet_patched = "_sentry_jinja2templates_init" not in str( + old_jinja2templates_init + ) + + if not_yet_patched: + + def _sentry_jinja2templates_init( + self: "Jinja2Templates", *args: "Any", **kwargs: "Any" + ) -> None: + def add_sentry_trace_meta(request: "Request") -> "Dict[str, Any]": + trace_meta = Markup( + sentry_sdk.get_current_scope().trace_propagation_meta() + ) + return { + "sentry_trace_meta": trace_meta, + } + + kwargs.setdefault("context_processors", []) + + if add_sentry_trace_meta not in kwargs["context_processors"]: + kwargs["context_processors"].append(add_sentry_trace_meta) + + return old_jinja2templates_init(self, *args, **kwargs) + + Jinja2Templates.__init__ = _sentry_jinja2templates_init + + +class StarletteRequestExtractor: + """ + Extracts useful information from the Starlette request + (like form data or cookies) and adds it to the Sentry event. + """ + + request: "Request" = None + + def __init__(self: "StarletteRequestExtractor", request: "Request") -> None: + self.request = request + + def extract_cookies_from_request( + self: "StarletteRequestExtractor", + ) -> "Optional[Dict[str, Any]]": + cookies: "Optional[Dict[str, Any]]" = None + if should_send_default_pii(): + cookies = self.cookies() + + return cookies + + async def extract_request_info( + self: "StarletteRequestExtractor", + ) -> "Optional[Dict[str, Any]]": + client = sentry_sdk.get_client() + + request_info: "Dict[str, Any]" = {} + + with capture_internal_exceptions(): + # Add cookies + if should_send_default_pii(): + request_info["cookies"] = self.cookies() + + # If there is no body, just return the cookies + content_length = await self.content_length() + if not content_length: + return request_info + + # Add annotation if body is too big + if content_length and not request_body_within_bounds( + client, content_length + ): + request_info["data"] = AnnotatedValue.removed_because_over_size_limit() + return request_info + + # Add JSON body, if it is a JSON request + json = await self.json() + if json: + request_info["data"] = json + return request_info + + # Add form as key/value pairs, if request has form data + form = await self.form() + if form: + form_data = {} + for key, val in form.items(): + is_file = isinstance(val, UploadFile) + form_data[key] = ( + val + if not is_file + else AnnotatedValue.removed_because_raw_data() + ) + + request_info["data"] = form_data + return request_info + + # Raw data, do not add body just an annotation + request_info["data"] = AnnotatedValue.removed_because_raw_data() + return request_info + + async def content_length(self: "StarletteRequestExtractor") -> "Optional[int]": + if "content-length" in self.request.headers: + return int(self.request.headers["content-length"]) + + return None + + def cookies(self: "StarletteRequestExtractor") -> "Dict[str, Any]": + return self.request.cookies + + async def form(self: "StarletteRequestExtractor") -> "Any": + if multipart is None: + return None + + # Parse the body first to get it cached, as Starlette does not cache form() as it + # does with body() and json() https://github.com/encode/starlette/discussions/1933 + # Calling `.form()` without calling `.body()` first will + # potentially break the users project. + await self.request.body() + + return await self.request.form() + + def is_json(self: "StarletteRequestExtractor") -> bool: + return _is_json_content_type(self.request.headers.get("content-type")) + + async def json(self: "StarletteRequestExtractor") -> "Optional[Dict[str, Any]]": + if not self.is_json(): + return None + try: + return await self.request.json() + except JSONDecodeError: + return None + + +def _transaction_name_from_router(scope: "StarletteScope") -> "Optional[str]": + router = scope.get("router") + if not router: + return None + + for route in router.routes: + match = route.matches(scope) + if match[0] == Match.FULL: + try: + return route.path + except AttributeError: + # routes added via app.host() won't have a path attribute + return scope.get("path") + + return None + + +def _set_transaction_name_and_source( + scope: "sentry_sdk.Scope", transaction_style: str, request: "Any" +) -> None: + name = None + source = SOURCE_FOR_STYLE[transaction_style] + + if transaction_style == "endpoint": + endpoint = request.scope.get("endpoint") + if endpoint: + name = transaction_from_function(endpoint) or None + + elif transaction_style == "url": + name = _transaction_name_from_router(request.scope) + + if name is None: + name = _DEFAULT_TRANSACTION_NAME + source = TransactionSource.ROUTE + + scope.set_transaction_name(name, source=source) + + +def _get_transaction_from_middleware( + app: "Any", asgi_scope: "Dict[str, Any]", integration: "StarletteIntegration" +) -> "Tuple[Optional[str], Optional[str]]": + name = None + source = None + + if integration.transaction_style == "endpoint": + name = transaction_from_function(app.__class__) + source = TransactionSource.COMPONENT + elif integration.transaction_style == "url": + name = _transaction_name_from_router(asgi_scope) + source = TransactionSource.ROUTE + + return name, source diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/starlite.py b/lib/python3.12/site-packages/sentry_sdk/integrations/starlite.py new file mode 100644 index 0000000000000000000000000000000000000000..af66d37faef223abd63e06c5f1992b8790a6264b --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/starlite.py @@ -0,0 +1,290 @@ +from copy import deepcopy + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.integrations.asgi import SentryAsgiMiddleware +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource +from sentry_sdk.utils import ( + ensure_integration_enabled, + event_from_exception, + transaction_from_function, +) + +try: + from starlite import Request, Starlite, State # type: ignore + from starlite.handlers.base import BaseRouteHandler # type: ignore + from starlite.middleware import DefineMiddleware # type: ignore + from starlite.plugins.base import get_plugin_for_value # type: ignore + from starlite.routes.http import HTTPRoute # type: ignore + from starlite.utils import ConnectionDataExtractor, is_async_callable, Ref # type: ignore + from pydantic import BaseModel # type: ignore +except ImportError: + raise DidNotEnable("Starlite is not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Optional, Union + from starlite.types import ( # type: ignore + ASGIApp, + Hint, + HTTPReceiveMessage, + HTTPScope, + Message, + Middleware, + Receive, + Scope as StarliteScope, + Send, + WebSocketReceiveMessage, + ) + from starlite import MiddlewareProtocol + from sentry_sdk._types import Event + + +_DEFAULT_TRANSACTION_NAME = "generic Starlite request" + + +class StarliteIntegration(Integration): + identifier = "starlite" + origin = f"auto.http.{identifier}" + + @staticmethod + def setup_once() -> None: + patch_app_init() + patch_middlewares() + patch_http_route_handle() + + +class SentryStarliteASGIMiddleware(SentryAsgiMiddleware): + def __init__( + self, app: "ASGIApp", span_origin: str = StarliteIntegration.origin + ) -> None: + super().__init__( + app=app, + unsafe_context_data=False, + transaction_style="endpoint", + mechanism_type="asgi", + span_origin=span_origin, + asgi_version=3, + ) + + +def patch_app_init() -> None: + """ + Replaces the Starlite class's `__init__` function in order to inject `after_exception` handlers and set the + `SentryStarliteASGIMiddleware` as the outmost middleware in the stack. + See: + - https://starlite-api.github.io/starlite/usage/0-the-starlite-app/5-application-hooks/#after-exception + - https://starlite-api.github.io/starlite/usage/7-middleware/0-middleware-intro/ + """ + old__init__ = Starlite.__init__ + + @ensure_integration_enabled(StarliteIntegration, old__init__) + def injection_wrapper(self: "Starlite", *args: "Any", **kwargs: "Any") -> None: + after_exception = kwargs.pop("after_exception", []) + kwargs.update( + after_exception=[ + exception_handler, + *( + after_exception + if isinstance(after_exception, list) + else [after_exception] + ), + ] + ) + + middleware = kwargs.get("middleware") or [] + kwargs["middleware"] = [SentryStarliteASGIMiddleware, *middleware] + old__init__(self, *args, **kwargs) + + Starlite.__init__ = injection_wrapper + + +def patch_middlewares() -> None: + old_resolve_middleware_stack = BaseRouteHandler.resolve_middleware + + @ensure_integration_enabled(StarliteIntegration, old_resolve_middleware_stack) + def resolve_middleware_wrapper(self: "BaseRouteHandler") -> "list[Middleware]": + return [ + enable_span_for_middleware(middleware) + for middleware in old_resolve_middleware_stack(self) + ] + + BaseRouteHandler.resolve_middleware = resolve_middleware_wrapper + + +def enable_span_for_middleware(middleware: "Middleware") -> "Middleware": + if ( + not hasattr(middleware, "__call__") # noqa: B004 + or middleware is SentryStarliteASGIMiddleware + ): + return middleware + + if isinstance(middleware, DefineMiddleware): + old_call: "ASGIApp" = middleware.middleware.__call__ + else: + old_call = middleware.__call__ + + async def _create_span_call( + self: "MiddlewareProtocol", + scope: "StarliteScope", + receive: "Receive", + send: "Send", + ) -> None: + if sentry_sdk.get_client().get_integration(StarliteIntegration) is None: + return await old_call(self, scope, receive, send) + + middleware_name = self.__class__.__name__ + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE, + name=middleware_name, + origin=StarliteIntegration.origin, + ) as middleware_span: + middleware_span.set_tag("starlite.middleware_name", middleware_name) + + # Creating spans for the "receive" callback + async def _sentry_receive( + *args: "Any", **kwargs: "Any" + ) -> "Union[HTTPReceiveMessage, WebSocketReceiveMessage]": + if sentry_sdk.get_client().get_integration(StarliteIntegration) is None: + return await receive(*args, **kwargs) + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE_RECEIVE, + name=getattr(receive, "__qualname__", str(receive)), + origin=StarliteIntegration.origin, + ) as span: + span.set_tag("starlite.middleware_name", middleware_name) + return await receive(*args, **kwargs) + + receive_name = getattr(receive, "__name__", str(receive)) + receive_patched = receive_name == "_sentry_receive" + new_receive = _sentry_receive if not receive_patched else receive + + # Creating spans for the "send" callback + async def _sentry_send(message: "Message") -> None: + if sentry_sdk.get_client().get_integration(StarliteIntegration) is None: + return await send(message) + with sentry_sdk.start_span( + op=OP.MIDDLEWARE_STARLITE_SEND, + name=getattr(send, "__qualname__", str(send)), + origin=StarliteIntegration.origin, + ) as span: + span.set_tag("starlite.middleware_name", middleware_name) + return await send(message) + + send_name = getattr(send, "__name__", str(send)) + send_patched = send_name == "_sentry_send" + new_send = _sentry_send if not send_patched else send + + return await old_call(self, scope, new_receive, new_send) + + not_yet_patched = old_call.__name__ not in ["_create_span_call"] + + if not_yet_patched: + if isinstance(middleware, DefineMiddleware): + middleware.middleware.__call__ = _create_span_call + else: + middleware.__call__ = _create_span_call + + return middleware + + +def patch_http_route_handle() -> None: + old_handle = HTTPRoute.handle + + async def handle_wrapper( + self: "HTTPRoute", scope: "HTTPScope", receive: "Receive", send: "Send" + ) -> None: + if sentry_sdk.get_client().get_integration(StarliteIntegration) is None: + return await old_handle(self, scope, receive, send) + + sentry_scope = sentry_sdk.get_isolation_scope() + request: "Request[Any, Any]" = scope["app"].request_class( + scope=scope, receive=receive, send=send + ) + extracted_request_data = ConnectionDataExtractor( + parse_body=True, parse_query=True + )(request) + body = extracted_request_data.pop("body") + + request_data = await body + + def event_processor(event: "Event", _: "Hint") -> "Event": + route_handler = scope.get("route_handler") + + request_info = event.get("request", {}) + request_info["content_length"] = len(scope.get("_body", b"")) + if should_send_default_pii(): + request_info["cookies"] = extracted_request_data["cookies"] + if request_data is not None: + request_info["data"] = request_data + + func = None + if route_handler.name is not None: + tx_name = route_handler.name + elif isinstance(route_handler.fn, Ref): + func = route_handler.fn.value + else: + func = route_handler.fn + if func is not None: + tx_name = transaction_from_function(func) + + tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]} + + if not tx_name: + tx_name = _DEFAULT_TRANSACTION_NAME + tx_info = {"source": TransactionSource.ROUTE} + + event.update( + { + "request": deepcopy(request_info), + "transaction": tx_name, + "transaction_info": tx_info, + } + ) + return event + + sentry_scope._name = StarliteIntegration.identifier + sentry_scope.add_event_processor(event_processor) + + return await old_handle(self, scope, receive, send) + + HTTPRoute.handle = handle_wrapper + + +def retrieve_user_from_scope(scope: "StarliteScope") -> "Optional[dict[str, Any]]": + scope_user = scope.get("user") + if not scope_user: + return None + if isinstance(scope_user, dict): + return scope_user + if isinstance(scope_user, BaseModel): + return scope_user.dict() + if hasattr(scope_user, "asdict"): # dataclasses + return scope_user.asdict() + + plugin = get_plugin_for_value(scope_user) + if plugin and not is_async_callable(plugin.to_dict): + return plugin.to_dict(scope_user) + + return None + + +@ensure_integration_enabled(StarliteIntegration) +def exception_handler(exc: Exception, scope: "StarliteScope", _: "State") -> None: + user_info: "Optional[dict[str, Any]]" = None + if should_send_default_pii(): + user_info = retrieve_user_from_scope(scope) + if user_info and isinstance(user_info, dict): + sentry_scope = sentry_sdk.get_isolation_scope() + sentry_scope.set_user(user_info) + + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": StarliteIntegration.identifier, "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/statsig.py b/lib/python3.12/site-packages/sentry_sdk/integrations/statsig.py new file mode 100644 index 0000000000000000000000000000000000000000..42e71a50f55234be96d09ca13674fb72c9ca4668 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/statsig.py @@ -0,0 +1,37 @@ +from functools import wraps +from typing import Any, TYPE_CHECKING + +from sentry_sdk.feature_flags import add_feature_flag +from sentry_sdk.integrations import Integration, DidNotEnable, _check_minimum_version +from sentry_sdk.utils import parse_version + +try: + from statsig import statsig as statsig_module + from statsig.version import __version__ as STATSIG_VERSION +except ImportError: + raise DidNotEnable("statsig is not installed") + +if TYPE_CHECKING: + from statsig.statsig_user import StatsigUser + + +class StatsigIntegration(Integration): + identifier = "statsig" + + @staticmethod + def setup_once() -> None: + version = parse_version(STATSIG_VERSION) + _check_minimum_version(StatsigIntegration, version, "statsig") + + # Wrap and patch evaluation method(s) in the statsig module + old_check_gate = statsig_module.check_gate + + @wraps(old_check_gate) + def sentry_check_gate( + user: "StatsigUser", gate: str, *args: "Any", **kwargs: "Any" + ) -> "Any": + enabled = old_check_gate(user, gate, *args, **kwargs) + add_feature_flag(gate, enabled) + return enabled + + statsig_module.check_gate = sentry_check_gate diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/stdlib.py b/lib/python3.12/site-packages/sentry_sdk/integrations/stdlib.py new file mode 100644 index 0000000000000000000000000000000000000000..e3120a3b32a04e3234ac6b03b77371eac017daec --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/stdlib.py @@ -0,0 +1,289 @@ +import os +import subprocess +import sys +import platform +from http.client import HTTPConnection + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import add_global_event_processor +from sentry_sdk.tracing_utils import ( + EnvironHeaders, + should_propagate_trace, + add_http_request_source, +) +from sentry_sdk.utils import ( + SENSITIVE_DATA_SUBSTITUTE, + capture_internal_exceptions, + ensure_integration_enabled, + is_sentry_url, + logger, + safe_repr, + parse_url, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Callable + from typing import Dict + from typing import Optional + from typing import List + + from sentry_sdk._types import Event, Hint + + +_RUNTIME_CONTEXT: "dict[str, object]" = { + "name": platform.python_implementation(), + "version": "%s.%s.%s" % (sys.version_info[:3]), + "build": sys.version, +} + + +class StdlibIntegration(Integration): + identifier = "stdlib" + + @staticmethod + def setup_once() -> None: + _install_httplib() + _install_subprocess() + + @add_global_event_processor + def add_python_runtime_context( + event: "Event", hint: "Hint" + ) -> "Optional[Event]": + if sentry_sdk.get_client().get_integration(StdlibIntegration) is not None: + contexts = event.setdefault("contexts", {}) + if isinstance(contexts, dict) and "runtime" not in contexts: + contexts["runtime"] = _RUNTIME_CONTEXT + + return event + + +def _install_httplib() -> None: + real_putrequest = HTTPConnection.putrequest + real_getresponse = HTTPConnection.getresponse + + def putrequest( + self: "HTTPConnection", method: str, url: str, *args: "Any", **kwargs: "Any" + ) -> "Any": + default_port = self.default_port + + # proxies go through set_tunnel + tunnel_host = getattr(self, "_tunnel_host", None) + if tunnel_host: + host = tunnel_host + port = getattr(self, "_tunnel_port", default_port) + else: + host = self.host + port = self.port + + client = sentry_sdk.get_client() + if client.get_integration(StdlibIntegration) is None or is_sentry_url( + client, host + ): + return real_putrequest(self, method, url, *args, **kwargs) + + real_url = url + if real_url is None or not real_url.startswith(("http://", "https://")): + real_url = "%s://%s%s%s" % ( + default_port == 443 and "https" or "http", + host, + port != default_port and ":%s" % port or "", + url, + ) + + parsed_url = None + with capture_internal_exceptions(): + parsed_url = parse_url(real_url, sanitize=False) + + span = sentry_sdk.start_span( + op=OP.HTTP_CLIENT, + name="%s %s" + % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE), + origin="auto.http.stdlib.httplib", + ) + span.set_data(SPANDATA.HTTP_METHOD, method) + if parsed_url is not None: + span.set_data("url", parsed_url.url) + span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query) + span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment) + + # for proxies, these point to the proxy host/port + if tunnel_host: + span.set_data(SPANDATA.NETWORK_PEER_ADDRESS, self.host) + span.set_data(SPANDATA.NETWORK_PEER_PORT, self.port) + + rv = real_putrequest(self, method, url, *args, **kwargs) + + if should_propagate_trace(client, real_url): + for ( + key, + value, + ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers( + span=span + ): + logger.debug( + "[Tracing] Adding `{key}` header {value} to outgoing request to {real_url}.".format( + key=key, value=value, real_url=real_url + ) + ) + self.putheader(key, value) + + self._sentrysdk_span = span # type: ignore[attr-defined] + + return rv + + def getresponse(self: "HTTPConnection", *args: "Any", **kwargs: "Any") -> "Any": + span = getattr(self, "_sentrysdk_span", None) + + if span is None: + return real_getresponse(self, *args, **kwargs) + + try: + rv = real_getresponse(self, *args, **kwargs) + + span.set_http_status(int(rv.status)) + span.set_data("reason", rv.reason) + finally: + span.finish() + + with capture_internal_exceptions(): + add_http_request_source(span) + + return rv + + HTTPConnection.putrequest = putrequest # type: ignore[method-assign] + HTTPConnection.getresponse = getresponse # type: ignore[method-assign] + + +def _init_argument( + args: "List[Any]", + kwargs: "Dict[Any, Any]", + name: str, + position: int, + setdefault_callback: "Optional[Callable[[Any], Any]]" = None, +) -> "Any": + """ + given (*args, **kwargs) of a function call, retrieve (and optionally set a + default for) an argument by either name or position. + + This is useful for wrapping functions with complex type signatures and + extracting a few arguments without needing to redefine that function's + entire type signature. + """ + + if name in kwargs: + rv = kwargs[name] + if setdefault_callback is not None: + rv = setdefault_callback(rv) + if rv is not None: + kwargs[name] = rv + elif position < len(args): + rv = args[position] + if setdefault_callback is not None: + rv = setdefault_callback(rv) + if rv is not None: + args[position] = rv + else: + rv = setdefault_callback and setdefault_callback(None) + if rv is not None: + kwargs[name] = rv + + return rv + + +def _install_subprocess() -> None: + old_popen_init = subprocess.Popen.__init__ + + @ensure_integration_enabled(StdlibIntegration, old_popen_init) + def sentry_patched_popen_init( + self: "subprocess.Popen[Any]", *a: "Any", **kw: "Any" + ) -> None: + # Convert from tuple to list to be able to set values. + a = list(a) + + args = _init_argument(a, kw, "args", 0) or [] + cwd = _init_argument(a, kw, "cwd", 9) + + # if args is not a list or tuple (and e.g. some iterator instead), + # let's not use it at all. There are too many things that can go wrong + # when trying to collect an iterator into a list and setting that list + # into `a` again. + # + # Also invocations where `args` is not a sequence are not actually + # legal. They just happen to work under CPython. + description = None + + if isinstance(args, (list, tuple)) and len(args) < 100: + with capture_internal_exceptions(): + description = " ".join(map(str, args)) + + if description is None: + description = safe_repr(args) + + env = None + + with sentry_sdk.start_span( + op=OP.SUBPROCESS, + name=description, + origin="auto.subprocess.stdlib.subprocess", + ) as span: + for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers( + span=span + ): + if env is None: + env = _init_argument( + a, + kw, + "env", + 10, + lambda x: dict(x if x is not None else os.environ), + ) + env["SUBPROCESS_" + k.upper().replace("-", "_")] = v + + if cwd: + span.set_data("subprocess.cwd", cwd) + + rv = old_popen_init(self, *a, **kw) + + span.set_tag("subprocess.pid", self.pid) + return rv + + subprocess.Popen.__init__ = sentry_patched_popen_init # type: ignore + + old_popen_wait = subprocess.Popen.wait + + @ensure_integration_enabled(StdlibIntegration, old_popen_wait) + def sentry_patched_popen_wait( + self: "subprocess.Popen[Any]", *a: "Any", **kw: "Any" + ) -> "Any": + with sentry_sdk.start_span( + op=OP.SUBPROCESS_WAIT, + origin="auto.subprocess.stdlib.subprocess", + ) as span: + span.set_tag("subprocess.pid", self.pid) + return old_popen_wait(self, *a, **kw) + + subprocess.Popen.wait = sentry_patched_popen_wait # type: ignore + + old_popen_communicate = subprocess.Popen.communicate + + @ensure_integration_enabled(StdlibIntegration, old_popen_communicate) + def sentry_patched_popen_communicate( + self: "subprocess.Popen[Any]", *a: "Any", **kw: "Any" + ) -> "Any": + with sentry_sdk.start_span( + op=OP.SUBPROCESS_COMMUNICATE, + origin="auto.subprocess.stdlib.subprocess", + ) as span: + span.set_tag("subprocess.pid", self.pid) + return old_popen_communicate(self, *a, **kw) + + subprocess.Popen.communicate = sentry_patched_popen_communicate # type: ignore + + +def get_subprocess_traceparent_headers() -> "EnvironHeaders": + return EnvironHeaders(os.environ, prefix="SUBPROCESS_") diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/strawberry.py b/lib/python3.12/site-packages/sentry_sdk/integrations/strawberry.py new file mode 100644 index 0000000000000000000000000000000000000000..da3c31a967dcd484c3ef79f1a7d2811ed13b399b --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/strawberry.py @@ -0,0 +1,406 @@ +import functools +import hashlib +import warnings +from inspect import isawaitable + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.integrations.logging import ignore_logger +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + capture_internal_exceptions, + ensure_integration_enabled, + event_from_exception, + logger, + package_version, + _get_installed_modules, +) + +try: + from functools import cached_property +except ImportError: + # The strawberry integration requires Python 3.8+. functools.cached_property + # was added in 3.8, so this check is technically not needed, but since this + # is an auto-enabling integration, we might get to executing this import in + # lower Python versions, so we need to deal with it. + raise DidNotEnable("strawberry-graphql integration requires Python 3.8 or newer") + +try: + from strawberry import Schema + from strawberry.extensions import SchemaExtension + from strawberry.extensions.tracing.utils import ( + should_skip_tracing as strawberry_should_skip_tracing, + ) + from strawberry.http import async_base_view, sync_base_view +except ImportError: + raise DidNotEnable("strawberry-graphql is not installed") + +try: + from strawberry.extensions.tracing import ( + SentryTracingExtension as StrawberrySentryAsyncExtension, + SentryTracingExtensionSync as StrawberrySentrySyncExtension, + ) +except ImportError: + StrawberrySentryAsyncExtension = None + StrawberrySentrySyncExtension = None + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Callable, Generator, List, Optional + from graphql import GraphQLError, GraphQLResolveInfo + from strawberry.http import GraphQLHTTPResponse + from strawberry.types import ExecutionContext + from sentry_sdk._types import Event, EventProcessor + + +ignore_logger("strawberry.execution") + + +class StrawberryIntegration(Integration): + identifier = "strawberry" + origin = f"auto.graphql.{identifier}" + + def __init__(self, async_execution: "Optional[bool]" = None) -> None: + if async_execution not in (None, False, True): + raise ValueError( + 'Invalid value for async_execution: "{}" (must be bool)'.format( + async_execution + ) + ) + self.async_execution = async_execution + + @staticmethod + def setup_once() -> None: + version = package_version("strawberry-graphql") + _check_minimum_version(StrawberryIntegration, version, "strawberry-graphql") + + _patch_schema_init() + _patch_views() + + +def _patch_schema_init() -> None: + old_schema_init = Schema.__init__ + + @functools.wraps(old_schema_init) + def _sentry_patched_schema_init( + self: "Schema", *args: "Any", **kwargs: "Any" + ) -> None: + integration = sentry_sdk.get_client().get_integration(StrawberryIntegration) + if integration is None: + return old_schema_init(self, *args, **kwargs) + + extensions = kwargs.get("extensions") or [] + + should_use_async_extension: "Optional[bool]" = None + if integration.async_execution is not None: + should_use_async_extension = integration.async_execution + else: + # try to figure it out ourselves + should_use_async_extension = _guess_if_using_async(extensions) + + if should_use_async_extension is None: + warnings.warn( + "Assuming strawberry is running sync. If not, initialize the integration as StrawberryIntegration(async_execution=True).", + stacklevel=2, + ) + should_use_async_extension = False + + # remove the built in strawberry sentry extension, if present + extensions = [ + extension + for extension in extensions + if extension + not in (StrawberrySentryAsyncExtension, StrawberrySentrySyncExtension) + ] + + # add our extension + extensions.append( + SentryAsyncExtension if should_use_async_extension else SentrySyncExtension + ) + + kwargs["extensions"] = extensions + + return old_schema_init(self, *args, **kwargs) + + Schema.__init__ = _sentry_patched_schema_init # type: ignore[method-assign] + + +class SentryAsyncExtension(SchemaExtension): + def __init__( + self: "Any", + *, + execution_context: "Optional[ExecutionContext]" = None, + ) -> None: + if execution_context: + self.execution_context = execution_context + + @cached_property + def _resource_name(self) -> str: + query_hash = self.hash_query(self.execution_context.query) # type: ignore + + if self.execution_context.operation_name: + return "{}:{}".format(self.execution_context.operation_name, query_hash) + + return query_hash + + def hash_query(self, query: str) -> str: + return hashlib.md5(query.encode("utf-8")).hexdigest() + + def on_operation(self) -> "Generator[None, None, None]": + self._operation_name = self.execution_context.operation_name + + operation_type = "query" + op = OP.GRAPHQL_QUERY + + if self.execution_context.query is None: + self.execution_context.query = "" + + if self.execution_context.query.strip().startswith("mutation"): + operation_type = "mutation" + op = OP.GRAPHQL_MUTATION + elif self.execution_context.query.strip().startswith("subscription"): + operation_type = "subscription" + op = OP.GRAPHQL_SUBSCRIPTION + + description = operation_type + if self._operation_name: + description += " {}".format(self._operation_name) + + sentry_sdk.add_breadcrumb( + category="graphql.operation", + data={ + "operation_name": self._operation_name, + "operation_type": operation_type, + }, + ) + + scope = sentry_sdk.get_isolation_scope() + event_processor = _make_request_event_processor(self.execution_context) + scope.add_event_processor(event_processor) + + span = sentry_sdk.get_current_span() + if span: + self.graphql_span = span.start_child( + op=op, + name=description, + origin=StrawberryIntegration.origin, + ) + else: + self.graphql_span = sentry_sdk.start_span( + op=op, + name=description, + origin=StrawberryIntegration.origin, + ) + + self.graphql_span.set_data("graphql.operation.type", operation_type) + self.graphql_span.set_data("graphql.operation.name", self._operation_name) + self.graphql_span.set_data("graphql.document", self.execution_context.query) + self.graphql_span.set_data("graphql.resource_name", self._resource_name) + + yield + + transaction = self.graphql_span.containing_transaction + if transaction and self.execution_context.operation_name: + transaction.name = self.execution_context.operation_name + transaction.source = TransactionSource.COMPONENT + transaction.op = op + + self.graphql_span.finish() + + def on_validate(self) -> "Generator[None, None, None]": + self.validation_span = self.graphql_span.start_child( + op=OP.GRAPHQL_VALIDATE, + name="validation", + origin=StrawberryIntegration.origin, + ) + + yield + + self.validation_span.finish() + + def on_parse(self) -> "Generator[None, None, None]": + self.parsing_span = self.graphql_span.start_child( + op=OP.GRAPHQL_PARSE, + name="parsing", + origin=StrawberryIntegration.origin, + ) + + yield + + self.parsing_span.finish() + + def should_skip_tracing( + self, + _next: "Callable[[Any, GraphQLResolveInfo, Any, Any], Any]", + info: "GraphQLResolveInfo", + ) -> bool: + return strawberry_should_skip_tracing(_next, info) + + async def _resolve( + self, + _next: "Callable[[Any, GraphQLResolveInfo, Any, Any], Any]", + root: "Any", + info: "GraphQLResolveInfo", + *args: str, + **kwargs: "Any", + ) -> "Any": + result = _next(root, info, *args, **kwargs) + + if isawaitable(result): + result = await result + + return result + + async def resolve( + self, + _next: "Callable[[Any, GraphQLResolveInfo, Any, Any], Any]", + root: "Any", + info: "GraphQLResolveInfo", + *args: str, + **kwargs: "Any", + ) -> "Any": + if self.should_skip_tracing(_next, info): + return await self._resolve(_next, root, info, *args, **kwargs) + + field_path = "{}.{}".format(info.parent_type, info.field_name) + + with self.graphql_span.start_child( + op=OP.GRAPHQL_RESOLVE, + name="resolving {}".format(field_path), + origin=StrawberryIntegration.origin, + ) as span: + span.set_data("graphql.field_name", info.field_name) + span.set_data("graphql.parent_type", info.parent_type.name) + span.set_data("graphql.field_path", field_path) + span.set_data("graphql.path", ".".join(map(str, info.path.as_list()))) + + return await self._resolve(_next, root, info, *args, **kwargs) + + +class SentrySyncExtension(SentryAsyncExtension): + def resolve( + self, + _next: "Callable[[Any, Any, Any, Any], Any]", + root: "Any", + info: "GraphQLResolveInfo", + *args: str, + **kwargs: "Any", + ) -> "Any": + if self.should_skip_tracing(_next, info): + return _next(root, info, *args, **kwargs) + + field_path = "{}.{}".format(info.parent_type, info.field_name) + + with self.graphql_span.start_child( + op=OP.GRAPHQL_RESOLVE, + name="resolving {}".format(field_path), + origin=StrawberryIntegration.origin, + ) as span: + span.set_data("graphql.field_name", info.field_name) + span.set_data("graphql.parent_type", info.parent_type.name) + span.set_data("graphql.field_path", field_path) + span.set_data("graphql.path", ".".join(map(str, info.path.as_list()))) + + return _next(root, info, *args, **kwargs) + + +def _patch_views() -> None: + old_async_view_handle_errors = async_base_view.AsyncBaseHTTPView._handle_errors + old_sync_view_handle_errors = sync_base_view.SyncBaseHTTPView._handle_errors + + def _sentry_patched_async_view_handle_errors( + self: "Any", errors: "List[GraphQLError]", response_data: "GraphQLHTTPResponse" + ) -> None: + old_async_view_handle_errors(self, errors, response_data) + _sentry_patched_handle_errors(self, errors, response_data) + + def _sentry_patched_sync_view_handle_errors( + self: "Any", errors: "List[GraphQLError]", response_data: "GraphQLHTTPResponse" + ) -> None: + old_sync_view_handle_errors(self, errors, response_data) + _sentry_patched_handle_errors(self, errors, response_data) + + @ensure_integration_enabled(StrawberryIntegration) + def _sentry_patched_handle_errors( + self: "Any", errors: "List[GraphQLError]", response_data: "GraphQLHTTPResponse" + ) -> None: + if not errors: + return + + scope = sentry_sdk.get_isolation_scope() + event_processor = _make_response_event_processor(response_data) + scope.add_event_processor(event_processor) + + with capture_internal_exceptions(): + for error in errors: + event, hint = event_from_exception( + error, + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": StrawberryIntegration.identifier, + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + + async_base_view.AsyncBaseHTTPView._handle_errors = ( # type: ignore[method-assign] + _sentry_patched_async_view_handle_errors + ) + sync_base_view.SyncBaseHTTPView._handle_errors = ( # type: ignore[method-assign] + _sentry_patched_sync_view_handle_errors + ) + + +def _make_request_event_processor( + execution_context: "ExecutionContext", +) -> "EventProcessor": + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + with capture_internal_exceptions(): + if should_send_default_pii(): + request_data = event.setdefault("request", {}) + request_data["api_target"] = "graphql" + + if not request_data.get("data"): + data: "dict[str, Any]" = {"query": execution_context.query} + if execution_context.variables: + data["variables"] = execution_context.variables + if execution_context.operation_name: + data["operationName"] = execution_context.operation_name + + request_data["data"] = data + + else: + try: + del event["request"]["data"] + except (KeyError, TypeError): + pass + + return event + + return inner + + +def _make_response_event_processor( + response_data: "GraphQLHTTPResponse", +) -> "EventProcessor": + def inner(event: "Event", hint: "dict[str, Any]") -> "Event": + with capture_internal_exceptions(): + if should_send_default_pii(): + contexts = event.setdefault("contexts", {}) + contexts["response"] = {"data": response_data} + + return event + + return inner + + +def _guess_if_using_async(extensions: "List[SchemaExtension]") -> "Optional[bool]": + if StrawberrySentryAsyncExtension in extensions: + return True + elif StrawberrySentrySyncExtension in extensions: + return False + + return None diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/sys_exit.py b/lib/python3.12/site-packages/sentry_sdk/integrations/sys_exit.py new file mode 100644 index 0000000000000000000000000000000000000000..120576ed94c46946666f1a18e24dae97fcb4b091 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/sys_exit.py @@ -0,0 +1,65 @@ +import functools +import sys + +import sentry_sdk +from sentry_sdk.utils import capture_internal_exceptions, event_from_exception +from sentry_sdk.integrations import Integration +from sentry_sdk._types import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + from typing import NoReturn, Union + + +class SysExitIntegration(Integration): + """Captures sys.exit calls and sends them as events to Sentry. + + By default, SystemExit exceptions are not captured by the SDK. Enabling this integration will capture SystemExit + exceptions generated by sys.exit calls and send them to Sentry. + + This integration, in its default configuration, only captures the sys.exit call if the exit code is a non-zero and + non-None value (unsuccessful exits). Pass `capture_successful_exits=True` to capture successful exits as well. + Note that the integration does not capture SystemExit exceptions raised outside a call to sys.exit. + """ + + identifier = "sys_exit" + + def __init__(self, *, capture_successful_exits: bool = False) -> None: + self._capture_successful_exits = capture_successful_exits + + @staticmethod + def setup_once() -> None: + SysExitIntegration._patch_sys_exit() + + @staticmethod + def _patch_sys_exit() -> None: + old_exit: "Callable[[Union[str, int, None]], NoReturn]" = sys.exit + + @functools.wraps(old_exit) + def sentry_patched_exit(__status: "Union[str, int, None]" = 0) -> "NoReturn": + # @ensure_integration_enabled ensures that this is non-None + integration = sentry_sdk.get_client().get_integration(SysExitIntegration) + if integration is None: + old_exit(__status) + + try: + old_exit(__status) + except SystemExit as e: + with capture_internal_exceptions(): + if integration._capture_successful_exits or __status not in ( + 0, + None, + ): + _capture_exception(e) + raise e + + sys.exit = sentry_patched_exit + + +def _capture_exception(exc: "SystemExit") -> None: + event, hint = event_from_exception( + exc, + client_options=sentry_sdk.get_client().options, + mechanism={"type": SysExitIntegration.identifier, "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/threading.py b/lib/python3.12/site-packages/sentry_sdk/integrations/threading.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5633507f24a28dfffcd8fbe823ec8fe031eba2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/threading.py @@ -0,0 +1,200 @@ +import sys +import warnings +from functools import wraps +from threading import Thread, current_thread +from concurrent.futures import ThreadPoolExecutor, Future + +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.scope import use_isolation_scope, use_scope +from sentry_sdk.utils import ( + event_from_exception, + capture_internal_exceptions, + logger, + reraise, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import TypeVar + from typing import Callable + from typing import Optional + + from sentry_sdk._types import ExcInfo + + F = TypeVar("F", bound=Callable[..., Any]) + T = TypeVar("T", bound=Any) + + +class ThreadingIntegration(Integration): + identifier = "threading" + + def __init__( + self, propagate_hub: "Optional[bool]" = None, propagate_scope: bool = True + ) -> None: + if propagate_hub is not None: + logger.warning( + "Deprecated: propagate_hub is deprecated. This will be removed in the future." + ) + + # Note: propagate_hub did not have any effect on propagation of scope data + # scope data was always propagated no matter what the value of propagate_hub was + # This is why the default for propagate_scope is True + + self.propagate_scope = propagate_scope + + if propagate_hub is not None: + self.propagate_scope = propagate_hub + + @staticmethod + def setup_once() -> None: + old_start = Thread.start + + try: + from django import VERSION as django_version # noqa: N811 + except ImportError: + django_version = None + + try: + import channels # type: ignore[import-untyped] + + channels_version = channels.__version__ + except (ImportError, AttributeError): + channels_version = None + + is_async_emulated_with_threads = ( + sys.version_info < (3, 9) + and channels_version is not None + and channels_version < "4.0.0" + and django_version is not None + and django_version >= (3, 0) + and django_version < (4, 0) + ) + + @wraps(old_start) + def sentry_start(self: "Thread", *a: "Any", **kw: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(ThreadingIntegration) + if integration is None: + return old_start(self, *a, **kw) + + if integration.propagate_scope: + if is_async_emulated_with_threads: + warnings.warn( + "There is a known issue with Django channels 2.x and 3.x when using Python 3.8 or older. " + "(Async support is emulated using threads and some Sentry data may be leaked between those threads.) " + "Please either upgrade to Django channels 4.0+, use Django's async features " + "available in Django 3.1+ instead of Django channels, or upgrade to Python 3.9+.", + stacklevel=2, + ) + isolation_scope = sentry_sdk.get_isolation_scope() + current_scope = sentry_sdk.get_current_scope() + + else: + isolation_scope = sentry_sdk.get_isolation_scope().fork() + current_scope = sentry_sdk.get_current_scope().fork() + else: + isolation_scope = None + current_scope = None + + # Patching instance methods in `start()` creates a reference cycle if + # done in a naive way. See + # https://github.com/getsentry/sentry-python/pull/434 + # + # In threading module, using current_thread API will access current thread instance + # without holding it to avoid a reference cycle in an easier way. + with capture_internal_exceptions(): + new_run = _wrap_run( + isolation_scope, + current_scope, + getattr(self.run, "__func__", self.run), + ) + self.run = new_run # type: ignore + + return old_start(self, *a, **kw) + + Thread.start = sentry_start # type: ignore + ThreadPoolExecutor.submit = _wrap_threadpool_executor_submit( # type: ignore + ThreadPoolExecutor.submit, is_async_emulated_with_threads + ) + + +def _wrap_run( + isolation_scope_to_use: "Optional[sentry_sdk.Scope]", + current_scope_to_use: "Optional[sentry_sdk.Scope]", + old_run_func: "F", +) -> "F": + @wraps(old_run_func) + def run(*a: "Any", **kw: "Any") -> "Any": + def _run_old_run_func() -> "Any": + try: + self = current_thread() + return old_run_func(self, *a[1:], **kw) + except Exception: + reraise(*_capture_exception()) + + if isolation_scope_to_use is not None and current_scope_to_use is not None: + with use_isolation_scope(isolation_scope_to_use): + with use_scope(current_scope_to_use): + return _run_old_run_func() + else: + return _run_old_run_func() + + return run # type: ignore + + +def _wrap_threadpool_executor_submit( + func: "Callable[..., Future[T]]", is_async_emulated_with_threads: bool +) -> "Callable[..., Future[T]]": + """ + Wrap submit call to propagate scopes on task submission. + """ + + @wraps(func) + def sentry_submit( + self: "ThreadPoolExecutor", + fn: "Callable[..., T]", + *args: "Any", + **kwargs: "Any", + ) -> "Future[T]": + integration = sentry_sdk.get_client().get_integration(ThreadingIntegration) + if integration is None: + return func(self, fn, *args, **kwargs) + + if integration.propagate_scope and is_async_emulated_with_threads: + isolation_scope = sentry_sdk.get_isolation_scope() + current_scope = sentry_sdk.get_current_scope() + elif integration.propagate_scope: + isolation_scope = sentry_sdk.get_isolation_scope().fork() + current_scope = sentry_sdk.get_current_scope().fork() + else: + isolation_scope = None + current_scope = None + + def wrapped_fn(*args: "Any", **kwargs: "Any") -> "Any": + if isolation_scope is not None and current_scope is not None: + with use_isolation_scope(isolation_scope): + with use_scope(current_scope): + return fn(*args, **kwargs) + + return fn(*args, **kwargs) + + return func(self, wrapped_fn, *args, **kwargs) + + return sentry_submit + + +def _capture_exception() -> "ExcInfo": + exc_info = sys.exc_info() + + client = sentry_sdk.get_client() + if client.get_integration(ThreadingIntegration) is not None: + event, hint = event_from_exception( + exc_info, + client_options=client.options, + mechanism={"type": "threading", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return exc_info diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/tornado.py b/lib/python3.12/site-packages/sentry_sdk/integrations/tornado.py new file mode 100644 index 0000000000000000000000000000000000000000..96a7629c534f00ff1a8bb311cc596f3c3dea67cb --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/tornado.py @@ -0,0 +1,223 @@ +import weakref +import contextlib +from inspect import iscoroutinefunction + +import sentry_sdk +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.tracing import TransactionSource +from sentry_sdk.utils import ( + HAS_REAL_CONTEXTVARS, + CONTEXTVARS_ERROR_MESSAGE, + ensure_integration_enabled, + event_from_exception, + capture_internal_exceptions, + transaction_from_function, +) +from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable +from sentry_sdk.integrations._wsgi_common import ( + RequestExtractor, + _filter_headers, + _is_json_content_type, +) +from sentry_sdk.integrations.logging import ignore_logger + +try: + from tornado import version_info as TORNADO_VERSION + from tornado.web import RequestHandler, HTTPError + from tornado.gen import coroutine +except ImportError: + raise DidNotEnable("Tornado not installed") + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Optional + from typing import Dict + from typing import Callable + from typing import Generator + + from sentry_sdk._types import Event, EventProcessor + + +class TornadoIntegration(Integration): + identifier = "tornado" + origin = f"auto.http.{identifier}" + + @staticmethod + def setup_once() -> None: + _check_minimum_version(TornadoIntegration, TORNADO_VERSION) + + if not HAS_REAL_CONTEXTVARS: + # Tornado is async. We better have contextvars or we're going to leak + # state between requests. + raise DidNotEnable( + "The tornado integration for Sentry requires Python 3.7+ or the aiocontextvars package" + + CONTEXTVARS_ERROR_MESSAGE + ) + + ignore_logger("tornado.access") + + old_execute = RequestHandler._execute + + awaitable = iscoroutinefunction(old_execute) + + if awaitable: + # Starting Tornado 6 RequestHandler._execute method is a standard Python coroutine (async/await) + # In that case our method should be a coroutine function too + async def sentry_execute_request_handler( + self: "RequestHandler", *args: "Any", **kwargs: "Any" + ) -> "Any": + with _handle_request_impl(self): + return await old_execute(self, *args, **kwargs) + + else: + + @coroutine # type: ignore + def sentry_execute_request_handler( + self: "RequestHandler", *args: "Any", **kwargs: "Any" + ) -> "Any": + with _handle_request_impl(self): + result = yield from old_execute(self, *args, **kwargs) + return result + + RequestHandler._execute = sentry_execute_request_handler + + old_log_exception = RequestHandler.log_exception + + def sentry_log_exception( + self: "Any", + ty: type, + value: BaseException, + tb: "Any", + *args: "Any", + **kwargs: "Any", + ) -> "Optional[Any]": + _capture_exception(ty, value, tb) + return old_log_exception(self, ty, value, tb, *args, **kwargs) + + RequestHandler.log_exception = sentry_log_exception + + +@contextlib.contextmanager +def _handle_request_impl(self: "RequestHandler") -> "Generator[None, None, None]": + integration = sentry_sdk.get_client().get_integration(TornadoIntegration) + + if integration is None: + yield + + weak_handler = weakref.ref(self) + + with sentry_sdk.isolation_scope() as scope: + headers = self.request.headers + + scope.clear_breadcrumbs() + processor = _make_event_processor(weak_handler) + scope.add_event_processor(processor) + + transaction = continue_trace( + headers, + op=OP.HTTP_SERVER, + # Like with all other integrations, this is our + # fallback transaction in case there is no route. + # sentry_urldispatcher_resolve is responsible for + # setting a transaction name later. + name="generic Tornado request", + source=TransactionSource.ROUTE, + origin=TornadoIntegration.origin, + ) + + with sentry_sdk.start_transaction( + transaction, custom_sampling_context={"tornado_request": self.request} + ): + yield + + +@ensure_integration_enabled(TornadoIntegration) +def _capture_exception(ty: type, value: BaseException, tb: "Any") -> None: + if isinstance(value, HTTPError): + return + + event, hint = event_from_exception( + (ty, value, tb), + client_options=sentry_sdk.get_client().options, + mechanism={"type": "tornado", "handled": False}, + ) + + sentry_sdk.capture_event(event, hint=hint) + + +def _make_event_processor( + weak_handler: "Callable[[], RequestHandler]", +) -> "EventProcessor": + def tornado_processor(event: "Event", hint: "dict[str, Any]") -> "Event": + handler = weak_handler() + if handler is None: + return event + + request = handler.request + + with capture_internal_exceptions(): + method = getattr(handler, handler.request.method.lower()) + event["transaction"] = transaction_from_function(method) or "" + event["transaction_info"] = {"source": TransactionSource.COMPONENT} + + with capture_internal_exceptions(): + extractor = TornadoRequestExtractor(request) + extractor.extract_into_event(event) + + request_info = event["request"] + + request_info["url"] = "%s://%s%s" % ( + request.protocol, + request.host, + request.path, + ) + + request_info["query_string"] = request.query + request_info["method"] = request.method + request_info["env"] = {"REMOTE_ADDR": request.remote_ip} + request_info["headers"] = _filter_headers(dict(request.headers)) + + if should_send_default_pii(): + try: + current_user = handler.current_user + except Exception: + current_user = None + + if current_user: + event.setdefault("user", {}).setdefault("is_authenticated", True) + + return event + + return tornado_processor + + +class TornadoRequestExtractor(RequestExtractor): + def content_length(self) -> int: + if self.request.body is None: + return 0 + return len(self.request.body) + + def cookies(self) -> "Dict[str, str]": + return {k: v.value for k, v in self.request.cookies.items()} + + def raw_data(self) -> bytes: + return self.request.body + + def form(self) -> "Dict[str, Any]": + return { + k: [v.decode("latin1", "replace") for v in vs] + for k, vs in self.request.body_arguments.items() + } + + def is_json(self) -> bool: + return _is_json_content_type(self.request.headers.get("content-type")) + + def files(self) -> "Dict[str, Any]": + return {k: v[0] for k, v in self.request.files.items() if v} + + def size_of_file(self, file: "Any") -> int: + return len(file.body or ()) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/trytond.py b/lib/python3.12/site-packages/sentry_sdk/integrations/trytond.py new file mode 100644 index 0000000000000000000000000000000000000000..382e7385e239ad86bbc59f894ff72391f74086b2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/trytond.py @@ -0,0 +1,53 @@ +import sentry_sdk +from sentry_sdk.integrations import Integration +from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware +from sentry_sdk.utils import ensure_integration_enabled, event_from_exception +from sentry_sdk.integrations import DidNotEnable + +try: + from trytond.exceptions import TrytonException # type: ignore + from trytond.wsgi import app # type: ignore +except ImportError: + raise DidNotEnable("Trytond is not installed.") + +# TODO: trytond-worker, trytond-cron and trytond-admin intergations + + +class TrytondWSGIIntegration(Integration): + identifier = "trytond_wsgi" + origin = f"auto.http.{identifier}" + + def __init__(self) -> None: + pass + + @staticmethod + def setup_once() -> None: + app.wsgi_app = SentryWsgiMiddleware( + app.wsgi_app, + span_origin=TrytondWSGIIntegration.origin, + ) + + @ensure_integration_enabled(TrytondWSGIIntegration) + def error_handler(e: Exception) -> None: + if isinstance(e, TrytonException): + return + else: + client = sentry_sdk.get_client() + event, hint = event_from_exception( + e, + client_options=client.options, + mechanism={"type": "trytond", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + # Expected error handlers signature was changed + # when the error_handler decorator was introduced + # in Tryton-5.4 + if hasattr(app, "error_handler"): + + @app.error_handler + def _(app, request, e): # type: ignore + error_handler(e) + + else: + app.error_handlers.append(error_handler) diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/typer.py b/lib/python3.12/site-packages/sentry_sdk/integrations/typer.py new file mode 100644 index 0000000000000000000000000000000000000000..2865f3ee6f8365f710275e9622e7b3010e53819d --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/typer.py @@ -0,0 +1,62 @@ +import sentry_sdk +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, +) +from sentry_sdk.integrations import Integration, DidNotEnable + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Callable + from typing import Any + from typing import Type + from typing import Optional + + from types import TracebackType + + Excepthook = Callable[ + [Type[BaseException], BaseException, Optional[TracebackType]], + Any, + ] + +try: + import typer + from typer.main import except_hook +except ImportError: + raise DidNotEnable("Typer not installed") + + +class TyperIntegration(Integration): + identifier = "typer" + + @staticmethod + def setup_once() -> None: + typer.main.except_hook = _make_excepthook(except_hook) # type: ignore + + +def _make_excepthook(old_excepthook: "Excepthook") -> "Excepthook": + def sentry_sdk_excepthook( + type_: "Type[BaseException]", + value: BaseException, + traceback: "Optional[TracebackType]", + ) -> None: + integration = sentry_sdk.get_client().get_integration(TyperIntegration) + + # Note: If we replace this with ensure_integration_enabled then + # we break the exceptiongroup backport; + # See: https://github.com/getsentry/sentry-python/issues/3097 + if integration is None: + return old_excepthook(type_, value, traceback) + + with capture_internal_exceptions(): + event, hint = event_from_exception( + (type_, value, traceback), + client_options=sentry_sdk.get_client().options, + mechanism={"type": "typer", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return old_excepthook(type_, value, traceback) + + return sentry_sdk_excepthook diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/unleash.py b/lib/python3.12/site-packages/sentry_sdk/integrations/unleash.py new file mode 100644 index 0000000000000000000000000000000000000000..304f5c3bd17623064378279e58a8f627f8675b9f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/unleash.py @@ -0,0 +1,33 @@ +from functools import wraps +from typing import Any + +from sentry_sdk.feature_flags import add_feature_flag +from sentry_sdk.integrations import Integration, DidNotEnable + +try: + from UnleashClient import UnleashClient +except ImportError: + raise DidNotEnable("UnleashClient is not installed") + + +class UnleashIntegration(Integration): + identifier = "unleash" + + @staticmethod + def setup_once() -> None: + # Wrap and patch evaluation methods (class methods) + old_is_enabled = UnleashClient.is_enabled + + @wraps(old_is_enabled) + def sentry_is_enabled( + self: "UnleashClient", feature: str, *args: "Any", **kwargs: "Any" + ) -> "Any": + enabled = old_is_enabled(self, feature, *args, **kwargs) + + # We have no way of knowing what type of unleash feature this is, so we have to treat + # it as a boolean / toggle feature. + add_feature_flag(feature, enabled) + + return enabled + + UnleashClient.is_enabled = sentry_is_enabled # type: ignore diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/unraisablehook.py b/lib/python3.12/site-packages/sentry_sdk/integrations/unraisablehook.py new file mode 100644 index 0000000000000000000000000000000000000000..61ef8a008ccc13af34c52b75227255b23d1cb295 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/unraisablehook.py @@ -0,0 +1,52 @@ +import sys + +import sentry_sdk +from sentry_sdk.utils import ( + capture_internal_exceptions, + event_from_exception, +) +from sentry_sdk.integrations import Integration + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Callable + from typing import Any + + +class UnraisablehookIntegration(Integration): + identifier = "unraisablehook" + + @staticmethod + def setup_once() -> None: + sys.unraisablehook = _make_unraisable(sys.unraisablehook) + + +def _make_unraisable( + old_unraisablehook: "Callable[[sys.UnraisableHookArgs], Any]", +) -> "Callable[[sys.UnraisableHookArgs], Any]": + def sentry_sdk_unraisablehook(unraisable: "sys.UnraisableHookArgs") -> None: + integration = sentry_sdk.get_client().get_integration(UnraisablehookIntegration) + + # Note: If we replace this with ensure_integration_enabled then + # we break the exceptiongroup backport; + # See: https://github.com/getsentry/sentry-python/issues/3097 + if integration is None: + return old_unraisablehook(unraisable) + + if unraisable.exc_value and unraisable.exc_traceback: + with capture_internal_exceptions(): + event, hint = event_from_exception( + ( + unraisable.exc_type, + unraisable.exc_value, + unraisable.exc_traceback, + ), + client_options=sentry_sdk.get_client().options, + mechanism={"type": "unraisablehook", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return old_unraisablehook(unraisable) + + return sentry_sdk_unraisablehook diff --git a/lib/python3.12/site-packages/sentry_sdk/integrations/wsgi.py b/lib/python3.12/site-packages/sentry_sdk/integrations/wsgi.py new file mode 100644 index 0000000000000000000000000000000000000000..1576e21a175e309dc6e7ff1ed92e91fee19c3ef5 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/integrations/wsgi.py @@ -0,0 +1,301 @@ +import sys +from functools import partial +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk._werkzeug import _get_headers, get_host +from sentry_sdk.api import continue_trace +from sentry_sdk.consts import OP +from sentry_sdk.integrations._wsgi_common import ( + DEFAULT_HTTP_METHODS_TO_CAPTURE, + _filter_headers, + nullcontext, +) +from sentry_sdk.scope import should_send_default_pii, use_isolation_scope +from sentry_sdk.sessions import track_session +from sentry_sdk.tracing import Transaction, TransactionSource +from sentry_sdk.utils import ( + ContextVar, + capture_internal_exceptions, + event_from_exception, + reraise, +) + +if TYPE_CHECKING: + from typing import Any, Callable, Dict, Iterator, Optional, Protocol, Tuple, TypeVar + + from sentry_sdk._types import Event, EventProcessor + from sentry_sdk.utils import ExcInfo + + WsgiResponseIter = TypeVar("WsgiResponseIter") + WsgiResponseHeaders = TypeVar("WsgiResponseHeaders") + WsgiExcInfo = TypeVar("WsgiExcInfo") + + class StartResponse(Protocol): + def __call__( + self, + status: str, + response_headers: "WsgiResponseHeaders", + exc_info: "Optional[WsgiExcInfo]" = None, + ) -> "WsgiResponseIter": # type: ignore + pass + + +_wsgi_middleware_applied = ContextVar("sentry_wsgi_middleware_applied") + + +def wsgi_decoding_dance(s: str, charset: str = "utf-8", errors: str = "replace") -> str: + return s.encode("latin1").decode(charset, errors) + + +def get_request_url( + environ: "Dict[str, str]", use_x_forwarded_for: bool = False +) -> str: + """Return the absolute URL without query string for the given WSGI + environment.""" + script_name = environ.get("SCRIPT_NAME", "").rstrip("/") + path_info = environ.get("PATH_INFO", "").lstrip("/") + path = f"{script_name}/{path_info}" + + return "%s://%s/%s" % ( + environ.get("wsgi.url_scheme"), + get_host(environ, use_x_forwarded_for), + wsgi_decoding_dance(path).lstrip("/"), + ) + + +class SentryWsgiMiddleware: + __slots__ = ( + "app", + "use_x_forwarded_for", + "span_origin", + "http_methods_to_capture", + ) + + def __init__( + self, + app: "Callable[[Dict[str, str], Callable[..., Any]], Any]", + use_x_forwarded_for: bool = False, + span_origin: str = "manual", + http_methods_to_capture: "Tuple[str, ...]" = DEFAULT_HTTP_METHODS_TO_CAPTURE, + ) -> None: + self.app = app + self.use_x_forwarded_for = use_x_forwarded_for + self.span_origin = span_origin + self.http_methods_to_capture = http_methods_to_capture + + def __call__( + self, environ: "Dict[str, str]", start_response: "Callable[..., Any]" + ) -> "_ScopedResponse": + if _wsgi_middleware_applied.get(False): + return self.app(environ, start_response) + + _wsgi_middleware_applied.set(True) + try: + with sentry_sdk.isolation_scope() as scope: + with track_session(scope, session_mode="request"): + with capture_internal_exceptions(): + scope.clear_breadcrumbs() + scope._name = "wsgi" + scope.add_event_processor( + _make_wsgi_event_processor( + environ, self.use_x_forwarded_for + ) + ) + + method = environ.get("REQUEST_METHOD", "").upper() + transaction = None + if method in self.http_methods_to_capture: + transaction = continue_trace( + environ, + op=OP.HTTP_SERVER, + name="generic WSGI request", + source=TransactionSource.ROUTE, + origin=self.span_origin, + ) + + transaction_context = ( + sentry_sdk.start_transaction( + transaction, + custom_sampling_context={"wsgi_environ": environ}, + ) + if transaction is not None + else nullcontext() + ) + with transaction_context: + try: + response = self.app( + environ, + partial( + _sentry_start_response, start_response, transaction + ), + ) + except BaseException: + reraise(*_capture_exception()) + finally: + _wsgi_middleware_applied.set(False) + + return _ScopedResponse(scope, response) + + +def _sentry_start_response( + old_start_response: "StartResponse", + transaction: "Optional[Transaction]", + status: str, + response_headers: "WsgiResponseHeaders", + exc_info: "Optional[WsgiExcInfo]" = None, +) -> "WsgiResponseIter": # type: ignore[type-var] + with capture_internal_exceptions(): + status_int = int(status.split(" ", 1)[0]) + if transaction is not None: + transaction.set_http_status(status_int) + + if exc_info is None: + # The Django Rest Framework WSGI test client, and likely other + # (incorrect) implementations, cannot deal with the exc_info argument + # if one is present. Avoid providing a third argument if not necessary. + return old_start_response(status, response_headers) + else: + return old_start_response(status, response_headers, exc_info) + + +def _get_environ(environ: "Dict[str, str]") -> "Iterator[Tuple[str, str]]": + """ + Returns our explicitly included environment variables we want to + capture (server name, port and remote addr if pii is enabled). + """ + keys = ["SERVER_NAME", "SERVER_PORT"] + if should_send_default_pii(): + # make debugging of proxy setup easier. Proxy headers are + # in headers. + keys += ["REMOTE_ADDR"] + + for key in keys: + if key in environ: + yield key, environ[key] + + +def get_client_ip(environ: "Dict[str, str]") -> "Optional[Any]": + """ + Infer the user IP address from various headers. This cannot be used in + security sensitive situations since the value may be forged from a client, + but it's good enough for the event payload. + """ + try: + return environ["HTTP_X_FORWARDED_FOR"].split(",")[0].strip() + except (KeyError, IndexError): + pass + + try: + return environ["HTTP_X_REAL_IP"] + except KeyError: + pass + + return environ.get("REMOTE_ADDR") + + +def _capture_exception() -> "ExcInfo": + """ + Captures the current exception and sends it to Sentry. + Returns the ExcInfo tuple to it can be reraised afterwards. + """ + exc_info = sys.exc_info() + e = exc_info[1] + + # SystemExit(0) is the only uncaught exception that is expected behavior + should_skip_capture = isinstance(e, SystemExit) and e.code in (0, None) + if not should_skip_capture: + event, hint = event_from_exception( + exc_info, + client_options=sentry_sdk.get_client().options, + mechanism={"type": "wsgi", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + + return exc_info + + +class _ScopedResponse: + """ + Users a separate scope for each response chunk. + + This will make WSGI apps more tolerant against: + - WSGI servers streaming responses from a different thread/from + different threads than the one that called start_response + - close() not being called + - WSGI servers streaming responses interleaved from the same thread + """ + + __slots__ = ("_response", "_scope") + + def __init__( + self, scope: "sentry_sdk.scope.Scope", response: "Iterator[bytes]" + ) -> None: + self._scope = scope + self._response = response + + def __iter__(self) -> "Iterator[bytes]": + iterator = iter(self._response) + + while True: + with use_isolation_scope(self._scope): + try: + chunk = next(iterator) + except StopIteration: + break + except BaseException: + reraise(*_capture_exception()) + + yield chunk + + def close(self) -> None: + with use_isolation_scope(self._scope): + try: + self._response.close() # type: ignore + except AttributeError: + pass + except BaseException: + reraise(*_capture_exception()) + + +def _make_wsgi_event_processor( + environ: "Dict[str, str]", use_x_forwarded_for: bool +) -> "EventProcessor": + # It's a bit unfortunate that we have to extract and parse the request data + # from the environ so eagerly, but there are a few good reasons for this. + # + # We might be in a situation where the scope never gets torn down + # properly. In that case we will have an unnecessary strong reference to + # all objects in the environ (some of which may take a lot of memory) when + # we're really just interested in a few of them. + # + # Keeping the environment around for longer than the request lifecycle is + # also not necessarily something uWSGI can deal with: + # https://github.com/unbit/uwsgi/issues/1950 + + client_ip = get_client_ip(environ) + request_url = get_request_url(environ, use_x_forwarded_for) + query_string = environ.get("QUERY_STRING") + method = environ.get("REQUEST_METHOD") + env = dict(_get_environ(environ)) + headers = _filter_headers(dict(_get_headers(environ))) + + def event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event": + with capture_internal_exceptions(): + # if the code below fails halfway through we at least have some data + request_info = event.setdefault("request", {}) + + if should_send_default_pii(): + user_info = event.setdefault("user", {}) + if client_ip: + user_info.setdefault("ip_address", client_ip) + + request_info["url"] = request_url + request_info["query_string"] = query_string + request_info["method"] = method + request_info["env"] = env + request_info["headers"] = headers + + return event + + return event_processor diff --git a/lib/python3.12/site-packages/sentry_sdk/logger.py b/lib/python3.12/site-packages/sentry_sdk/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..4a90fef70b06a52de53f9c47077a025c09d998db --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/logger.py @@ -0,0 +1,88 @@ +# NOTE: this is the logger sentry exposes to users, not some generic logger. +import functools +import time +from typing import Any, TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.utils import format_attribute, safe_repr, capture_internal_exceptions + +if TYPE_CHECKING: + from sentry_sdk._types import Attributes, Log + + +OTEL_RANGES = [ + # ((severity level range), severity text) + # https://opentelemetry.io/docs/specs/otel/logs/data-model + ((1, 4), "trace"), + ((5, 8), "debug"), + ((9, 12), "info"), + ((13, 16), "warn"), + ((17, 20), "error"), + ((21, 24), "fatal"), +] + + +class _dict_default_key(dict): # type: ignore[type-arg] + """dict that returns the key if missing.""" + + def __missing__(self, key: str) -> str: + return "{" + key + "}" + + +def _capture_log( + severity_text: str, severity_number: int, template: str, **kwargs: "Any" +) -> None: + body = template + + attributes: "Attributes" = {} + + if "attributes" in kwargs: + provided_attributes = kwargs.pop("attributes") or {} + for attribute, value in provided_attributes.items(): + attributes[attribute] = format_attribute(value) + + for k, v in kwargs.items(): + attributes[f"sentry.message.parameter.{k}"] = format_attribute(v) + + if kwargs: + # only attach template if there are parameters + attributes["sentry.message.template"] = format_attribute(template) + + with capture_internal_exceptions(): + body = template.format_map(_dict_default_key(kwargs)) + + sentry_sdk.get_current_scope()._capture_log( + { + "severity_text": severity_text, + "severity_number": severity_number, + "attributes": attributes, + "body": body, + "time_unix_nano": time.time_ns(), + "trace_id": None, + "span_id": None, + } + ) + + +trace = functools.partial(_capture_log, "trace", 1) +debug = functools.partial(_capture_log, "debug", 5) +info = functools.partial(_capture_log, "info", 9) +warning = functools.partial(_capture_log, "warn", 13) +error = functools.partial(_capture_log, "error", 17) +fatal = functools.partial(_capture_log, "fatal", 21) + + +def _otel_severity_text(otel_severity_number: int) -> str: + for (lower, upper), severity in OTEL_RANGES: + if lower <= otel_severity_number <= upper: + return severity + + return "default" + + +def _log_level_to_otel(level: int, mapping: "dict[Any, int]") -> "tuple[int, str]": + for py_level, otel_severity_number in sorted(mapping.items(), reverse=True): + if level >= py_level: + return otel_severity_number, _otel_severity_text(otel_severity_number) + + return 0, "default" diff --git a/lib/python3.12/site-packages/sentry_sdk/metrics.py b/lib/python3.12/site-packages/sentry_sdk/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..167e49da05d65b5c8ffce501edef5404dd4a1641 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/metrics.py @@ -0,0 +1,62 @@ +import time +from typing import Any, Optional, TYPE_CHECKING, Union + +import sentry_sdk +from sentry_sdk.utils import format_attribute, safe_repr + +if TYPE_CHECKING: + from sentry_sdk._types import Attributes, Metric, MetricType + + +def _capture_metric( + name: str, + metric_type: "MetricType", + value: float, + unit: "Optional[str]" = None, + attributes: "Optional[Attributes]" = None, +) -> None: + attrs: "Attributes" = {} + + if attributes: + for k, v in attributes.items(): + attrs[k] = format_attribute(v) + + metric: "Metric" = { + "timestamp": time.time(), + "trace_id": None, + "span_id": None, + "name": name, + "type": metric_type, + "value": float(value), + "unit": unit, + "attributes": attrs, + } + + sentry_sdk.get_current_scope()._capture_metric(metric) + + +def count( + name: str, + value: float, + unit: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, +) -> None: + _capture_metric(name, "counter", value, unit, attributes) + + +def gauge( + name: str, + value: float, + unit: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, +) -> None: + _capture_metric(name, "gauge", value, unit, attributes) + + +def distribution( + name: str, + value: float, + unit: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, +) -> None: + _capture_metric(name, "distribution", value, unit, attributes) diff --git a/lib/python3.12/site-packages/sentry_sdk/monitor.py b/lib/python3.12/site-packages/sentry_sdk/monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..eeb262a84a172415ffee0005204e33a3bb9ce539 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/monitor.py @@ -0,0 +1,113 @@ +import os +import time +from threading import Thread, Lock + +import sentry_sdk +from sentry_sdk.utils import logger + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional + + +MAX_DOWNSAMPLE_FACTOR = 10 + + +class Monitor: + """ + Performs health checks in a separate thread once every interval seconds + and updates the internal state. Other parts of the SDK only read this state + and act accordingly. + """ + + name = "sentry.monitor" + + def __init__( + self, transport: "sentry_sdk.transport.Transport", interval: float = 10 + ) -> None: + self.transport: "sentry_sdk.transport.Transport" = transport + self.interval: float = interval + + self._healthy = True + self._downsample_factor: int = 0 + + self._thread: "Optional[Thread]" = None + self._thread_lock = Lock() + self._thread_for_pid: "Optional[int]" = None + self._running = True + + def _ensure_running(self) -> None: + """ + Check that the monitor has an active thread to run in, or create one if not. + + Note that this might fail (e.g. in Python 3.12 it's not possible to + spawn new threads at interpreter shutdown). In that case self._running + will be False after running this function. + """ + if self._thread_for_pid == os.getpid() and self._thread is not None: + return None + + with self._thread_lock: + if self._thread_for_pid == os.getpid() and self._thread is not None: + return None + + def _thread() -> None: + while self._running: + time.sleep(self.interval) + if self._running: + self.run() + + thread = Thread(name=self.name, target=_thread) + thread.daemon = True + try: + thread.start() + except RuntimeError: + # Unfortunately at this point the interpreter is in a state that no + # longer allows us to spawn a thread and we have to bail. + self._running = False + return None + + self._thread = thread + self._thread_for_pid = os.getpid() + + return None + + def run(self) -> None: + self.check_health() + self.set_downsample_factor() + + def set_downsample_factor(self) -> None: + if self._healthy: + if self._downsample_factor > 0: + logger.debug( + "[Monitor] health check positive, reverting to normal sampling" + ) + self._downsample_factor = 0 + else: + if self.downsample_factor < MAX_DOWNSAMPLE_FACTOR: + self._downsample_factor += 1 + logger.debug( + "[Monitor] health check negative, downsampling with a factor of %d", + self._downsample_factor, + ) + + def check_health(self) -> None: + """ + Perform the actual health checks, + currently only checks if the transport is rate-limited. + TODO: augment in the future with more checks. + """ + self._healthy = self.transport.is_healthy() + + def is_healthy(self) -> bool: + self._ensure_running() + return self._healthy + + @property + def downsample_factor(self) -> int: + self._ensure_running() + return self._downsample_factor + + def kill(self) -> None: + self._running = False diff --git a/lib/python3.12/site-packages/sentry_sdk/py.typed b/lib/python3.12/site-packages/sentry_sdk/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sentry_sdk/scope.py b/lib/python3.12/site-packages/sentry_sdk/scope.py new file mode 100644 index 0000000000000000000000000000000000000000..3bc51c1af07e8459a4a65db7e34f34b541fcf8a0 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/scope.py @@ -0,0 +1,1980 @@ +import os +import sys +import warnings +from copy import copy, deepcopy +from collections import deque +from contextlib import contextmanager +from enum import Enum +from datetime import datetime, timezone +from functools import wraps +from itertools import chain + +import sentry_sdk +from sentry_sdk._types import AnnotatedValue +from sentry_sdk.attachments import Attachment +from sentry_sdk.consts import ( + DEFAULT_MAX_BREADCRUMBS, + FALSE_VALUES, + INSTRUMENTER, + SPANDATA, +) +from sentry_sdk.feature_flags import FlagBuffer, DEFAULT_FLAG_CAPACITY +from sentry_sdk.profiler.continuous_profiler import ( + get_profiler_id, + try_autostart_continuous_profiler, + try_profile_lifecycle_trace_start, +) +from sentry_sdk.profiler.transaction_profiler import Profile +from sentry_sdk.session import Session +from sentry_sdk.tracing_utils import ( + Baggage, + has_tracing_enabled, + has_span_streaming_enabled, + normalize_incoming_data, + PropagationContext, +) +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SENTRY_TRACE_HEADER_NAME, + NoOpSpan, + Span, + Transaction, +) +from sentry_sdk.utils import ( + capture_internal_exception, + capture_internal_exceptions, + ContextVar, + datetime_from_isoformat, + disable_capture_event, + event_from_exception, + exc_info_from_error, + format_attribute, + logger, + has_logs_enabled, + has_metrics_enabled, +) + +import typing +from typing import TYPE_CHECKING, cast + +if TYPE_CHECKING: + from collections.abc import Mapping + + from typing import Any + from typing import Callable + from typing import Deque + from typing import Dict + from typing import Generator + from typing import Iterator + from typing import List + from typing import Optional + from typing import ParamSpec + from typing import Tuple + from typing import TypeVar + from typing import Union + + from typing_extensions import Unpack + + from sentry_sdk._types import ( + Attributes, + AttributeValue, + Breadcrumb, + BreadcrumbHint, + ErrorProcessor, + Event, + EventProcessor, + ExcInfo, + Hint, + Log, + LogLevelStr, + Metric, + SamplingContext, + Type, + ) + + from sentry_sdk.tracing import TransactionKwargs + + import sentry_sdk + + P = ParamSpec("P") + R = TypeVar("R") + + F = TypeVar("F", bound=Callable[..., Any]) + T = TypeVar("T") + + +# Holds data that will be added to **all** events sent by this process. +# In case this is a http server (think web framework) with multiple users +# the data will be added to events of all users. +# Typically this is used for process wide data such as the release. +_global_scope: "Optional[Scope]" = None + +# Holds data for the active request. +# This is used to isolate data for different requests or users. +# The isolation scope is usually created by integrations, but may also +# be created manually +_isolation_scope = ContextVar("isolation_scope", default=None) + +# Holds data for the active span. +# This can be used to manually add additional data to a span. +_current_scope = ContextVar("current_scope", default=None) + +global_event_processors: "List[EventProcessor]" = [] + +# A function returning a (trace_id, span_id) tuple +# from an external tracing source (such as otel) +_external_propagation_context_fn: "Optional[Callable[[], Optional[Tuple[str, str]]]]" = None + + +class ScopeType(Enum): + CURRENT = "current" + ISOLATION = "isolation" + GLOBAL = "global" + MERGED = "merged" + + +class _ScopeManager: + def __init__(self, hub: "Optional[Any]" = None) -> None: + self._old_scopes: "List[Scope]" = [] + + def __enter__(self) -> "Scope": + isolation_scope = Scope.get_isolation_scope() + + self._old_scopes.append(isolation_scope) + + forked_scope = isolation_scope.fork() + _isolation_scope.set(forked_scope) + + return forked_scope + + def __exit__(self, exc_type: "Any", exc_value: "Any", tb: "Any") -> None: + old_scope = self._old_scopes.pop() + _isolation_scope.set(old_scope) + + +def add_global_event_processor(processor: "EventProcessor") -> None: + global_event_processors.append(processor) + + +def register_external_propagation_context( + fn: "Callable[[], Optional[Tuple[str, str]]]", +) -> None: + global _external_propagation_context_fn + _external_propagation_context_fn = fn + + +def remove_external_propagation_context() -> None: + global _external_propagation_context_fn + _external_propagation_context_fn = None + + +def get_external_propagation_context() -> "Optional[Tuple[str, str]]": + return ( + _external_propagation_context_fn() if _external_propagation_context_fn else None + ) + + +def has_external_propagation_context() -> bool: + return _external_propagation_context_fn is not None + + +def _attr_setter(fn: "Any") -> "Any": + return property(fset=fn, doc=fn.__doc__) + + +def _disable_capture(fn: "F") -> "F": + @wraps(fn) + def wrapper(self: "Any", *args: "Dict[str, Any]", **kwargs: "Any") -> "Any": + if not self._should_capture: + return + try: + self._should_capture = False + return fn(self, *args, **kwargs) + finally: + self._should_capture = True + + return wrapper # type: ignore + + +class Scope: + """The scope holds extra information that should be sent with all + events that belong to it. + """ + + # NOTE: Even though it should not happen, the scope needs to not crash when + # accessed by multiple threads. It's fine if it's full of races, but those + # races should never make the user application crash. + # + # The same needs to hold for any accesses of the scope the SDK makes. + + __slots__ = ( + "_level", + "_name", + "_fingerprint", + # note that for legacy reasons, _transaction is the transaction *name*, + # not a Transaction object (the object is stored in _span) + "_transaction", + "_transaction_info", + "_user", + "_tags", + "_contexts", + "_extras", + "_breadcrumbs", + "_n_breadcrumbs_truncated", + "_gen_ai_original_message_count", + "_gen_ai_conversation_id", + "_event_processors", + "_error_processors", + "_should_capture", + "_span", + "_session", + "_attachments", + "_force_auto_session_tracking", + "_profile", + "_propagation_context", + "client", + "_type", + "_last_event_id", + "_flags", + "_attributes", + ) + + def __init__( + self, + ty: "Optional[ScopeType]" = None, + client: "Optional[sentry_sdk.Client]" = None, + ) -> None: + self._type = ty + + self._event_processors: "List[EventProcessor]" = [] + self._error_processors: "List[ErrorProcessor]" = [] + + self._name: "Optional[str]" = None + self._propagation_context: "Optional[PropagationContext]" = None + self._n_breadcrumbs_truncated: int = 0 + self._gen_ai_original_message_count: "Dict[str, int]" = {} + + self.client: "sentry_sdk.client.BaseClient" = NonRecordingClient() + + if client is not None: + self.set_client(client) + + self.clear() + + incoming_trace_information = self._load_trace_data_from_env() + self.generate_propagation_context(incoming_data=incoming_trace_information) + + def __copy__(self) -> "Scope": + """ + Returns a copy of this scope. + This also creates a copy of all referenced data structures. + """ + rv: "Scope" = object.__new__(self.__class__) + + rv._type = self._type + rv.client = self.client + rv._level = self._level + rv._name = self._name + rv._fingerprint = self._fingerprint + rv._transaction = self._transaction + rv._transaction_info = self._transaction_info.copy() + rv._user = self._user + + rv._tags = self._tags.copy() + rv._contexts = self._contexts.copy() + rv._extras = self._extras.copy() + + rv._breadcrumbs = copy(self._breadcrumbs) + rv._n_breadcrumbs_truncated = self._n_breadcrumbs_truncated + rv._gen_ai_original_message_count = self._gen_ai_original_message_count.copy() + rv._event_processors = self._event_processors.copy() + rv._error_processors = self._error_processors.copy() + rv._propagation_context = self._propagation_context + + rv._should_capture = self._should_capture + rv._span = self._span + rv._session = self._session + rv._force_auto_session_tracking = self._force_auto_session_tracking + rv._attachments = self._attachments.copy() + + rv._profile = self._profile + + rv._last_event_id = self._last_event_id + + rv._flags = deepcopy(self._flags) + + rv._attributes = self._attributes.copy() + + rv._gen_ai_conversation_id = self._gen_ai_conversation_id + + return rv + + @classmethod + def get_current_scope(cls) -> "Scope": + """ + .. versionadded:: 2.0.0 + + Returns the current scope. + """ + current_scope = _current_scope.get() + if current_scope is None: + current_scope = Scope(ty=ScopeType.CURRENT) + _current_scope.set(current_scope) + + return current_scope + + @classmethod + def set_current_scope(cls, new_current_scope: "Scope") -> None: + """ + .. versionadded:: 2.0.0 + + Sets the given scope as the new current scope overwriting the existing current scope. + :param new_current_scope: The scope to set as the new current scope. + """ + _current_scope.set(new_current_scope) + + @classmethod + def get_isolation_scope(cls) -> "Scope": + """ + .. versionadded:: 2.0.0 + + Returns the isolation scope. + """ + isolation_scope = _isolation_scope.get() + if isolation_scope is None: + isolation_scope = Scope(ty=ScopeType.ISOLATION) + _isolation_scope.set(isolation_scope) + + return isolation_scope + + @classmethod + def set_isolation_scope(cls, new_isolation_scope: "Scope") -> None: + """ + .. versionadded:: 2.0.0 + + Sets the given scope as the new isolation scope overwriting the existing isolation scope. + :param new_isolation_scope: The scope to set as the new isolation scope. + """ + _isolation_scope.set(new_isolation_scope) + + @classmethod + def get_global_scope(cls) -> "Scope": + """ + .. versionadded:: 2.0.0 + + Returns the global scope. + """ + global _global_scope + if _global_scope is None: + _global_scope = Scope(ty=ScopeType.GLOBAL) + + return _global_scope + + def set_global_attributes(self) -> None: + from sentry_sdk.client import SDK_INFO + + self.set_attribute("sentry.sdk.name", SDK_INFO["name"]) + self.set_attribute("sentry.sdk.version", SDK_INFO["version"]) + + options = sentry_sdk.get_client().options + + server_name = options.get("server_name") + if server_name: + self.set_attribute(SPANDATA.SERVER_ADDRESS, server_name) + + environment = options.get("environment") + if environment: + self.set_attribute("sentry.environment", environment) + + release = options.get("release") + if release: + self.set_attribute("sentry.release", release) + + @classmethod + def last_event_id(cls) -> "Optional[str]": + """ + .. versionadded:: 2.2.0 + + Returns event ID of the event most recently captured by the isolation scope, or None if no event + has been captured. We do not consider events that are dropped, e.g. by a before_send hook. + Transactions also are not considered events in this context. + + The event corresponding to the returned event ID is NOT guaranteed to actually be sent to Sentry; + whether the event is sent depends on the transport. The event could be sent later or not at all. + Even a sent event could fail to arrive in Sentry due to network issues, exhausted quotas, or + various other reasons. + """ + return cls.get_isolation_scope()._last_event_id + + def _merge_scopes( + self, + additional_scope: "Optional[Scope]" = None, + additional_scope_kwargs: "Optional[Dict[str, Any]]" = None, + ) -> "Scope": + """ + Merges global, isolation and current scope into a new scope and + adds the given additional scope or additional scope kwargs to it. + """ + if additional_scope and additional_scope_kwargs: + raise TypeError("cannot provide scope and kwargs") + + final_scope = copy(_global_scope) if _global_scope is not None else Scope() + final_scope._type = ScopeType.MERGED + + isolation_scope = _isolation_scope.get() + if isolation_scope is not None: + final_scope.update_from_scope(isolation_scope) + + current_scope = _current_scope.get() + if current_scope is not None: + final_scope.update_from_scope(current_scope) + + if self != current_scope and self != isolation_scope: + final_scope.update_from_scope(self) + + if additional_scope is not None: + if callable(additional_scope): + additional_scope(final_scope) + else: + final_scope.update_from_scope(additional_scope) + + elif additional_scope_kwargs: + final_scope.update_from_kwargs(**additional_scope_kwargs) + + return final_scope + + @classmethod + def get_client(cls) -> "sentry_sdk.client.BaseClient": + """ + .. versionadded:: 2.0.0 + + Returns the currently used :py:class:`sentry_sdk.Client`. + This checks the current scope, the isolation scope and the global scope for a client. + If no client is available a :py:class:`sentry_sdk.client.NonRecordingClient` is returned. + """ + current_scope = _current_scope.get() + try: + client = current_scope.client + except AttributeError: + client = None + + if client is not None and client.is_active(): + return client + + isolation_scope = _isolation_scope.get() + try: + client = isolation_scope.client + except AttributeError: + client = None + + if client is not None and client.is_active(): + return client + + try: + client = _global_scope.client # type: ignore + except AttributeError: + client = None + + if client is not None and client.is_active(): + return client + + return NonRecordingClient() + + def set_client( + self, client: "Optional[sentry_sdk.client.BaseClient]" = None + ) -> None: + """ + .. versionadded:: 2.0.0 + + Sets the client for this scope. + + :param client: The client to use in this scope. + If `None` the client of the scope will be replaced by a :py:class:`sentry_sdk.NonRecordingClient`. + + """ + if client is not None: + self.client = client + # We need a client to set the initial global attributes on the global + # scope since they mostly come from client options, so populate them + # as soon as a client is set + sentry_sdk.get_global_scope().set_global_attributes() + else: + self.client = NonRecordingClient() + + def fork(self) -> "Scope": + """ + .. versionadded:: 2.0.0 + + Returns a fork of this scope. + """ + forked_scope = copy(self) + return forked_scope + + def _load_trace_data_from_env(self) -> "Optional[Dict[str, str]]": + """ + Load Sentry trace id and baggage from environment variables. + Can be disabled by setting SENTRY_USE_ENVIRONMENT to "false". + """ + incoming_trace_information = None + + sentry_use_environment = ( + os.environ.get("SENTRY_USE_ENVIRONMENT") or "" + ).lower() + use_environment = sentry_use_environment not in FALSE_VALUES + if use_environment: + incoming_trace_information = {} + + if os.environ.get("SENTRY_TRACE"): + incoming_trace_information[SENTRY_TRACE_HEADER_NAME] = ( + os.environ.get("SENTRY_TRACE") or "" + ) + + if os.environ.get("SENTRY_BAGGAGE"): + incoming_trace_information[BAGGAGE_HEADER_NAME] = ( + os.environ.get("SENTRY_BAGGAGE") or "" + ) + + return incoming_trace_information or None + + def set_new_propagation_context(self) -> None: + """ + Creates a new propagation context and sets it as `_propagation_context`. Overwriting existing one. + """ + self._propagation_context = PropagationContext() + + def generate_propagation_context( + self, incoming_data: "Optional[Dict[str, str]]" = None + ) -> None: + """ + Makes sure the propagation context is set on the scope. + If there is `incoming_data` overwrite existing propagation context. + If there is no `incoming_data` create new propagation context, but do NOT overwrite if already existing. + """ + if incoming_data is not None: + self._propagation_context = PropagationContext.from_incoming_data( + incoming_data + ) + + # TODO-neel this below is a BIG code smell but requires a bunch of other refactoring + if self._type != ScopeType.CURRENT: + if self._propagation_context is None: + self.set_new_propagation_context() + + def get_dynamic_sampling_context(self) -> "Optional[Dict[str, str]]": + """ + Returns the Dynamic Sampling Context from the Propagation Context. + If not existing, creates a new one. + + Deprecated: Logic moved to PropagationContext, don't use directly. + """ + if self._propagation_context is None: + return None + + return self._propagation_context.dynamic_sampling_context + + def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]": + """ + Returns the Sentry "sentry-trace" header (aka the traceparent) from the + currently active span or the scopes Propagation Context. + """ + client = self.get_client() + + # If we have an active span, return traceparent from there + if has_tracing_enabled(client.options) and self.span is not None: + return self.span.to_traceparent() + + # else return traceparent from the propagation context + return self.get_active_propagation_context().to_traceparent() + + def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]": + """ + Returns the Sentry "baggage" header containing trace information from the + currently active span or the scopes Propagation Context. + """ + client = self.get_client() + + # If we have an active span, return baggage from there + if has_tracing_enabled(client.options) and self.span is not None: + return self.span.to_baggage() + + # else return baggage from the propagation context + return self.get_active_propagation_context().get_baggage() + + def get_trace_context(self) -> "Dict[str, Any]": + """ + Returns the Sentry "trace" context from the Propagation Context. + """ + if has_tracing_enabled(self.get_client().options) and self._span is not None: + return self._span.get_trace_context() + + # if we are tracing externally (otel), those values take precedence + external_propagation_context = get_external_propagation_context() + if external_propagation_context: + trace_id, span_id = external_propagation_context + return {"trace_id": trace_id, "span_id": span_id} + + propagation_context = self.get_active_propagation_context() + + return { + "trace_id": propagation_context.trace_id, + "span_id": propagation_context.span_id, + "parent_span_id": propagation_context.parent_span_id, + "dynamic_sampling_context": propagation_context.dynamic_sampling_context, + } + + def trace_propagation_meta(self, *args: "Any", **kwargs: "Any") -> str: + """ + Return meta tags which should be injected into HTML templates + to allow propagation of trace information. + """ + span = kwargs.pop("span", None) + if span is not None: + logger.warning( + "The parameter `span` in trace_propagation_meta() is deprecated and will be removed in the future." + ) + + meta = "" + + for name, content in self.iter_trace_propagation_headers(): + meta += f'' + + return meta + + def iter_headers(self) -> "Iterator[Tuple[str, str]]": + """ + Creates a generator which returns the `sentry-trace` and `baggage` headers from the Propagation Context. + Deprecated: use PropagationContext.iter_headers instead. + """ + if self._propagation_context is not None: + yield from self._propagation_context.iter_headers() + + def iter_trace_propagation_headers( + self, *args: "Any", **kwargs: "Any" + ) -> "Generator[Tuple[str, str], None, None]": + """ + Return HTTP headers which allow propagation of trace data. + + If a span is given, the trace data will taken from the span. + If no span is given, the trace data is taken from the scope. + """ + client = self.get_client() + if not client.options.get("propagate_traces"): + warnings.warn( + "The `propagate_traces` parameter is deprecated. Please use `trace_propagation_targets` instead.", + DeprecationWarning, + stacklevel=2, + ) + return + + span = kwargs.pop("span", None) + span = span or self.span + + if has_tracing_enabled(client.options) and span is not None: + for header in span.iter_headers(): + yield header + elif has_external_propagation_context(): + # when we have an external_propagation_context (otlp) + # we leave outgoing propagation to the propagator + return + else: + for header in self.get_active_propagation_context().iter_headers(): + yield header + + def get_active_propagation_context(self) -> "PropagationContext": + if self._propagation_context is not None: + return self._propagation_context + + current_scope = self.get_current_scope() + if current_scope._propagation_context is not None: + return current_scope._propagation_context + + isolation_scope = self.get_isolation_scope() + # should actually never happen, but just in case someone calls scope.clear + if isolation_scope._propagation_context is None: + isolation_scope._propagation_context = PropagationContext() + return isolation_scope._propagation_context + + def clear(self) -> None: + """Clears the entire scope.""" + self._level: "Optional[LogLevelStr]" = None + self._fingerprint: "Optional[List[str]]" = None + self._transaction: "Optional[str]" = None + self._transaction_info: "dict[str, str]" = {} + self._user: "Optional[Dict[str, Any]]" = None + + self._tags: "Dict[str, Any]" = {} + self._contexts: "Dict[str, Dict[str, Any]]" = {} + self._extras: "dict[str, Any]" = {} + self._attachments: "List[Attachment]" = [] + + self.clear_breadcrumbs() + self._should_capture: bool = True + + self._span: "Optional[Span]" = None + self._session: "Optional[Session]" = None + self._force_auto_session_tracking: "Optional[bool]" = None + + self._profile: "Optional[Profile]" = None + + self._propagation_context = None + + # self._last_event_id is only applicable to isolation scopes + self._last_event_id: "Optional[str]" = None + self._flags: "Optional[FlagBuffer]" = None + + self._attributes: "Attributes" = {} + + self._gen_ai_conversation_id: "Optional[str]" = None + + @_attr_setter + def level(self, value: "LogLevelStr") -> None: + """ + When set this overrides the level. + + .. deprecated:: 1.0.0 + Use :func:`set_level` instead. + + :param value: The level to set. + """ + logger.warning( + "Deprecated: use .set_level() instead. This will be removed in the future." + ) + + self._level = value + + def set_level(self, value: "LogLevelStr") -> None: + """ + Sets the level for the scope. + + :param value: The level to set. + """ + self._level = value + + @_attr_setter + def fingerprint(self, value: "Optional[List[str]]") -> None: + """When set this overrides the default fingerprint.""" + self._fingerprint = value + + @property + def transaction(self) -> "Any": + # would be type: () -> Optional[Transaction], see https://github.com/python/mypy/issues/3004 + """Return the transaction (root span) in the scope, if any.""" + + # there is no span/transaction on the scope + if self._span is None: + return None + + # there is an orphan span on the scope + if self._span.containing_transaction is None: + return None + + # there is either a transaction (which is its own containing + # transaction) or a non-orphan span on the scope + return self._span.containing_transaction + + @transaction.setter + def transaction(self, value: "Any") -> None: + # would be type: (Optional[str]) -> None, see https://github.com/python/mypy/issues/3004 + """When set this forces a specific transaction name to be set. + + Deprecated: use set_transaction_name instead.""" + + # XXX: the docstring above is misleading. The implementation of + # apply_to_event prefers an existing value of event.transaction over + # anything set in the scope. + # XXX: note that with the introduction of the Scope.transaction getter, + # there is a semantic and type mismatch between getter and setter. The + # getter returns a Transaction, the setter sets a transaction name. + # Without breaking version compatibility, we could make the setter set a + # transaction name or transaction (self._span) depending on the type of + # the value argument. + + logger.warning( + "Assigning to scope.transaction directly is deprecated: use scope.set_transaction_name() instead." + ) + self._transaction = value + if self._span and self._span.containing_transaction: + self._span.containing_transaction.name = value + + def set_transaction_name(self, name: str, source: "Optional[str]" = None) -> None: + """Set the transaction name and optionally the transaction source.""" + self._transaction = name + + if self._span and self._span.containing_transaction: + self._span.containing_transaction.name = name + if source: + self._span.containing_transaction.source = source + + if source: + self._transaction_info["source"] = source + + @_attr_setter + def user(self, value: "Optional[Dict[str, Any]]") -> None: + """When set a specific user is bound to the scope. Deprecated in favor of set_user.""" + warnings.warn( + "The `Scope.user` setter is deprecated in favor of `Scope.set_user()`.", + DeprecationWarning, + stacklevel=2, + ) + self.set_user(value) + + def set_user(self, value: "Optional[Dict[str, Any]]") -> None: + """Sets a user for the scope.""" + self._user = value + session = self.get_isolation_scope()._session + if session is not None: + session.update(user=value) + + @property + def span(self) -> "Optional[Span]": + """Get/set current tracing span or transaction.""" + return self._span + + @span.setter + def span(self, span: "Optional[Span]") -> None: + self._span = span + # XXX: this differs from the implementation in JS, there Scope.setSpan + # does not set Scope._transactionName. + if isinstance(span, Transaction): + transaction = span + if transaction.name: + self._transaction = transaction.name + if transaction.source: + self._transaction_info["source"] = transaction.source + + @property + def profile(self) -> "Optional[Profile]": + return self._profile + + @profile.setter + def profile(self, profile: "Optional[Profile]") -> None: + self._profile = profile + + def set_tag(self, key: str, value: "Any") -> None: + """ + Sets a tag for a key to a specific value. + + :param key: Key of the tag to set. + + :param value: Value of the tag to set. + """ + self._tags[key] = value + + def set_tags(self, tags: "Mapping[str, object]") -> None: + """Sets multiple tags at once. + + This method updates multiple tags at once. The tags are passed as a dictionary + or other mapping type. + + Calling this method is equivalent to calling `set_tag` on each key-value pair + in the mapping. If a tag key already exists in the scope, its value will be + updated. If the tag key does not exist in the scope, the key-value pair will + be added to the scope. + + This method only modifies tag keys in the `tags` mapping passed to the method. + `scope.set_tags({})` is, therefore, a no-op. + + :param tags: A mapping of tag keys to tag values to set. + """ + self._tags.update(tags) + + def remove_tag(self, key: str) -> None: + """ + Removes a specific tag. + + :param key: Key of the tag to remove. + """ + self._tags.pop(key, None) + + def set_context( + self, + key: str, + value: "Dict[str, Any]", + ) -> None: + """ + Binds a context at a certain key to a specific value. + """ + self._contexts[key] = value + + def remove_context( + self, + key: str, + ) -> None: + """Removes a context.""" + self._contexts.pop(key, None) + + def set_extra( + self, + key: str, + value: "Any", + ) -> None: + """Sets an extra key to a specific value.""" + self._extras[key] = value + + def remove_extra( + self, + key: str, + ) -> None: + """Removes a specific extra key.""" + self._extras.pop(key, None) + + def set_conversation_id(self, conversation_id: str) -> None: + """ + Sets the conversation ID for gen_ai spans. + + :param conversation_id: The conversation ID to set. + """ + self._gen_ai_conversation_id = conversation_id + + def get_conversation_id(self) -> "Optional[str]": + """ + Gets the conversation ID for gen_ai spans. + + :returns: The conversation ID, or None if not set. + """ + return self._gen_ai_conversation_id + + def remove_conversation_id(self) -> None: + """Removes the conversation ID.""" + self._gen_ai_conversation_id = None + + def clear_breadcrumbs(self) -> None: + """Clears breadcrumb buffer.""" + self._breadcrumbs: "Deque[Breadcrumb]" = deque() + self._n_breadcrumbs_truncated = 0 + + def add_attachment( + self, + bytes: "Union[None, bytes, Callable[[], bytes]]" = None, + filename: "Optional[str]" = None, + path: "Optional[str]" = None, + content_type: "Optional[str]" = None, + add_to_transactions: bool = False, + ) -> None: + """Adds an attachment to future events sent from this scope. + + The parameters are the same as for the :py:class:`sentry_sdk.attachments.Attachment` constructor. + """ + self._attachments.append( + Attachment( + bytes=bytes, + path=path, + filename=filename, + content_type=content_type, + add_to_transactions=add_to_transactions, + ) + ) + + def add_breadcrumb( + self, + crumb: "Optional[Breadcrumb]" = None, + hint: "Optional[BreadcrumbHint]" = None, + **kwargs: "Any", + ) -> None: + """ + Adds a breadcrumb. + + :param crumb: Dictionary with the data as the sentry v7/v8 protocol expects. + + :param hint: An optional value that can be used by `before_breadcrumb` + to customize the breadcrumbs that are emitted. + """ + client = self.get_client() + + if not client.is_active(): + logger.info("Dropped breadcrumb because no client bound") + return + + before_breadcrumb = client.options.get("before_breadcrumb") + max_breadcrumbs = client.options.get("max_breadcrumbs", DEFAULT_MAX_BREADCRUMBS) + + crumb: "Breadcrumb" = dict(crumb or ()) + crumb.update(kwargs) + if not crumb: + return + + hint: "Hint" = dict(hint or ()) + + if crumb.get("timestamp") is None: + crumb["timestamp"] = datetime.now(timezone.utc) + if crumb.get("type") is None: + crumb["type"] = "default" + + if before_breadcrumb is not None: + new_crumb = before_breadcrumb(crumb, hint) + else: + new_crumb = crumb + + if new_crumb is not None: + self._breadcrumbs.append(new_crumb) + else: + logger.info("before breadcrumb dropped breadcrumb (%s)", crumb) + + while len(self._breadcrumbs) > max_breadcrumbs: + self._breadcrumbs.popleft() + self._n_breadcrumbs_truncated += 1 + + def start_transaction( + self, + transaction: "Optional[Transaction]" = None, + instrumenter: str = INSTRUMENTER.SENTRY, + custom_sampling_context: "Optional[SamplingContext]" = None, + **kwargs: "Unpack[TransactionKwargs]", + ) -> "Union[Transaction, NoOpSpan]": + """ + Start and return a transaction. + + Start an existing transaction if given, otherwise create and start a new + transaction with kwargs. + + This is the entry point to manual tracing instrumentation. + + A tree structure can be built by adding child spans to the transaction, + and child spans to other spans. To start a new child span within the + transaction or any span, call the respective `.start_child()` method. + + Every child span must be finished before the transaction is finished, + otherwise the unfinished spans are discarded. + + When used as context managers, spans and transactions are automatically + finished at the end of the `with` block. If not using context managers, + call the `.finish()` method. + + When the transaction is finished, it will be sent to Sentry with all its + finished child spans. + + :param transaction: The transaction to start. If omitted, we create and + start a new transaction. + :param instrumenter: This parameter is meant for internal use only. It + will be removed in the next major version. + :param custom_sampling_context: The transaction's custom sampling context. + :param kwargs: Optional keyword arguments to be passed to the Transaction + constructor. See :py:class:`sentry_sdk.tracing.Transaction` for + available arguments. + """ + kwargs.setdefault("scope", self) + + client = self.get_client() + + configuration_instrumenter = client.options["instrumenter"] + + if instrumenter != configuration_instrumenter: + return NoOpSpan() + + try_autostart_continuous_profiler() + + custom_sampling_context = custom_sampling_context or {} + + # kwargs at this point has type TransactionKwargs, since we have removed + # the client and custom_sampling_context from it. + transaction_kwargs: "TransactionKwargs" = kwargs + + # if we haven't been given a transaction, make one + if transaction is None: + transaction = Transaction(**transaction_kwargs) + + # use traces_sample_rate, traces_sampler, and/or inheritance to make a + # sampling decision + sampling_context = { + "transaction_context": transaction.to_json(), + "parent_sampled": transaction.parent_sampled, + } + sampling_context.update(custom_sampling_context) + transaction._set_initial_sampling_decision(sampling_context=sampling_context) + + # update the sample rate in the dsc + if transaction.sample_rate is not None: + propagation_context = self.get_active_propagation_context() + baggage = propagation_context.baggage + + if baggage is not None: + baggage.sentry_items["sample_rate"] = str(transaction.sample_rate) + + if transaction._baggage: + transaction._baggage.sentry_items["sample_rate"] = str( + transaction.sample_rate + ) + + if transaction.sampled: + profile = Profile( + transaction.sampled, transaction._start_timestamp_monotonic_ns + ) + profile._set_initial_sampling_decision(sampling_context=sampling_context) + + transaction._profile = profile + + transaction._continuous_profile = try_profile_lifecycle_trace_start() + + # Typically, the profiler is set when the transaction is created. But when + # using the auto lifecycle, the profiler isn't running when the first + # transaction is started. So make sure we update the profiler id on it. + if transaction._continuous_profile is not None: + transaction.set_profiler_id(get_profiler_id()) + + # we don't bother to keep spans if we already know we're not going to + # send the transaction + max_spans = (client.options["_experiments"].get("max_spans")) or 1000 + transaction.init_span_recorder(maxlen=max_spans) + + return transaction + + def start_span( + self, instrumenter: str = INSTRUMENTER.SENTRY, **kwargs: "Any" + ) -> "Span": + """ + Start a span whose parent is the currently active span or transaction, if any. + + The return value is a :py:class:`sentry_sdk.tracing.Span` instance, + typically used as a context manager to start and stop timing in a `with` + block. + + Only spans contained in a transaction are sent to Sentry. Most + integrations start a transaction at the appropriate time, for example + for every incoming HTTP request. Use + :py:meth:`sentry_sdk.start_transaction` to start a new transaction when + one is not already in progress. + + For supported `**kwargs` see :py:class:`sentry_sdk.tracing.Span`. + + The instrumenter parameter is deprecated for user code, and it will + be removed in the next major version. Going forward, it should only + be used by the SDK itself. + """ + if kwargs.get("description") is not None: + warnings.warn( + "The `description` parameter is deprecated. Please use `name` instead.", + DeprecationWarning, + stacklevel=2, + ) + + with new_scope(): + kwargs.setdefault("scope", self) + + client = self.get_client() + + configuration_instrumenter = client.options["instrumenter"] + + if instrumenter != configuration_instrumenter: + return NoOpSpan() + + # get current span or transaction + span = self.span or self.get_isolation_scope().span + + if span is None: + # New spans get the `trace_id` from the scope + if "trace_id" not in kwargs: + propagation_context = self.get_active_propagation_context() + kwargs["trace_id"] = propagation_context.trace_id + + span = Span(**kwargs) + else: + # Children take `trace_id`` from the parent span. + span = span.start_child(**kwargs) + + return span + + def continue_trace( + self, + environ_or_headers: "Dict[str, Any]", + op: "Optional[str]" = None, + name: "Optional[str]" = None, + source: "Optional[str]" = None, + origin: str = "manual", + ) -> "Transaction": + """ + Sets the propagation context from environment or headers and returns a transaction. + """ + self.generate_propagation_context(environ_or_headers) + + # generate_propagation_context ensures that the propagation_context is not None. + propagation_context = cast(PropagationContext, self._propagation_context) + + optional_kwargs = {} + if name: + optional_kwargs["name"] = name + if source: + optional_kwargs["source"] = source + + return Transaction( + op=op, + origin=origin, + baggage=propagation_context.baggage, + parent_sampled=propagation_context.parent_sampled, + trace_id=propagation_context.trace_id, + parent_span_id=propagation_context.parent_span_id, + same_process_as_parent=False, + **optional_kwargs, + ) + + def capture_event( + self, + event: "Event", + hint: "Optional[Hint]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """ + Captures an event. + + Merges given scope data and calls :py:meth:`sentry_sdk.client._Client.capture_event`. + + :param event: A ready-made event that can be directly sent to Sentry. + + :param hint: Contains metadata about the event that can be read from `before_send`, such as the original exception object or a HTTP request object. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`). + """ + if disable_capture_event.get(False): + return None + + scope = self._merge_scopes(scope, scope_kwargs) + + event_id = self.get_client().capture_event(event=event, hint=hint, scope=scope) + + if event_id is not None and event.get("type") != "transaction": + self.get_isolation_scope()._last_event_id = event_id + + return event_id + + def _capture_log(self, log: "Optional[Log]") -> None: + if log is None: + return + + client = self.get_client() + if not has_logs_enabled(client.options): + return + + merged_scope = self._merge_scopes() + + debug = client.options.get("debug", False) + if debug: + logger.debug( + f"[Sentry Logs] [{log.get('severity_text')}] {log.get('body')}" + ) + + client._capture_log(log, scope=merged_scope) + + def _capture_metric(self, metric: "Optional[Metric]") -> None: + if metric is None: + return + + client = self.get_client() + if not has_metrics_enabled(client.options): + return + + merged_scope = self._merge_scopes() + + debug = client.options.get("debug", False) + if debug: + logger.debug( + f"[Sentry Metrics] [{metric.get('type')}] {metric.get('name')}: {metric.get('value')}" + ) + + client._capture_metric(metric, scope=merged_scope) + + def _capture_span(self, span: "Optional[StreamedSpan]") -> None: + if span is None: + return + + client = self.get_client() + if not has_span_streaming_enabled(client.options): + return + + merged_scope = self._merge_scopes() + client._capture_span(span, scope=merged_scope) + + def capture_message( + self, + message: str, + level: "Optional[LogLevelStr]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """ + Captures a message. + + :param message: The string to send as the message. + + :param level: If no level is provided, the default level is `info`. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`). + """ + if disable_capture_event.get(False): + return None + + if level is None: + level = "info" + + event: "Event" = { + "message": message, + "level": level, + } + + return self.capture_event(event, scope=scope, **scope_kwargs) + + def capture_exception( + self, + error: "Optional[Union[BaseException, ExcInfo]]" = None, + scope: "Optional[Scope]" = None, + **scope_kwargs: "Any", + ) -> "Optional[str]": + """Captures an exception. + + :param error: An exception to capture. If `None`, `sys.exc_info()` will be used. + + :param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :param scope_kwargs: Optional data to apply to event. + For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`. + The `scope` and `scope_kwargs` parameters are mutually exclusive. + + :returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`). + """ + if disable_capture_event.get(False): + return None + + if error is not None: + exc_info = exc_info_from_error(error) + else: + exc_info = sys.exc_info() + + event, hint = event_from_exception( + exc_info, client_options=self.get_client().options + ) + + try: + return self.capture_event(event, hint=hint, scope=scope, **scope_kwargs) + except Exception: + capture_internal_exception(sys.exc_info()) + + return None + + def start_session(self, *args: "Any", **kwargs: "Any") -> None: + """Starts a new session.""" + session_mode = kwargs.pop("session_mode", "application") + + self.end_session() + + client = self.get_client() + self._session = Session( + release=client.options.get("release"), + environment=client.options.get("environment"), + user=self._user, + session_mode=session_mode, + ) + + def end_session(self, *args: "Any", **kwargs: "Any") -> None: + """Ends the current session if there is one.""" + session = self._session + self._session = None + + if session is not None: + session.close() + self.get_client().capture_session(session) + + def stop_auto_session_tracking(self, *args: "Any", **kwargs: "Any") -> None: + """Stops automatic session tracking. + + This temporarily session tracking for the current scope when called. + To resume session tracking call `resume_auto_session_tracking`. + """ + self.end_session() + self._force_auto_session_tracking = False + + def resume_auto_session_tracking(self) -> None: + """Resumes automatic session tracking for the current scope if + disabled earlier. This requires that generally automatic session + tracking is enabled. + """ + self._force_auto_session_tracking = None + + def add_event_processor( + self, + func: "EventProcessor", + ) -> None: + """Register a scope local event processor on the scope. + + :param func: This function behaves like `before_send.` + """ + if len(self._event_processors) > 20: + logger.warning( + "Too many event processors on scope! Clearing list to free up some memory: %r", + self._event_processors, + ) + del self._event_processors[:] + + self._event_processors.append(func) + + def add_error_processor( + self, + func: "ErrorProcessor", + cls: "Optional[Type[BaseException]]" = None, + ) -> None: + """Register a scope local error processor on the scope. + + :param func: A callback that works similar to an event processor but is invoked with the original exception info triple as second argument. + + :param cls: Optionally, only process exceptions of this type. + """ + if cls is not None: + cls_ = cls # For mypy. + real_func = func + + def func(event: "Event", exc_info: "ExcInfo") -> "Optional[Event]": + try: + is_inst = isinstance(exc_info[1], cls_) + except Exception: + is_inst = False + if is_inst: + return real_func(event, exc_info) + return event + + self._error_processors.append(func) + + def _apply_level_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if self._level is not None: + event["level"] = self._level + + def _apply_breadcrumbs_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + event.setdefault("breadcrumbs", {}) + + # This check is just for mypy - + if not isinstance(event["breadcrumbs"], AnnotatedValue): + event["breadcrumbs"].setdefault("values", []) + event["breadcrumbs"]["values"].extend(self._breadcrumbs) + + # Attempt to sort timestamps + try: + if not isinstance(event["breadcrumbs"], AnnotatedValue): + for crumb in event["breadcrumbs"]["values"]: + if isinstance(crumb["timestamp"], str): + crumb["timestamp"] = datetime_from_isoformat(crumb["timestamp"]) + + event["breadcrumbs"]["values"].sort( + key=lambda crumb: crumb["timestamp"] + ) + except Exception as err: + logger.debug("Error when sorting breadcrumbs", exc_info=err) + pass + + def _apply_user_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if event.get("user") is None and self._user is not None: + event["user"] = self._user + + def _apply_transaction_name_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if event.get("transaction") is None and self._transaction is not None: + event["transaction"] = self._transaction + + def _apply_transaction_info_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if event.get("transaction_info") is None and self._transaction_info is not None: + event["transaction_info"] = self._transaction_info + + def _apply_fingerprint_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if event.get("fingerprint") is None and self._fingerprint is not None: + event["fingerprint"] = self._fingerprint + + def _apply_extra_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if self._extras: + event.setdefault("extra", {}).update(self._extras) + + def _apply_tags_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if self._tags: + event.setdefault("tags", {}).update(self._tags) + + def _apply_contexts_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + if self._contexts: + event.setdefault("contexts", {}).update(self._contexts) + + contexts = event.setdefault("contexts", {}) + + # Add "trace" context + if contexts.get("trace") is None: + contexts["trace"] = self.get_trace_context() + + def _apply_flags_to_event( + self, event: "Event", hint: "Hint", options: "Optional[Dict[str, Any]]" + ) -> None: + flags = self.flags.get() + if len(flags) > 0: + event.setdefault("contexts", {}).setdefault("flags", {}).update( + {"values": flags} + ) + + def _apply_scope_attributes_to_telemetry( + self, telemetry: "Union[Log, Metric, StreamedSpan]" + ) -> None: + # TODO: turn Logs, Metrics into actual classes + if isinstance(telemetry, dict): + attributes = telemetry["attributes"] + else: + attributes = telemetry._attributes + + for attribute, value in self._attributes.items(): + if attribute not in attributes: + attributes[attribute] = value + + def _apply_user_attributes_to_telemetry( + self, telemetry: "Union[Log, Metric, StreamedSpan]" + ) -> None: + if isinstance(telemetry, dict): + attributes = telemetry["attributes"] + else: + attributes = telemetry._attributes + + if not should_send_default_pii() or self._user is None: + return + + for attribute_name, user_attribute in ( + ("user.id", "id"), + ("user.name", "username"), + ("user.email", "email"), + ): + if user_attribute in self._user and attribute_name not in attributes: + attributes[attribute_name] = self._user[user_attribute] + + def _drop(self, cause: "Any", ty: str) -> "Optional[Any]": + logger.info("%s (%s) dropped event", ty, cause) + return None + + def run_error_processors(self, event: "Event", hint: "Hint") -> "Optional[Event]": + """ + Runs the error processors on the event and returns the modified event. + """ + exc_info = hint.get("exc_info") + if exc_info is not None: + error_processors = chain( + self.get_global_scope()._error_processors, + self.get_isolation_scope()._error_processors, + self.get_current_scope()._error_processors, + ) + + for error_processor in error_processors: + new_event = error_processor(event, exc_info) + if new_event is None: + return self._drop(error_processor, "error processor") + + event = new_event + + return event + + def run_event_processors(self, event: "Event", hint: "Hint") -> "Optional[Event]": + """ + Runs the event processors on the event and returns the modified event. + """ + ty = event.get("type") + is_check_in = ty == "check_in" + + if not is_check_in: + # Get scopes without creating them to prevent infinite recursion + isolation_scope = _isolation_scope.get() + current_scope = _current_scope.get() + + event_processors = chain( + global_event_processors, + _global_scope and _global_scope._event_processors or [], + isolation_scope and isolation_scope._event_processors or [], + current_scope and current_scope._event_processors or [], + ) + + for event_processor in event_processors: + new_event = event + with capture_internal_exceptions(): + new_event = event_processor(event, hint) + if new_event is None: + return self._drop(event_processor, "event processor") + event = new_event + + return event + + @_disable_capture + def apply_to_event( + self, + event: "Event", + hint: "Hint", + options: "Optional[Dict[str, Any]]" = None, + ) -> "Optional[Event]": + """Applies the information contained on the scope to the given event.""" + ty = event.get("type") + is_transaction = ty == "transaction" + is_check_in = ty == "check_in" + + # put all attachments into the hint. This lets callbacks play around + # with attachments. We also later pull this out of the hint when we + # create the envelope. + attachments_to_send = hint.get("attachments") or [] + for attachment in self._attachments: + if not is_transaction or attachment.add_to_transactions: + attachments_to_send.append(attachment) + hint["attachments"] = attachments_to_send + + self._apply_contexts_to_event(event, hint, options) + + if is_check_in: + # Check-ins only support the trace context, strip all others + event["contexts"] = { + "trace": event.setdefault("contexts", {}).get("trace", {}) + } + + if not is_check_in: + self._apply_level_to_event(event, hint, options) + self._apply_fingerprint_to_event(event, hint, options) + self._apply_user_to_event(event, hint, options) + self._apply_transaction_name_to_event(event, hint, options) + self._apply_transaction_info_to_event(event, hint, options) + self._apply_tags_to_event(event, hint, options) + self._apply_extra_to_event(event, hint, options) + + if not is_transaction and not is_check_in: + self._apply_breadcrumbs_to_event(event, hint, options) + self._apply_flags_to_event(event, hint, options) + + event = self.run_error_processors(event, hint) + if event is None: + return None + + event = self.run_event_processors(event, hint) + if event is None: + return None + + return event + + @_disable_capture + def apply_to_telemetry(self, telemetry: "Union[Log, Metric, StreamedSpan]") -> None: + # Attributes-based events and telemetry go through here (logs, metrics, + # spansV2) + if not isinstance(telemetry, StreamedSpan): + trace_context = self.get_trace_context() + trace_id = trace_context.get("trace_id") + if telemetry.get("trace_id") is None: + telemetry["trace_id"] = ( + trace_id or "00000000-0000-0000-0000-000000000000" + ) + span_id = trace_context.get("span_id") + if telemetry.get("span_id") is None and span_id: + telemetry["span_id"] = span_id + + self._apply_scope_attributes_to_telemetry(telemetry) + self._apply_user_attributes_to_telemetry(telemetry) + + def update_from_scope(self, scope: "Scope") -> None: + """Update the scope with another scope's data.""" + if scope._level is not None: + self._level = scope._level + if scope._fingerprint is not None: + self._fingerprint = scope._fingerprint + if scope._transaction is not None: + self._transaction = scope._transaction + if scope._transaction_info is not None: + self._transaction_info.update(scope._transaction_info) + if scope._user is not None: + self._user = scope._user + if scope._tags: + self._tags.update(scope._tags) + if scope._contexts: + self._contexts.update(scope._contexts) + if scope._extras: + self._extras.update(scope._extras) + if scope._breadcrumbs: + self._breadcrumbs.extend(scope._breadcrumbs) + if scope._n_breadcrumbs_truncated: + self._n_breadcrumbs_truncated = ( + self._n_breadcrumbs_truncated + scope._n_breadcrumbs_truncated + ) + if scope._gen_ai_original_message_count: + self._gen_ai_original_message_count.update( + scope._gen_ai_original_message_count + ) + if scope._gen_ai_conversation_id: + self._gen_ai_conversation_id = scope._gen_ai_conversation_id + if scope._span: + self._span = scope._span + if scope._attachments: + self._attachments.extend(scope._attachments) + if scope._profile: + self._profile = scope._profile + if scope._propagation_context: + self._propagation_context = scope._propagation_context + if scope._session: + self._session = scope._session + if scope._flags: + if not self._flags: + self._flags = deepcopy(scope._flags) + else: + for flag in scope._flags.get(): + self._flags.set(flag["flag"], flag["result"]) + if scope._attributes: + self._attributes.update(scope._attributes) + + def update_from_kwargs( + self, + user: "Optional[Any]" = None, + level: "Optional[LogLevelStr]" = None, + extras: "Optional[Dict[str, Any]]" = None, + contexts: "Optional[Dict[str, Dict[str, Any]]]" = None, + tags: "Optional[Dict[str, str]]" = None, + fingerprint: "Optional[List[str]]" = None, + attributes: "Optional[Attributes]" = None, + ) -> None: + """Update the scope's attributes.""" + if level is not None: + self._level = level + if user is not None: + self._user = user + if extras is not None: + self._extras.update(extras) + if contexts is not None: + self._contexts.update(contexts) + if tags is not None: + self._tags.update(tags) + if fingerprint is not None: + self._fingerprint = fingerprint + if attributes is not None: + self._attributes.update(attributes) + + def __repr__(self) -> str: + return "<%s id=%s name=%s type=%s>" % ( + self.__class__.__name__, + hex(id(self)), + self._name, + self._type, + ) + + @property + def flags(self) -> "FlagBuffer": + if self._flags is None: + max_flags = ( + self.get_client().options["_experiments"].get("max_flags") + or DEFAULT_FLAG_CAPACITY + ) + self._flags = FlagBuffer(capacity=max_flags) + return self._flags + + def set_attribute(self, attribute: str, value: "AttributeValue") -> None: + """ + Set an attribute on the scope. + + Any attributes-based telemetry (logs, metrics) captured while this scope + is active will inherit attributes set on the scope. + """ + self._attributes[attribute] = format_attribute(value) + + def remove_attribute(self, attribute: str) -> None: + """Remove an attribute if set on the scope. No-op if there is no such attribute.""" + try: + del self._attributes[attribute] + except KeyError: + pass + + +@contextmanager +def new_scope() -> "Generator[Scope, None, None]": + """ + .. versionadded:: 2.0.0 + + Context manager that forks the current scope and runs the wrapped code in it. + After the wrapped code is executed, the original scope is restored. + + Example Usage: + + .. code-block:: python + + import sentry_sdk + + with sentry_sdk.new_scope() as scope: + scope.set_tag("color", "green") + sentry_sdk.capture_message("hello") # will include `color` tag. + + sentry_sdk.capture_message("hello, again") # will NOT include `color` tag. + + """ + # fork current scope + current_scope = Scope.get_current_scope() + new_scope = current_scope.fork() + token = _current_scope.set(new_scope) + + try: + yield new_scope + + finally: + try: + # restore original scope + _current_scope.reset(token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + +@contextmanager +def use_scope(scope: "Scope") -> "Generator[Scope, None, None]": + """ + .. versionadded:: 2.0.0 + + Context manager that uses the given `scope` and runs the wrapped code in it. + After the wrapped code is executed, the original scope is restored. + + Example Usage: + Suppose the variable `scope` contains a `Scope` object, which is not currently + the active scope. + + .. code-block:: python + + import sentry_sdk + + with sentry_sdk.use_scope(scope): + scope.set_tag("color", "green") + sentry_sdk.capture_message("hello") # will include `color` tag. + + sentry_sdk.capture_message("hello, again") # will NOT include `color` tag. + + """ + # set given scope as current scope + token = _current_scope.set(scope) + + try: + yield scope + + finally: + try: + # restore original scope + _current_scope.reset(token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + +@contextmanager +def isolation_scope() -> "Generator[Scope, None, None]": + """ + .. versionadded:: 2.0.0 + + Context manager that forks the current isolation scope and runs the wrapped code in it. + The current scope is also forked to not bleed data into the existing current scope. + After the wrapped code is executed, the original scopes are restored. + + Example Usage: + + .. code-block:: python + + import sentry_sdk + + with sentry_sdk.isolation_scope() as scope: + scope.set_tag("color", "green") + sentry_sdk.capture_message("hello") # will include `color` tag. + + sentry_sdk.capture_message("hello, again") # will NOT include `color` tag. + + """ + # fork current scope + current_scope = Scope.get_current_scope() + forked_current_scope = current_scope.fork() + current_token = _current_scope.set(forked_current_scope) + + # fork isolation scope + isolation_scope = Scope.get_isolation_scope() + new_isolation_scope = isolation_scope.fork() + isolation_token = _isolation_scope.set(new_isolation_scope) + + try: + yield new_isolation_scope + + finally: + # restore original scopes + try: + _current_scope.reset(current_token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + try: + _isolation_scope.reset(isolation_token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + +@contextmanager +def use_isolation_scope(isolation_scope: "Scope") -> "Generator[Scope, None, None]": + """ + .. versionadded:: 2.0.0 + + Context manager that uses the given `isolation_scope` and runs the wrapped code in it. + The current scope is also forked to not bleed data into the existing current scope. + After the wrapped code is executed, the original scopes are restored. + + Example Usage: + + .. code-block:: python + + import sentry_sdk + + with sentry_sdk.isolation_scope() as scope: + scope.set_tag("color", "green") + sentry_sdk.capture_message("hello") # will include `color` tag. + + sentry_sdk.capture_message("hello, again") # will NOT include `color` tag. + + """ + # fork current scope + current_scope = Scope.get_current_scope() + forked_current_scope = current_scope.fork() + current_token = _current_scope.set(forked_current_scope) + + # set given scope as isolation scope + isolation_token = _isolation_scope.set(isolation_scope) + + try: + yield isolation_scope + + finally: + # restore original scopes + try: + _current_scope.reset(current_token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + try: + _isolation_scope.reset(isolation_token) + except (LookupError, ValueError): + capture_internal_exception(sys.exc_info()) + + +def should_send_default_pii() -> bool: + """Shortcut for `Scope.get_client().should_send_default_pii()`.""" + return Scope.get_client().should_send_default_pii() + + +# Circular imports +from sentry_sdk.client import NonRecordingClient + +if TYPE_CHECKING: + import sentry_sdk.client diff --git a/lib/python3.12/site-packages/sentry_sdk/scrubber.py b/lib/python3.12/site-packages/sentry_sdk/scrubber.py new file mode 100644 index 0000000000000000000000000000000000000000..2857c4edaa8af5a633d7e5e8d6de891490709cef --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/scrubber.py @@ -0,0 +1,171 @@ +from sentry_sdk.utils import ( + capture_internal_exceptions, + AnnotatedValue, + iter_event_frames, +) + +from typing import TYPE_CHECKING, cast, List, Dict + +if TYPE_CHECKING: + from sentry_sdk._types import Event + from typing import Optional + + +DEFAULT_DENYLIST = [ + # stolen from relay + "password", + "passwd", + "secret", + "api_key", + "apikey", + "auth", + "credentials", + "mysql_pwd", + "privatekey", + "private_key", + "token", + "session", + # django + "csrftoken", + "sessionid", + # wsgi + "x_csrftoken", + "x_forwarded_for", + "set_cookie", + "cookie", + "authorization", + "x_api_key", + # other common names used in the wild + "aiohttp_session", # aiohttp + "connect.sid", # Express + "csrf_token", # Pyramid + "csrf", # (this is a cookie name used in accepted answers on stack overflow) + "_csrf", # Express + "_csrf_token", # Bottle + "PHPSESSID", # PHP + "_session", # Sanic + "symfony", # Symfony + "user_session", # Vue + "_xsrf", # Tornado + "XSRF-TOKEN", # Angular, Laravel +] + +DEFAULT_PII_DENYLIST = [ + "x_forwarded_for", + "x_real_ip", + "ip_address", + "remote_addr", +] + + +class EventScrubber: + def __init__( + self, + denylist: "Optional[List[str]]" = None, + recursive: bool = False, + send_default_pii: bool = False, + pii_denylist: "Optional[List[str]]" = None, + ) -> None: + """ + A scrubber that goes through the event payload and removes sensitive data configured through denylists. + + :param denylist: A security denylist that is always scrubbed, defaults to DEFAULT_DENYLIST. + :param recursive: Whether to scrub the event payload recursively, default False. + :param send_default_pii: Whether pii is sending is on, pii fields are not scrubbed. + :param pii_denylist: The denylist to use for scrubbing when pii is not sent, defaults to DEFAULT_PII_DENYLIST. + """ + self.denylist = DEFAULT_DENYLIST.copy() if denylist is None else denylist + + if not send_default_pii: + pii_denylist = ( + DEFAULT_PII_DENYLIST.copy() if pii_denylist is None else pii_denylist + ) + self.denylist += pii_denylist + + self.denylist = [x.lower() for x in self.denylist] + self.recursive = recursive + + def scrub_list(self, lst: object) -> None: + """ + If a list is passed to this method, the method recursively searches the list and any + nested lists for any dictionaries. The method calls scrub_dict on all dictionaries + it finds. + If the parameter passed to this method is not a list, the method does nothing. + """ + if not isinstance(lst, list): + return + + for v in lst: + self.scrub_dict(v) # no-op unless v is a dict + self.scrub_list(v) # no-op unless v is a list + + def scrub_dict(self, d: object) -> None: + """ + If a dictionary is passed to this method, the method scrubs the dictionary of any + sensitive data. The method calls itself recursively on any nested dictionaries ( + including dictionaries nested in lists) if self.recursive is True. + This method does nothing if the parameter passed to it is not a dictionary. + """ + if not isinstance(d, dict): + return + + for k, v in d.items(): + # The cast is needed because mypy is not smart enough to figure out that k must be a + # string after the isinstance check. + if isinstance(k, str) and k.lower() in self.denylist: + d[k] = AnnotatedValue.substituted_because_contains_sensitive_data() + elif self.recursive: + self.scrub_dict(v) # no-op unless v is a dict + self.scrub_list(v) # no-op unless v is a list + + def scrub_request(self, event: "Event") -> None: + with capture_internal_exceptions(): + if "request" in event: + if "headers" in event["request"]: + self.scrub_dict(event["request"]["headers"]) + if "cookies" in event["request"]: + self.scrub_dict(event["request"]["cookies"]) + if "data" in event["request"]: + self.scrub_dict(event["request"]["data"]) + + def scrub_extra(self, event: "Event") -> None: + with capture_internal_exceptions(): + if "extra" in event: + self.scrub_dict(event["extra"]) + + def scrub_user(self, event: "Event") -> None: + with capture_internal_exceptions(): + if "user" in event: + self.scrub_dict(event["user"]) + + def scrub_breadcrumbs(self, event: "Event") -> None: + with capture_internal_exceptions(): + if "breadcrumbs" in event: + if ( + not isinstance(event["breadcrumbs"], AnnotatedValue) + and "values" in event["breadcrumbs"] + ): + for value in event["breadcrumbs"]["values"]: + if "data" in value: + self.scrub_dict(value["data"]) + + def scrub_frames(self, event: "Event") -> None: + with capture_internal_exceptions(): + for frame in iter_event_frames(event): + if "vars" in frame: + self.scrub_dict(frame["vars"]) + + def scrub_spans(self, event: "Event") -> None: + with capture_internal_exceptions(): + if "spans" in event: + for span in cast(List[Dict[str, object]], event["spans"]): + if "data" in span: + self.scrub_dict(span["data"]) + + def scrub_event(self, event: "Event") -> None: + self.scrub_request(event) + self.scrub_extra(event) + self.scrub_user(event) + self.scrub_breadcrumbs(event) + self.scrub_frames(event) + self.scrub_spans(event) diff --git a/lib/python3.12/site-packages/sentry_sdk/serializer.py b/lib/python3.12/site-packages/sentry_sdk/serializer.py new file mode 100644 index 0000000000000000000000000000000000000000..9725d3ab532e34e4ffb3dee51107943f6e72401f --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/serializer.py @@ -0,0 +1,390 @@ +import sys +import math +from collections.abc import Mapping, Sequence, Set +from datetime import datetime + +from sentry_sdk.utils import ( + AnnotatedValue, + capture_internal_exception, + disable_capture_event, + format_timestamp, + safe_repr, + strip_string, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from types import TracebackType + + from typing import Any + from typing import Callable + from typing import ContextManager + from typing import Dict + from typing import List + from typing import Optional + from typing import Type + from typing import Union + + from sentry_sdk._types import NotImplementedType + + Span = Dict[str, Any] + + ReprProcessor = Callable[[Any, Dict[str, Any]], Union[NotImplementedType, str]] + Segment = Union[str, int] + + +# Bytes are technically not strings in Python 3, but we can serialize them +serializable_str_types = (str, bytes, bytearray, memoryview) + + +# Maximum length of JSON-serialized event payloads that can be safely sent +# before the server may reject the event due to its size. This is not intended +# to reflect actual values defined server-side, but rather only be an upper +# bound for events sent by the SDK. +# +# Can be overwritten if wanting to send more bytes, e.g. with a custom server. +# When changing this, keep in mind that events may be a little bit larger than +# this value due to attached metadata, so keep the number conservative. +MAX_EVENT_BYTES = 10**6 + +# Maximum depth and breadth of databags. Excess data will be trimmed. If +# max_request_body_size is "always", request bodies won't be trimmed. +MAX_DATABAG_DEPTH = 5 +MAX_DATABAG_BREADTH = 10 +CYCLE_MARKER = "" + + +global_repr_processors: "List[ReprProcessor]" = [] + + +def add_global_repr_processor(processor: "ReprProcessor") -> None: + global_repr_processors.append(processor) + + +sequence_types: "List[type]" = [Sequence, Set] + + +def add_repr_sequence_type(ty: type) -> None: + sequence_types.append(ty) + + +class Memo: + __slots__ = ("_ids", "_objs") + + def __init__(self) -> None: + self._ids: "Dict[int, Any]" = {} + self._objs: "List[Any]" = [] + + def memoize(self, obj: "Any") -> "ContextManager[bool]": + self._objs.append(obj) + return self + + def __enter__(self) -> bool: + obj = self._objs[-1] + if id(obj) in self._ids: + return True + else: + self._ids[id(obj)] = obj + return False + + def __exit__( + self, + ty: "Optional[Type[BaseException]]", + value: "Optional[BaseException]", + tb: "Optional[TracebackType]", + ) -> None: + self._ids.pop(id(self._objs.pop()), None) + + +def serialize(event: "Dict[str, Any]", **kwargs: "Any") -> "Dict[str, Any]": + """ + A very smart serializer that takes a dict and emits a json-friendly dict. + Currently used for serializing the final Event and also prematurely while fetching the stack + local variables for each frame in a stacktrace. + + It works internally with 'databags' which are arbitrary data structures like Mapping, Sequence and Set. + The algorithm itself is a recursive graph walk down the data structures it encounters. + + It has the following responsibilities: + * Trimming databags and keeping them within MAX_DATABAG_BREADTH and MAX_DATABAG_DEPTH. + * Calling safe_repr() on objects appropriately to keep them informative and readable in the final payload. + * Annotating the payload with the _meta field whenever trimming happens. + + :param max_request_body_size: If set to "always", will never trim request bodies. + :param max_value_length: The max length to strip strings to, defaults to sentry_sdk.consts.DEFAULT_MAX_VALUE_LENGTH + :param is_vars: If we're serializing vars early, we want to repr() things that are JSON-serializable to make their type more apparent. For example, it's useful to see the difference between a unicode-string and a bytestring when viewing a stacktrace. + :param custom_repr: A custom repr function that runs before safe_repr on the object to be serialized. If it returns None or throws internally, we will fallback to safe_repr. + + """ + memo = Memo() + path: "List[Segment]" = [] + meta_stack: "List[Dict[str, Any]]" = [] + + keep_request_bodies: bool = kwargs.pop("max_request_body_size", None) == "always" + max_value_length: "Optional[int]" = kwargs.pop("max_value_length", None) + is_vars = kwargs.pop("is_vars", False) + custom_repr: "Callable[..., Optional[str]]" = kwargs.pop("custom_repr", None) + + def _safe_repr_wrapper(value: "Any") -> str: + try: + repr_value = None + if custom_repr is not None: + repr_value = custom_repr(value) + return repr_value or safe_repr(value) + except Exception: + return safe_repr(value) + + def _annotate(**meta: "Any") -> None: + while len(meta_stack) <= len(path): + try: + segment = path[len(meta_stack) - 1] + node = meta_stack[-1].setdefault(str(segment), {}) + except IndexError: + node = {} + + meta_stack.append(node) + + meta_stack[-1].setdefault("", {}).update(meta) + + def _is_databag() -> "Optional[bool]": + """ + A databag is any value that we need to trim. + True for stuff like vars, request bodies, breadcrumbs and extra. + + :returns: `True` for "yes", `False` for :"no", `None` for "maybe soon". + """ + try: + if is_vars: + return True + + is_request_body = _is_request_body() + if is_request_body in (True, None): + return is_request_body + + p0 = path[0] + if p0 == "breadcrumbs" and path[1] == "values": + path[2] + return True + + if p0 == "extra": + return True + + except IndexError: + return None + + return False + + def _is_span_attribute() -> "Optional[bool]": + try: + if path[0] == "spans" and path[2] == "data": + return True + except IndexError: + return None + + return False + + def _is_request_body() -> "Optional[bool]": + try: + if path[0] == "request" and path[1] == "data": + return True + except IndexError: + return None + + return False + + def _serialize_node( + obj: "Any", + is_databag: "Optional[bool]" = None, + is_request_body: "Optional[bool]" = None, + should_repr_strings: "Optional[bool]" = None, + segment: "Optional[Segment]" = None, + remaining_breadth: "Optional[Union[int, float]]" = None, + remaining_depth: "Optional[Union[int, float]]" = None, + ) -> "Any": + if segment is not None: + path.append(segment) + + try: + with memo.memoize(obj) as result: + if result: + return CYCLE_MARKER + + return _serialize_node_impl( + obj, + is_databag=is_databag, + is_request_body=is_request_body, + should_repr_strings=should_repr_strings, + remaining_depth=remaining_depth, + remaining_breadth=remaining_breadth, + ) + except BaseException: + capture_internal_exception(sys.exc_info()) + + if is_databag: + return "" + + return None + finally: + if segment is not None: + path.pop() + del meta_stack[len(path) + 1 :] + + def _flatten_annotated(obj: "Any") -> "Any": + if isinstance(obj, AnnotatedValue): + _annotate(**obj.metadata) + obj = obj.value + return obj + + def _serialize_node_impl( + obj: "Any", + is_databag: "Optional[bool]", + is_request_body: "Optional[bool]", + should_repr_strings: "Optional[bool]", + remaining_depth: "Optional[Union[float, int]]", + remaining_breadth: "Optional[Union[float, int]]", + ) -> "Any": + if isinstance(obj, AnnotatedValue): + should_repr_strings = False + if should_repr_strings is None: + should_repr_strings = is_vars + + if is_databag is None: + is_databag = _is_databag() + + if is_request_body is None: + is_request_body = _is_request_body() + + if is_databag: + if is_request_body and keep_request_bodies: + remaining_depth = float("inf") + remaining_breadth = float("inf") + else: + if remaining_depth is None: + remaining_depth = MAX_DATABAG_DEPTH + if remaining_breadth is None: + remaining_breadth = MAX_DATABAG_BREADTH + + obj = _flatten_annotated(obj) + + if remaining_depth is not None and remaining_depth <= 0: + _annotate(rem=[["!limit", "x"]]) + if is_databag: + return _flatten_annotated( + strip_string(_safe_repr_wrapper(obj), max_length=max_value_length) + ) + return None + + is_span_attribute = _is_span_attribute() + if (is_databag or is_span_attribute) and global_repr_processors: + hints = {"memo": memo, "remaining_depth": remaining_depth} + for processor in global_repr_processors: + result = processor(obj, hints) + if result is not NotImplemented: + return _flatten_annotated(result) + + sentry_repr = getattr(type(obj), "__sentry_repr__", None) + + if obj is None or isinstance(obj, (bool, int, float)): + if should_repr_strings or ( + isinstance(obj, float) and (math.isinf(obj) or math.isnan(obj)) + ): + return _safe_repr_wrapper(obj) + else: + return obj + + elif callable(sentry_repr): + return sentry_repr(obj) + + elif isinstance(obj, datetime): + return ( + str(format_timestamp(obj)) + if not should_repr_strings + else _safe_repr_wrapper(obj) + ) + + elif isinstance(obj, Mapping): + # Create temporary copy here to avoid calling too much code that + # might mutate our dictionary while we're still iterating over it. + obj = dict(obj.items()) + + rv_dict: "Dict[str, Any]" = {} + i = 0 + + for k, v in obj.items(): + if remaining_breadth is not None and i >= remaining_breadth: + _annotate(len=len(obj)) + break + + str_k = str(k) + v = _serialize_node( + v, + segment=str_k, + should_repr_strings=should_repr_strings, + is_databag=is_databag, + is_request_body=is_request_body, + remaining_depth=( + remaining_depth - 1 if remaining_depth is not None else None + ), + remaining_breadth=remaining_breadth, + ) + rv_dict[str_k] = v + i += 1 + + return rv_dict + + elif not isinstance(obj, serializable_str_types) and isinstance( + obj, tuple(sequence_types) + ): + rv_list = [] + + for i, v in enumerate(obj): + if remaining_breadth is not None and i >= remaining_breadth: + _annotate(len=len(obj)) + break + + rv_list.append( + _serialize_node( + v, + segment=i, + should_repr_strings=should_repr_strings, + is_databag=is_databag, + is_request_body=is_request_body, + remaining_depth=( + remaining_depth - 1 if remaining_depth is not None else None + ), + remaining_breadth=remaining_breadth, + ) + ) + + return rv_list + + if should_repr_strings: + obj = _safe_repr_wrapper(obj) + else: + if isinstance(obj, bytes) or isinstance(obj, bytearray): + obj = obj.decode("utf-8", "replace") + + if not isinstance(obj, str): + obj = _safe_repr_wrapper(obj) + + is_span_description = ( + len(path) == 3 and path[0] == "spans" and path[-1] == "description" + ) + if is_span_description: + return obj + + return _flatten_annotated(strip_string(obj, max_length=max_value_length)) + + # + # Start of serialize() function + # + disable_capture_event.set(True) + try: + serialized_event = _serialize_node(event, **kwargs) + if not is_vars and meta_stack and isinstance(serialized_event, dict): + serialized_event["_meta"] = meta_stack[0] + + return serialized_event + finally: + disable_capture_event.set(False) diff --git a/lib/python3.12/site-packages/sentry_sdk/session.py b/lib/python3.12/site-packages/sentry_sdk/session.py new file mode 100644 index 0000000000000000000000000000000000000000..315ba1bf9b06eff1d5083c7b3f1a61bfeeea2815 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/session.py @@ -0,0 +1,169 @@ +import uuid +from datetime import datetime, timezone + +from sentry_sdk.utils import format_timestamp + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Optional + from typing import Union + from typing import Any + from typing import Dict + + from sentry_sdk._types import SessionStatus + + +def _minute_trunc(ts: "datetime") -> "datetime": + return ts.replace(second=0, microsecond=0) + + +def _make_uuid( + val: "Union[str, uuid.UUID]", +) -> "uuid.UUID": + if isinstance(val, uuid.UUID): + return val + return uuid.UUID(val) + + +class Session: + def __init__( + self, + sid: "Optional[Union[str, uuid.UUID]]" = None, + did: "Optional[str]" = None, + timestamp: "Optional[datetime]" = None, + started: "Optional[datetime]" = None, + duration: "Optional[float]" = None, + status: "Optional[SessionStatus]" = None, + release: "Optional[str]" = None, + environment: "Optional[str]" = None, + user_agent: "Optional[str]" = None, + ip_address: "Optional[str]" = None, + errors: "Optional[int]" = None, + user: "Optional[Any]" = None, + session_mode: str = "application", + ) -> None: + if sid is None: + sid = uuid.uuid4() + if started is None: + started = datetime.now(timezone.utc) + if status is None: + status = "ok" + self.status = status + self.did: "Optional[str]" = None + self.started = started + self.release: "Optional[str]" = None + self.environment: "Optional[str]" = None + self.duration: "Optional[float]" = None + self.user_agent: "Optional[str]" = None + self.ip_address: "Optional[str]" = None + self.session_mode: str = session_mode + self.errors = 0 + + self.update( + sid=sid, + did=did, + timestamp=timestamp, + duration=duration, + release=release, + environment=environment, + user_agent=user_agent, + ip_address=ip_address, + errors=errors, + user=user, + ) + + @property + def truncated_started(self) -> "datetime": + return _minute_trunc(self.started) + + def update( + self, + sid: "Optional[Union[str, uuid.UUID]]" = None, + did: "Optional[str]" = None, + timestamp: "Optional[datetime]" = None, + started: "Optional[datetime]" = None, + duration: "Optional[float]" = None, + status: "Optional[SessionStatus]" = None, + release: "Optional[str]" = None, + environment: "Optional[str]" = None, + user_agent: "Optional[str]" = None, + ip_address: "Optional[str]" = None, + errors: "Optional[int]" = None, + user: "Optional[Any]" = None, + ) -> None: + # If a user is supplied we pull some data form it + if user: + if ip_address is None: + ip_address = user.get("ip_address") + if did is None: + did = user.get("id") or user.get("email") or user.get("username") + + if sid is not None: + self.sid = _make_uuid(sid) + if did is not None: + self.did = str(did) + if timestamp is None: + timestamp = datetime.now(timezone.utc) + self.timestamp = timestamp + if started is not None: + self.started = started + if duration is not None: + self.duration = duration + if release is not None: + self.release = release + if environment is not None: + self.environment = environment + if ip_address is not None: + self.ip_address = ip_address + if user_agent is not None: + self.user_agent = user_agent + if errors is not None: + self.errors = errors + + if status is not None: + self.status = status + + def close( + self, + status: "Optional[SessionStatus]" = None, + ) -> "Any": + if status is None and self.status == "ok": + status = "exited" + if status is not None: + self.update(status=status) + + def get_json_attrs( + self, + with_user_info: "Optional[bool]" = True, + ) -> "Any": + attrs = {} + if self.release is not None: + attrs["release"] = self.release + if self.environment is not None: + attrs["environment"] = self.environment + if with_user_info: + if self.ip_address is not None: + attrs["ip_address"] = self.ip_address + if self.user_agent is not None: + attrs["user_agent"] = self.user_agent + return attrs + + def to_json(self) -> "Any": + rv: "Dict[str, Any]" = { + "sid": str(self.sid), + "init": True, + "started": format_timestamp(self.started), + "timestamp": format_timestamp(self.timestamp), + "status": self.status, + } + if self.errors: + rv["errors"] = self.errors + if self.did is not None: + rv["did"] = self.did + if self.duration is not None: + rv["duration"] = self.duration + attrs = self.get_json_attrs() + if attrs: + rv["attrs"] = attrs + return rv diff --git a/lib/python3.12/site-packages/sentry_sdk/sessions.py b/lib/python3.12/site-packages/sentry_sdk/sessions.py new file mode 100644 index 0000000000000000000000000000000000000000..2b7ed8487d7b600e9a078cd4cea3920380f29db2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/sessions.py @@ -0,0 +1,269 @@ +import os +import warnings +from threading import Thread, Lock, Event +from contextlib import contextmanager + +import sentry_sdk +from sentry_sdk.envelope import Envelope +from sentry_sdk.session import Session +from sentry_sdk.utils import format_timestamp + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Callable + from typing import Dict + from typing import Generator + from typing import List + from typing import Optional + from typing import Union + + +def is_auto_session_tracking_enabled( + hub: "Optional[sentry_sdk.Hub]" = None, +) -> "Union[Any, bool, None]": + """DEPRECATED: Utility function to find out if session tracking is enabled.""" + + # Internal callers should use private _is_auto_session_tracking_enabled, instead. + warnings.warn( + "This function is deprecated and will be removed in the next major release. " + "There is no public API replacement.", + DeprecationWarning, + stacklevel=2, + ) + + if hub is None: + hub = sentry_sdk.Hub.current + + should_track = hub.scope._force_auto_session_tracking + + if should_track is None: + client_options = hub.client.options if hub.client else {} + should_track = client_options.get("auto_session_tracking", False) + + return should_track + + +@contextmanager +def auto_session_tracking( + hub: "Optional[sentry_sdk.Hub]" = None, session_mode: str = "application" +) -> "Generator[None, None, None]": + """DEPRECATED: Use track_session instead + Starts and stops a session automatically around a block. + """ + warnings.warn( + "This function is deprecated and will be removed in the next major release. " + "Use track_session instead.", + DeprecationWarning, + stacklevel=2, + ) + + if hub is None: + hub = sentry_sdk.Hub.current + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + should_track = is_auto_session_tracking_enabled(hub) + if should_track: + hub.start_session(session_mode=session_mode) + try: + yield + finally: + if should_track: + hub.end_session() + + +def is_auto_session_tracking_enabled_scope(scope: "sentry_sdk.Scope") -> bool: + """ + DEPRECATED: Utility function to find out if session tracking is enabled. + """ + + warnings.warn( + "This function is deprecated and will be removed in the next major release. " + "There is no public API replacement.", + DeprecationWarning, + stacklevel=2, + ) + + # Internal callers should use private _is_auto_session_tracking_enabled, instead. + return _is_auto_session_tracking_enabled(scope) + + +def _is_auto_session_tracking_enabled(scope: "sentry_sdk.Scope") -> bool: + """ + Utility function to find out if session tracking is enabled. + """ + + should_track = scope._force_auto_session_tracking + if should_track is None: + client_options = sentry_sdk.get_client().options + should_track = client_options.get("auto_session_tracking", False) + + return should_track + + +@contextmanager +def auto_session_tracking_scope( + scope: "sentry_sdk.Scope", session_mode: str = "application" +) -> "Generator[None, None, None]": + """DEPRECATED: This function is a deprecated alias for track_session. + Starts and stops a session automatically around a block. + """ + + warnings.warn( + "This function is a deprecated alias for track_session and will be removed in the next major release.", + DeprecationWarning, + stacklevel=2, + ) + + with track_session(scope, session_mode=session_mode): + yield + + +@contextmanager +def track_session( + scope: "sentry_sdk.Scope", session_mode: str = "application" +) -> "Generator[None, None, None]": + """ + Start a new session in the provided scope, assuming session tracking is enabled. + This is a no-op context manager if session tracking is not enabled. + """ + + should_track = _is_auto_session_tracking_enabled(scope) + if should_track: + scope.start_session(session_mode=session_mode) + try: + yield + finally: + if should_track: + scope.end_session() + + +TERMINAL_SESSION_STATES = ("exited", "abnormal", "crashed") +MAX_ENVELOPE_ITEMS = 100 + + +def make_aggregate_envelope(aggregate_states: "Any", attrs: "Any") -> "Any": + return {"attrs": dict(attrs), "aggregates": list(aggregate_states.values())} + + +class SessionFlusher: + def __init__( + self, + capture_func: "Callable[[Envelope], None]", + flush_interval: int = 60, + ) -> None: + self.capture_func = capture_func + self.flush_interval = flush_interval + self.pending_sessions: "List[Any]" = [] + self.pending_aggregates: "Dict[Any, Any]" = {} + self._thread: "Optional[Thread]" = None + self._thread_lock = Lock() + self._aggregate_lock = Lock() + self._thread_for_pid: "Optional[int]" = None + self.__shutdown_requested = Event() + + def flush(self) -> None: + pending_sessions = self.pending_sessions + self.pending_sessions = [] + + with self._aggregate_lock: + pending_aggregates = self.pending_aggregates + self.pending_aggregates = {} + + envelope = Envelope() + for session in pending_sessions: + if len(envelope.items) == MAX_ENVELOPE_ITEMS: + self.capture_func(envelope) + envelope = Envelope() + + envelope.add_session(session) + + for attrs, states in pending_aggregates.items(): + if len(envelope.items) == MAX_ENVELOPE_ITEMS: + self.capture_func(envelope) + envelope = Envelope() + + envelope.add_sessions(make_aggregate_envelope(states, attrs)) + + if len(envelope.items) > 0: + self.capture_func(envelope) + + def _ensure_running(self) -> None: + """ + Check that we have an active thread to run in, or create one if not. + + Note that this might fail (e.g. in Python 3.12 it's not possible to + spawn new threads at interpreter shutdown). In that case self._running + will be False after running this function. + """ + if self._thread_for_pid == os.getpid() and self._thread is not None: + return None + with self._thread_lock: + if self._thread_for_pid == os.getpid() and self._thread is not None: + return None + + def _thread() -> None: + running = True + while running: + running = not self.__shutdown_requested.wait(self.flush_interval) + self.flush() + + thread = Thread(target=_thread) + thread.daemon = True + try: + thread.start() + except RuntimeError: + # Unfortunately at this point the interpreter is in a state that no + # longer allows us to spawn a thread and we have to bail. + self.__shutdown_requested.set() + return None + + self._thread = thread + self._thread_for_pid = os.getpid() + + return None + + def add_aggregate_session( + self, + session: "Session", + ) -> None: + # NOTE on `session.did`: + # the protocol can deal with buckets that have a distinct-id, however + # in practice we expect the python SDK to have an extremely high cardinality + # here, effectively making aggregation useless, therefore we do not + # aggregate per-did. + + # For this part we can get away with using the global interpreter lock + with self._aggregate_lock: + attrs = session.get_json_attrs(with_user_info=False) + primary_key = tuple(sorted(attrs.items())) + secondary_key = session.truncated_started # (, session.did) + states = self.pending_aggregates.setdefault(primary_key, {}) + state = states.setdefault(secondary_key, {}) + + if "started" not in state: + state["started"] = format_timestamp(session.truncated_started) + # if session.did is not None: + # state["did"] = session.did + if session.status == "crashed": + state["crashed"] = state.get("crashed", 0) + 1 + elif session.status == "abnormal": + state["abnormal"] = state.get("abnormal", 0) + 1 + elif session.errors > 0: + state["errored"] = state.get("errored", 0) + 1 + else: + state["exited"] = state.get("exited", 0) + 1 + + def add_session( + self, + session: "Session", + ) -> None: + if session.session_mode == "request": + self.add_aggregate_session(session) + else: + self.pending_sessions.append(session.to_json()) + self._ensure_running() + + def kill(self) -> None: + self.__shutdown_requested.set() diff --git a/lib/python3.12/site-packages/sentry_sdk/spotlight.py b/lib/python3.12/site-packages/sentry_sdk/spotlight.py new file mode 100644 index 0000000000000000000000000000000000000000..f70ea9d34198abebd614fdd9984b1cf8cf0cc895 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/spotlight.py @@ -0,0 +1,331 @@ +import io +import logging +import os +import time +import urllib.parse +import urllib.request +import urllib.error +import urllib3 +import sys + +from itertools import chain, product + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Callable + from typing import Dict + from typing import Optional + from typing import Self + +from sentry_sdk.utils import ( + logger as sentry_logger, + env_to_bool, + capture_internal_exceptions, +) +from sentry_sdk.envelope import Envelope + + +logger = logging.getLogger("spotlight") + + +DEFAULT_SPOTLIGHT_URL = "http://localhost:8969/stream" +DJANGO_SPOTLIGHT_MIDDLEWARE_PATH = "sentry_sdk.spotlight.SpotlightMiddleware" + + +class SpotlightClient: + """ + A client for sending envelopes to Sentry Spotlight. + + Implements exponential backoff retry logic per the SDK spec: + - Logs error at least once when server is unreachable + - Does not log for every failed envelope + - Uses exponential backoff to avoid hammering an unavailable server + - Never blocks normal Sentry operation + """ + + # Exponential backoff settings + INITIAL_RETRY_DELAY = 1.0 # Start with 1 second + MAX_RETRY_DELAY = 60.0 # Max 60 seconds + + def __init__(self, url: str) -> None: + self.url = url + self.http = urllib3.PoolManager() + self._retry_delay = self.INITIAL_RETRY_DELAY + self._last_error_time: float = 0.0 + + def capture_envelope(self, envelope: "Envelope") -> None: + # Check if we're in backoff period - skip sending to avoid blocking + if self._last_error_time > 0: + time_since_error = time.time() - self._last_error_time + if time_since_error < self._retry_delay: + # Still in backoff period, skip this envelope + return + + body = io.BytesIO() + envelope.serialize_into(body) + try: + req = self.http.request( + url=self.url, + body=body.getvalue(), + method="POST", + headers={ + "Content-Type": "application/x-sentry-envelope", + }, + ) + req.close() + # Success - reset backoff state + self._retry_delay = self.INITIAL_RETRY_DELAY + self._last_error_time = 0.0 + except Exception as e: + self._last_error_time = time.time() + + # Increase backoff delay exponentially first, so logged value matches actual wait + self._retry_delay = min(self._retry_delay * 2, self.MAX_RETRY_DELAY) + + # Log error once per backoff cycle (we skip sends during backoff, so only one failure per cycle) + sentry_logger.warning( + "Failed to send envelope to Spotlight at %s: %s. " + "Will retry after %.1f seconds.", + self.url, + e, + self._retry_delay, + ) + + +try: + from django.utils.deprecation import MiddlewareMixin + from django.http import HttpResponseServerError, HttpResponse, HttpRequest + from django.conf import settings + + SPOTLIGHT_JS_ENTRY_PATH = "/assets/main.js" + SPOTLIGHT_JS_SNIPPET_PATTERN = ( + "\n" + '\n' + ) + SPOTLIGHT_ERROR_PAGE_SNIPPET = ( + '\n' + '\n' + ) + CHARSET_PREFIX = "charset=" + BODY_TAG_NAME = "body" + BODY_CLOSE_TAG_POSSIBILITIES = tuple( + "".format("".join(chars)) + for chars in product(*zip(BODY_TAG_NAME.upper(), BODY_TAG_NAME.lower())) + ) + + class SpotlightMiddleware(MiddlewareMixin): # type: ignore[misc] + _spotlight_script: "Optional[str]" = None + _spotlight_url: "Optional[str]" = None + + def __init__(self: "Self", get_response: "Callable[..., HttpResponse]") -> None: + super().__init__(get_response) + + import sentry_sdk.api + + self.sentry_sdk = sentry_sdk.api + + spotlight_client = self.sentry_sdk.get_client().spotlight + if spotlight_client is None: + sentry_logger.warning( + "Cannot find Spotlight client from SpotlightMiddleware, disabling the middleware." + ) + return None + # Spotlight URL has a trailing `/stream` part at the end so split it off + self._spotlight_url = urllib.parse.urljoin(spotlight_client.url, "../") + + @property + def spotlight_script(self: "Self") -> "Optional[str]": + if self._spotlight_url is not None and self._spotlight_script is None: + try: + spotlight_js_url = urllib.parse.urljoin( + self._spotlight_url, SPOTLIGHT_JS_ENTRY_PATH + ) + req = urllib.request.Request( + spotlight_js_url, + method="HEAD", + ) + urllib.request.urlopen(req) + self._spotlight_script = SPOTLIGHT_JS_SNIPPET_PATTERN.format( + spotlight_url=self._spotlight_url, + spotlight_js_url=spotlight_js_url, + ) + except urllib.error.URLError as err: + sentry_logger.debug( + "Cannot get Spotlight JS to inject at %s. SpotlightMiddleware will not be very useful.", + spotlight_js_url, + exc_info=err, + ) + + return self._spotlight_script + + def process_response( + self: "Self", _request: "HttpRequest", response: "HttpResponse" + ) -> "Optional[HttpResponse]": + content_type_header = tuple( + p.strip() + for p in response.headers.get("Content-Type", "").lower().split(";") + ) + content_type = content_type_header[0] + if len(content_type_header) > 1 and content_type_header[1].startswith( + CHARSET_PREFIX + ): + encoding = content_type_header[1][len(CHARSET_PREFIX) :] + else: + encoding = "utf-8" + + if ( + self.spotlight_script is not None + and not response.streaming + and content_type == "text/html" + ): + content_length = len(response.content) + injection = self.spotlight_script.encode(encoding) + injection_site = next( + ( + idx + for idx in ( + response.content.rfind(body_variant.encode(encoding)) + for body_variant in BODY_CLOSE_TAG_POSSIBILITIES + ) + if idx > -1 + ), + content_length, + ) + + # This approach works even when we don't have a `` tag + response.content = ( + response.content[:injection_site] + + injection + + response.content[injection_site:] + ) + + if response.has_header("Content-Length"): + response.headers["Content-Length"] = content_length + len(injection) + + return response + + def process_exception( + self: "Self", _request: "HttpRequest", exception: Exception + ) -> "Optional[HttpResponseServerError]": + if not settings.DEBUG or not self._spotlight_url: + return None + + try: + spotlight = ( + urllib.request.urlopen(self._spotlight_url).read().decode("utf-8") + ) + except urllib.error.URLError: + return None + else: + event_id = self.sentry_sdk.capture_exception(exception) + return HttpResponseServerError( + spotlight.replace( + "", + SPOTLIGHT_ERROR_PAGE_SNIPPET.format( + spotlight_url=self._spotlight_url, event_id=event_id + ), + ) + ) + +except ImportError: + settings = None + + +def _resolve_spotlight_url( + spotlight_config: "Any", sentry_logger: "Any" +) -> "Optional[str]": + """ + Resolve the Spotlight URL based on config and environment variable. + + Implements precedence rules per the SDK spec: + https://develop.sentry.dev/sdk/expected-features/spotlight/ + + Returns the resolved URL string, or None if Spotlight should be disabled. + """ + spotlight_env_value = os.environ.get("SENTRY_SPOTLIGHT") + + # Parse env var to determine if it's a boolean or URL + spotlight_from_env: "Optional[bool]" = None + spotlight_env_url: "Optional[str]" = None + if spotlight_env_value: + parsed = env_to_bool(spotlight_env_value, strict=True) + if parsed is None: + # It's a URL string + spotlight_from_env = True + spotlight_env_url = spotlight_env_value + else: + spotlight_from_env = parsed + + # Apply precedence rules per spec: + # https://develop.sentry.dev/sdk/expected-features/spotlight/#precedence-rules + if spotlight_config is False: + # Config explicitly disables spotlight - warn if env var was set + if spotlight_from_env: + sentry_logger.warning( + "Spotlight is disabled via spotlight=False config option, " + "ignoring SENTRY_SPOTLIGHT environment variable." + ) + return None + elif spotlight_config is True: + # Config enables spotlight with boolean true + # If env var has URL, use env var URL per spec + if spotlight_env_url: + return spotlight_env_url + else: + return DEFAULT_SPOTLIGHT_URL + elif isinstance(spotlight_config, str): + # Config has URL string - use config URL, warn if env var differs + if spotlight_env_value and spotlight_env_value != spotlight_config: + sentry_logger.warning( + "Spotlight URL from config (%s) takes precedence over " + "SENTRY_SPOTLIGHT environment variable (%s).", + spotlight_config, + spotlight_env_value, + ) + return spotlight_config + elif spotlight_config is None: + # No config - use env var + if spotlight_env_url: + return spotlight_env_url + elif spotlight_from_env: + return DEFAULT_SPOTLIGHT_URL + # else: stays None (disabled) + + return None + + +def setup_spotlight(options: "Dict[str, Any]") -> "Optional[SpotlightClient]": + url = _resolve_spotlight_url(options.get("spotlight"), sentry_logger) + + if url is None: + return None + + # Only set up logging handler when spotlight is actually enabled + _handler = logging.StreamHandler(sys.stderr) + _handler.setFormatter(logging.Formatter(" [spotlight] %(levelname)s: %(message)s")) + logger.addHandler(_handler) + logger.setLevel(logging.INFO) + + # Update options with resolved URL for consistency + options["spotlight"] = url + + with capture_internal_exceptions(): + if ( + settings is not None + and settings.DEBUG + and env_to_bool(os.environ.get("SENTRY_SPOTLIGHT_ON_ERROR", "1")) + and env_to_bool(os.environ.get("SENTRY_SPOTLIGHT_MIDDLEWARE", "1")) + ): + middleware = settings.MIDDLEWARE + if DJANGO_SPOTLIGHT_MIDDLEWARE_PATH not in middleware: + settings.MIDDLEWARE = type(middleware)( + chain(middleware, (DJANGO_SPOTLIGHT_MIDDLEWARE_PATH,)) + ) + logger.info("Enabled Spotlight integration for Django") + + client = SpotlightClient(url) + logger.info("Enabled Spotlight using sidecar at %s", url) + + return client diff --git a/lib/python3.12/site-packages/sentry_sdk/traces.py b/lib/python3.12/site-packages/sentry_sdk/traces.py new file mode 100644 index 0000000000000000000000000000000000000000..a63a9bebb03daa28a9cfc3eada4be81d43ae18c4 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/traces.py @@ -0,0 +1,70 @@ +""" +The API in this file is only meant to be used in span streaming mode. + +You can enable span streaming mode via +sentry_sdk.init(_experiments={"trace_lifecycle": "stream"}). +""" + +import uuid +from typing import TYPE_CHECKING + +from sentry_sdk.utils import format_attribute + +if TYPE_CHECKING: + from typing import Optional + from sentry_sdk._types import Attributes, AttributeValue + + +class StreamedSpan: + """ + A span holds timing information of a block of code. + + Spans can have multiple child spans thus forming a span tree. + + This is the Span First span implementation. The original transaction-based + span implementation lives in tracing.Span. + """ + + __slots__ = ( + "name", + "_attributes", + "_trace_id", + ) + + def __init__( + self, + *, + name: str, + attributes: "Optional[Attributes]" = None, + trace_id: "Optional[str]" = None, + ): + self.name: str = name + self._attributes: "Attributes" = {} + if attributes: + for attribute, value in attributes.items(): + self.set_attribute(attribute, value) + + self._trace_id = trace_id + + def get_attributes(self) -> "Attributes": + return self._attributes + + def set_attribute(self, key: str, value: "AttributeValue") -> None: + self._attributes[key] = format_attribute(value) + + def set_attributes(self, attributes: "Attributes") -> None: + for key, value in attributes.items(): + self.set_attribute(key, value) + + def remove_attribute(self, key: str) -> None: + try: + del self._attributes[key] + except KeyError: + pass + + @property + def trace_id(self) -> str: + if not self._trace_id: + self._trace_id = uuid.uuid4().hex + + return self._trace_id diff --git a/lib/python3.12/site-packages/sentry_sdk/tracing.py b/lib/python3.12/site-packages/sentry_sdk/tracing.py new file mode 100644 index 0000000000000000000000000000000000000000..a778da73613a9cee8c2bd4998620a4a60d59ca46 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/tracing.py @@ -0,0 +1,1446 @@ +import uuid +import warnings +from datetime import datetime, timedelta, timezone +from enum import Enum +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk.consts import INSTRUMENTER, SPANDATA, SPANSTATUS, SPANTEMPLATE +from sentry_sdk.profiler.continuous_profiler import get_profiler_id +from sentry_sdk.utils import ( + capture_internal_exceptions, + get_current_thread_meta, + is_valid_sample_rate, + logger, + nanosecond_time, + should_be_treated_as_error, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping, MutableMapping + from typing import ( + Any, + Dict, + Iterator, + List, + Optional, + ParamSpec, + Set, + Tuple, + TypeVar, + Union, + overload, + ) + + from typing_extensions import TypedDict, Unpack + + P = ParamSpec("P") + R = TypeVar("R") + + from sentry_sdk._types import ( + Event, + MeasurementUnit, + MeasurementValue, + SamplingContext, + ) + from sentry_sdk.profiler.continuous_profiler import ContinuousProfile + from sentry_sdk.profiler.transaction_profiler import Profile + + class SpanKwargs(TypedDict, total=False): + trace_id: str + """ + The trace ID of the root span. If this new span is to be the root span, + omit this parameter, and a new trace ID will be generated. + """ + + span_id: str + """The span ID of this span. If omitted, a new span ID will be generated.""" + + parent_span_id: str + """The span ID of the parent span, if applicable.""" + + same_process_as_parent: bool + """Whether this span is in the same process as the parent span.""" + + sampled: bool + """ + Whether the span should be sampled. Overrides the default sampling decision + for this span when provided. + """ + + op: str + """ + The span's operation. A list of recommended values is available here: + https://develop.sentry.dev/sdk/performance/span-operations/ + """ + + description: str + """A description of what operation is being performed within the span. This argument is DEPRECATED. Please use the `name` parameter, instead.""" + + hub: "Optional[sentry_sdk.Hub]" + """The hub to use for this span. This argument is DEPRECATED. Please use the `scope` parameter, instead.""" + + status: str + """The span's status. Possible values are listed at https://develop.sentry.dev/sdk/event-payloads/span/""" + + containing_transaction: "Optional[Transaction]" + """The transaction that this span belongs to.""" + + start_timestamp: "Optional[Union[datetime, float]]" + """ + The timestamp when the span started. If omitted, the current time + will be used. + """ + + scope: "sentry_sdk.Scope" + """The scope to use for this span. If not provided, we use the current scope.""" + + origin: str + """ + The origin of the span. + See https://develop.sentry.dev/sdk/performance/trace-origin/ + Default "manual". + """ + + name: str + """A string describing what operation is being performed within the span/transaction.""" + + class TransactionKwargs(SpanKwargs, total=False): + source: str + """ + A string describing the source of the transaction name. This will be used to determine the transaction's type. + See https://develop.sentry.dev/sdk/event-payloads/transaction/#transaction-annotations for more information. + Default "custom". + """ + + parent_sampled: bool + """Whether the parent transaction was sampled. If True this transaction will be kept, if False it will be discarded.""" + + baggage: "Baggage" + """The W3C baggage header value. (see https://www.w3.org/TR/baggage/)""" + + ProfileContext = TypedDict( + "ProfileContext", + { + "profiler_id": str, + }, + ) + +BAGGAGE_HEADER_NAME = "baggage" +SENTRY_TRACE_HEADER_NAME = "sentry-trace" + + +# Transaction source +# see https://develop.sentry.dev/sdk/event-payloads/transaction/#transaction-annotations +class TransactionSource(str, Enum): + COMPONENT = "component" + CUSTOM = "custom" + ROUTE = "route" + TASK = "task" + URL = "url" + VIEW = "view" + + def __str__(self) -> str: + return self.value + + +# These are typically high cardinality and the server hates them +LOW_QUALITY_TRANSACTION_SOURCES = [ + TransactionSource.URL, +] + +SOURCE_FOR_STYLE = { + "endpoint": TransactionSource.COMPONENT, + "function_name": TransactionSource.COMPONENT, + "handler_name": TransactionSource.COMPONENT, + "method_and_path_pattern": TransactionSource.ROUTE, + "path": TransactionSource.URL, + "route_name": TransactionSource.COMPONENT, + "route_pattern": TransactionSource.ROUTE, + "uri_template": TransactionSource.ROUTE, + "url": TransactionSource.ROUTE, +} + + +def get_span_status_from_http_code(http_status_code: int) -> str: + """ + Returns the Sentry status corresponding to the given HTTP status code. + + See: https://develop.sentry.dev/sdk/event-payloads/contexts/#trace-context + """ + if http_status_code < 400: + return SPANSTATUS.OK + + elif 400 <= http_status_code < 500: + if http_status_code == 403: + return SPANSTATUS.PERMISSION_DENIED + elif http_status_code == 404: + return SPANSTATUS.NOT_FOUND + elif http_status_code == 429: + return SPANSTATUS.RESOURCE_EXHAUSTED + elif http_status_code == 413: + return SPANSTATUS.FAILED_PRECONDITION + elif http_status_code == 401: + return SPANSTATUS.UNAUTHENTICATED + elif http_status_code == 409: + return SPANSTATUS.ALREADY_EXISTS + else: + return SPANSTATUS.INVALID_ARGUMENT + + elif 500 <= http_status_code < 600: + if http_status_code == 504: + return SPANSTATUS.DEADLINE_EXCEEDED + elif http_status_code == 501: + return SPANSTATUS.UNIMPLEMENTED + elif http_status_code == 503: + return SPANSTATUS.UNAVAILABLE + else: + return SPANSTATUS.INTERNAL_ERROR + + return SPANSTATUS.UNKNOWN_ERROR + + +class _SpanRecorder: + """Limits the number of spans recorded in a transaction.""" + + __slots__ = ("maxlen", "spans", "dropped_spans") + + def __init__(self, maxlen: int) -> None: + # FIXME: this is `maxlen - 1` only to preserve historical behavior + # enforced by tests. + # Either this should be changed to `maxlen` or the JS SDK implementation + # should be changed to match a consistent interpretation of what maxlen + # limits: either transaction+spans or only child spans. + self.maxlen = maxlen - 1 + self.spans: "List[Span]" = [] + self.dropped_spans: int = 0 + + def add(self, span: "Span") -> None: + if len(self.spans) > self.maxlen: + span._span_recorder = None + self.dropped_spans += 1 + else: + self.spans.append(span) + + +class Span: + """A span holds timing information of a block of code. + Spans can have multiple child spans thus forming a span tree. + + :param trace_id: The trace ID of the root span. If this new span is to be the root span, + omit this parameter, and a new trace ID will be generated. + :param span_id: The span ID of this span. If omitted, a new span ID will be generated. + :param parent_span_id: The span ID of the parent span, if applicable. + :param same_process_as_parent: Whether this span is in the same process as the parent span. + :param sampled: Whether the span should be sampled. Overrides the default sampling decision + for this span when provided. + :param op: The span's operation. A list of recommended values is available here: + https://develop.sentry.dev/sdk/performance/span-operations/ + :param description: A description of what operation is being performed within the span. + + .. deprecated:: 2.15.0 + Please use the `name` parameter, instead. + :param name: A string describing what operation is being performed within the span. + :param hub: The hub to use for this span. + + .. deprecated:: 2.0.0 + Please use the `scope` parameter, instead. + :param status: The span's status. Possible values are listed at + https://develop.sentry.dev/sdk/event-payloads/span/ + :param containing_transaction: The transaction that this span belongs to. + :param start_timestamp: The timestamp when the span started. If omitted, the current time + will be used. + :param scope: The scope to use for this span. If not provided, we use the current scope. + """ + + __slots__ = ( + "_trace_id", + "_span_id", + "parent_span_id", + "same_process_as_parent", + "sampled", + "op", + "description", + "_measurements", + "start_timestamp", + "_start_timestamp_monotonic_ns", + "status", + "timestamp", + "_tags", + "_data", + "_span_recorder", + "hub", + "_context_manager_state", + "_containing_transaction", + "scope", + "origin", + "name", + "_flags", + "_flags_capacity", + ) + + def __init__( + self, + trace_id: "Optional[str]" = None, + span_id: "Optional[str]" = None, + parent_span_id: "Optional[str]" = None, + same_process_as_parent: bool = True, + sampled: "Optional[bool]" = None, + op: "Optional[str]" = None, + description: "Optional[str]" = None, + hub: "Optional[sentry_sdk.Hub]" = None, # deprecated + status: "Optional[str]" = None, + containing_transaction: "Optional[Transaction]" = None, + start_timestamp: "Optional[Union[datetime, float]]" = None, + scope: "Optional[sentry_sdk.Scope]" = None, + origin: str = "manual", + name: "Optional[str]" = None, + ) -> None: + self._trace_id = trace_id + self._span_id = span_id + self.parent_span_id = parent_span_id + self.same_process_as_parent = same_process_as_parent + self.sampled = sampled + self.op = op + self.description = name or description + self.status = status + self.hub = hub # backwards compatibility + self.scope = scope + self.origin = origin + self._measurements: "Dict[str, MeasurementValue]" = {} + self._tags: "MutableMapping[str, str]" = {} + self._data: "Dict[str, Any]" = {} + self._containing_transaction = containing_transaction + self._flags: "Dict[str, bool]" = {} + self._flags_capacity = 10 + + if hub is not None: + warnings.warn( + "The `hub` parameter is deprecated. Please use `scope` instead.", + DeprecationWarning, + stacklevel=2, + ) + + self.scope = self.scope or hub.scope + + if start_timestamp is None: + start_timestamp = datetime.now(timezone.utc) + elif isinstance(start_timestamp, float): + start_timestamp = datetime.fromtimestamp(start_timestamp, timezone.utc) + self.start_timestamp = start_timestamp + try: + # profiling depends on this value and requires that + # it is measured in nanoseconds + self._start_timestamp_monotonic_ns = nanosecond_time() + except AttributeError: + pass + + #: End timestamp of span + self.timestamp: "Optional[datetime]" = None + + self._span_recorder: "Optional[_SpanRecorder]" = None + + self.update_active_thread() + self.set_profiler_id(get_profiler_id()) + + # TODO this should really live on the Transaction class rather than the Span + # class + def init_span_recorder(self, maxlen: int) -> None: + if self._span_recorder is None: + self._span_recorder = _SpanRecorder(maxlen) + + @property + def trace_id(self) -> str: + if not self._trace_id: + self._trace_id = uuid.uuid4().hex + + return self._trace_id + + @trace_id.setter + def trace_id(self, value: str) -> None: + self._trace_id = value + + @property + def span_id(self) -> str: + if not self._span_id: + self._span_id = uuid.uuid4().hex[16:] + + return self._span_id + + @span_id.setter + def span_id(self, value: str) -> None: + self._span_id = value + + def __repr__(self) -> str: + return ( + "<%s(op=%r, description:%r, trace_id=%r, span_id=%r, parent_span_id=%r, sampled=%r, origin=%r)>" + % ( + self.__class__.__name__, + self.op, + self.description, + self.trace_id, + self.span_id, + self.parent_span_id, + self.sampled, + self.origin, + ) + ) + + def __enter__(self) -> "Span": + scope = self.scope or sentry_sdk.get_current_scope() + old_span = scope.span + scope.span = self + self._context_manager_state = (scope, old_span) + return self + + def __exit__( + self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]" + ) -> None: + if value is not None and should_be_treated_as_error(ty, value): + self.set_status(SPANSTATUS.INTERNAL_ERROR) + + with capture_internal_exceptions(): + scope, old_span = self._context_manager_state + del self._context_manager_state + self.finish(scope) + scope.span = old_span + + @property + def containing_transaction(self) -> "Optional[Transaction]": + """The ``Transaction`` that this span belongs to. + The ``Transaction`` is the root of the span tree, + so one could also think of this ``Transaction`` as the "root span".""" + + # this is a getter rather than a regular attribute so that transactions + # can return `self` here instead (as a way to prevent them circularly + # referencing themselves) + return self._containing_transaction + + def start_child( + self, instrumenter: str = INSTRUMENTER.SENTRY, **kwargs: "Any" + ) -> "Span": + """ + Start a sub-span from the current span or transaction. + + Takes the same arguments as the initializer of :py:class:`Span`. The + trace id, sampling decision, transaction pointer, and span recorder are + inherited from the current span/transaction. + + The instrumenter parameter is deprecated for user code, and it will + be removed in the next major version. Going forward, it should only + be used by the SDK itself. + """ + if kwargs.get("description") is not None: + warnings.warn( + "The `description` parameter is deprecated. Please use `name` instead.", + DeprecationWarning, + stacklevel=2, + ) + + configuration_instrumenter = sentry_sdk.get_client().options["instrumenter"] + + if instrumenter != configuration_instrumenter: + return NoOpSpan() + + kwargs.setdefault("sampled", self.sampled) + + child = Span( + trace_id=self.trace_id, + parent_span_id=self.span_id, + containing_transaction=self.containing_transaction, + **kwargs, + ) + + span_recorder = ( + self.containing_transaction and self.containing_transaction._span_recorder + ) + if span_recorder: + span_recorder.add(child) + + return child + + @classmethod + def continue_from_environ( + cls, + environ: "Mapping[str, str]", + **kwargs: "Any", + ) -> "Transaction": + """ + DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`. + + Create a Transaction with the given params, then add in data pulled from + the ``sentry-trace`` and ``baggage`` headers from the environ (if any) + before returning the Transaction. + + This is different from :py:meth:`~sentry_sdk.tracing.Span.continue_from_headers` + in that it assumes header names in the form ``HTTP_HEADER_NAME`` - + such as you would get from a WSGI/ASGI environ - + rather than the form ``header-name``. + + :param environ: The ASGI/WSGI environ to pull information from. + """ + return Transaction.continue_from_headers(EnvironHeaders(environ), **kwargs) + + @classmethod + def continue_from_headers( + cls, + headers: "Mapping[str, str]", + *, + _sample_rand: "Optional[str]" = None, + **kwargs: "Any", + ) -> "Transaction": + """ + DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`. + + Create a transaction with the given params (including any data pulled from + the ``sentry-trace`` and ``baggage`` headers). + + :param headers: The dictionary with the HTTP headers to pull information from. + :param _sample_rand: If provided, we override the sample_rand value from the + incoming headers with this value. (internal use only) + """ + logger.warning("Deprecated: use sentry_sdk.continue_trace instead.") + + # TODO-neel move away from this kwargs stuff, it's confusing and opaque + # make more explicit + baggage = Baggage.from_incoming_header( + headers.get(BAGGAGE_HEADER_NAME), _sample_rand=_sample_rand + ) + kwargs.update({BAGGAGE_HEADER_NAME: baggage}) + + sentrytrace_kwargs = extract_sentrytrace_data( + headers.get(SENTRY_TRACE_HEADER_NAME) + ) + + if sentrytrace_kwargs is not None: + kwargs.update(sentrytrace_kwargs) + + # If there's an incoming sentry-trace but no incoming baggage header, + # for instance in traces coming from older SDKs, + # baggage will be empty and immutable and won't be populated as head SDK. + baggage.freeze() + + transaction = Transaction(**kwargs) + transaction.same_process_as_parent = False + + return transaction + + def iter_headers(self) -> "Iterator[Tuple[str, str]]": + """ + Creates a generator which returns the span's ``sentry-trace`` and ``baggage`` headers. + If the span's containing transaction doesn't yet have a ``baggage`` value, + this will cause one to be generated and stored. + """ + if not self.containing_transaction: + # Do not propagate headers if there is no containing transaction. Otherwise, this + # span ends up being the root span of a new trace, and since it does not get sent + # to Sentry, the trace will be missing a root transaction. The dynamic sampling + # context will also be missing, breaking dynamic sampling & traces. + return + + yield SENTRY_TRACE_HEADER_NAME, self.to_traceparent() + + baggage = self.containing_transaction.get_baggage().serialize() + if baggage: + yield BAGGAGE_HEADER_NAME, baggage + + @classmethod + def from_traceparent( + cls, + traceparent: "Optional[str]", + **kwargs: "Any", + ) -> "Optional[Transaction]": + """ + DEPRECATED: Use :py:meth:`sentry_sdk.continue_trace`. + + Create a ``Transaction`` with the given params, then add in data pulled from + the given ``sentry-trace`` header value before returning the ``Transaction``. + """ + if not traceparent: + return None + + return cls.continue_from_headers( + {SENTRY_TRACE_HEADER_NAME: traceparent}, **kwargs + ) + + def to_traceparent(self) -> str: + if self.sampled is True: + sampled = "1" + elif self.sampled is False: + sampled = "0" + else: + sampled = None + + traceparent = "%s-%s" % (self.trace_id, self.span_id) + if sampled is not None: + traceparent += "-%s" % (sampled,) + + return traceparent + + def to_baggage(self) -> "Optional[Baggage]": + """Returns the :py:class:`~sentry_sdk.tracing_utils.Baggage` + associated with this ``Span``, if any. (Taken from the root of the span tree.) + """ + if self.containing_transaction: + return self.containing_transaction.get_baggage() + return None + + def set_tag(self, key: str, value: "Any") -> None: + self._tags[key] = value + + def set_data(self, key: str, value: "Any") -> None: + self._data[key] = value + + def update_data(self, data: "Dict[str, Any]") -> None: + self._data.update(data) + + def set_flag(self, flag: str, result: bool) -> None: + if len(self._flags) < self._flags_capacity: + self._flags[flag] = result + + def set_status(self, value: str) -> None: + self.status = value + + def set_measurement( + self, name: str, value: float, unit: "MeasurementUnit" = "" + ) -> None: + """ + .. deprecated:: 2.28.0 + This function is deprecated and will be removed in the next major release. + """ + + warnings.warn( + "`set_measurement()` is deprecated and will be removed in the next major version. Please use `set_data()` instead.", + DeprecationWarning, + stacklevel=2, + ) + self._measurements[name] = {"value": value, "unit": unit} + + def set_thread( + self, thread_id: "Optional[int]", thread_name: "Optional[str]" + ) -> None: + if thread_id is not None: + self.set_data(SPANDATA.THREAD_ID, str(thread_id)) + + if thread_name is not None: + self.set_data(SPANDATA.THREAD_NAME, thread_name) + + def set_profiler_id(self, profiler_id: "Optional[str]") -> None: + if profiler_id is not None: + self.set_data(SPANDATA.PROFILER_ID, profiler_id) + + def set_http_status(self, http_status: int) -> None: + self.set_tag( + "http.status_code", str(http_status) + ) # TODO-neel remove in major, we keep this for backwards compatibility + self.set_data(SPANDATA.HTTP_STATUS_CODE, http_status) + self.set_status(get_span_status_from_http_code(http_status)) + + def is_success(self) -> bool: + return self.status == "ok" + + def finish( + self, + scope: "Optional[sentry_sdk.Scope]" = None, + end_timestamp: "Optional[Union[float, datetime]]" = None, + ) -> "Optional[str]": + """ + Sets the end timestamp of the span. + + Additionally it also creates a breadcrumb from the span, + if the span represents a database or HTTP request. + + :param scope: The scope to use for this transaction. + If not provided, the current scope will be used. + :param end_timestamp: Optional timestamp that should + be used as timestamp instead of the current time. + + :return: Always ``None``. The type is ``Optional[str]`` to match + the return value of :py:meth:`sentry_sdk.tracing.Transaction.finish`. + """ + if self.timestamp is not None: + # This span is already finished, ignore. + return None + + try: + if end_timestamp: + if isinstance(end_timestamp, float): + end_timestamp = datetime.fromtimestamp(end_timestamp, timezone.utc) + self.timestamp = end_timestamp + else: + elapsed = nanosecond_time() - self._start_timestamp_monotonic_ns + self.timestamp = self.start_timestamp + timedelta( + microseconds=elapsed / 1000 + ) + except AttributeError: + self.timestamp = datetime.now(timezone.utc) + + scope = scope or sentry_sdk.get_current_scope() + + # Copy conversation_id from scope to span data if this is an AI span + conversation_id = scope.get_conversation_id() + if conversation_id: + has_ai_op = SPANDATA.GEN_AI_OPERATION_NAME in self._data + is_ai_span_op = self.op is not None and ( + self.op.startswith("ai.") or self.op.startswith("gen_ai.") + ) + if has_ai_op or is_ai_span_op: + self.set_data("gen_ai.conversation.id", conversation_id) + + maybe_create_breadcrumbs_from_span(scope, self) + + return None + + def to_json(self) -> "Dict[str, Any]": + """Returns a JSON-compatible representation of the span.""" + + rv: "Dict[str, Any]" = { + "trace_id": self.trace_id, + "span_id": self.span_id, + "parent_span_id": self.parent_span_id, + "same_process_as_parent": self.same_process_as_parent, + "op": self.op, + "description": self.description, + "start_timestamp": self.start_timestamp, + "timestamp": self.timestamp, + "origin": self.origin, + } + + if self.status: + rv["status"] = self.status + # TODO-neel remove redundant tag in major + self._tags["status"] = self.status + + if len(self._measurements) > 0: + rv["measurements"] = self._measurements + + tags = self._tags + if tags: + rv["tags"] = tags + + data = {} + data.update(self._flags) + data.update(self._data) + if data: + rv["data"] = data + + return rv + + def get_trace_context(self) -> "Any": + rv: "Dict[str, Any]" = { + "trace_id": self.trace_id, + "span_id": self.span_id, + "parent_span_id": self.parent_span_id, + "op": self.op, + "description": self.description, + "origin": self.origin, + } + if self.status: + rv["status"] = self.status + + if self.containing_transaction: + rv["dynamic_sampling_context"] = ( + self.containing_transaction.get_baggage().dynamic_sampling_context() + ) + + data = {} + + thread_id = self._data.get(SPANDATA.THREAD_ID) + if thread_id is not None: + data["thread.id"] = thread_id + + thread_name = self._data.get(SPANDATA.THREAD_NAME) + if thread_name is not None: + data["thread.name"] = thread_name + + if data: + rv["data"] = data + + return rv + + def get_profile_context(self) -> "Optional[ProfileContext]": + profiler_id = self._data.get(SPANDATA.PROFILER_ID) + if profiler_id is None: + return None + + return { + "profiler_id": profiler_id, + } + + def update_active_thread(self) -> None: + thread_id, thread_name = get_current_thread_meta() + self.set_thread(thread_id, thread_name) + + +class Transaction(Span): + """The Transaction is the root element that holds all the spans + for Sentry performance instrumentation. + + :param name: Identifier of the transaction. + Will show up in the Sentry UI. + :param parent_sampled: Whether the parent transaction was sampled. + If True this transaction will be kept, if False it will be discarded. + :param baggage: The W3C baggage header value. + (see https://www.w3.org/TR/baggage/) + :param source: A string describing the source of the transaction name. + This will be used to determine the transaction's type. + See https://develop.sentry.dev/sdk/event-payloads/transaction/#transaction-annotations + for more information. Default "custom". + :param kwargs: Additional arguments to be passed to the Span constructor. + See :py:class:`sentry_sdk.tracing.Span` for available arguments. + """ + + __slots__ = ( + "name", + "source", + "parent_sampled", + # used to create baggage value for head SDKs in dynamic sampling + "sample_rate", + "_measurements", + "_contexts", + "_profile", + "_continuous_profile", + "_baggage", + "_sample_rand", + ) + + def __init__( # type: ignore[misc] + self, + name: str = "", + parent_sampled: "Optional[bool]" = None, + baggage: "Optional[Baggage]" = None, + source: str = TransactionSource.CUSTOM, + **kwargs: "Unpack[SpanKwargs]", + ) -> None: + super().__init__(**kwargs) + + self.name = name + self.source = source + self.sample_rate: "Optional[float]" = None + self.parent_sampled = parent_sampled + self._measurements: "Dict[str, MeasurementValue]" = {} + self._contexts: "Dict[str, Any]" = {} + self._profile: "Optional[Profile]" = None + self._continuous_profile: "Optional[ContinuousProfile]" = None + self._baggage = baggage + + baggage_sample_rand = ( + None if self._baggage is None else self._baggage._sample_rand() + ) + if baggage_sample_rand is not None: + self._sample_rand = baggage_sample_rand + else: + self._sample_rand = _generate_sample_rand(self.trace_id) + + def __repr__(self) -> str: + return ( + "<%s(name=%r, op=%r, trace_id=%r, span_id=%r, parent_span_id=%r, sampled=%r, source=%r, origin=%r)>" + % ( + self.__class__.__name__, + self.name, + self.op, + self.trace_id, + self.span_id, + self.parent_span_id, + self.sampled, + self.source, + self.origin, + ) + ) + + def _possibly_started(self) -> bool: + """Returns whether the transaction might have been started. + + If this returns False, we know that the transaction was not started + with sentry_sdk.start_transaction, and therefore the transaction will + be discarded. + """ + + # We must explicitly check self.sampled is False since self.sampled can be None + return self._span_recorder is not None or self.sampled is False + + def __enter__(self) -> "Transaction": + if not self._possibly_started(): + logger.debug( + "Transaction was entered without being started with sentry_sdk.start_transaction." + "The transaction will not be sent to Sentry. To fix, start the transaction by" + "passing it to sentry_sdk.start_transaction." + ) + + super().__enter__() + + if self._profile is not None: + self._profile.__enter__() + + return self + + def __exit__( + self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]" + ) -> None: + if self._profile is not None: + self._profile.__exit__(ty, value, tb) + + if self._continuous_profile is not None: + self._continuous_profile.stop() + + super().__exit__(ty, value, tb) + + @property + def containing_transaction(self) -> "Transaction": + """The root element of the span tree. + In the case of a transaction it is the transaction itself. + """ + + # Transactions (as spans) belong to themselves (as transactions). This + # is a getter rather than a regular attribute to avoid having a circular + # reference. + return self + + def _get_scope_from_finish_args( + self, + scope_arg: "Optional[Union[sentry_sdk.Scope, sentry_sdk.Hub]]", + hub_arg: "Optional[Union[sentry_sdk.Scope, sentry_sdk.Hub]]", + ) -> "Optional[sentry_sdk.Scope]": + """ + Logic to get the scope from the arguments passed to finish. This + function exists for backwards compatibility with the old finish. + + TODO: Remove this function in the next major version. + """ + scope_or_hub = scope_arg + if hub_arg is not None: + warnings.warn( + "The `hub` parameter is deprecated. Please use the `scope` parameter, instead.", + DeprecationWarning, + stacklevel=3, + ) + + scope_or_hub = hub_arg + + if isinstance(scope_or_hub, sentry_sdk.Hub): + warnings.warn( + "Passing a Hub to finish is deprecated. Please pass a Scope, instead.", + DeprecationWarning, + stacklevel=3, + ) + + return scope_or_hub.scope + + return scope_or_hub + + def _get_log_representation(self) -> str: + return "{op}transaction <{name}>".format( + op=("<" + self.op + "> " if self.op else ""), name=self.name + ) + + def finish( + self, + scope: "Optional[sentry_sdk.Scope]" = None, + end_timestamp: "Optional[Union[float, datetime]]" = None, + *, + hub: "Optional[sentry_sdk.Hub]" = None, + ) -> "Optional[str]": + """Finishes the transaction and sends it to Sentry. + All finished spans in the transaction will also be sent to Sentry. + + :param scope: The Scope to use for this transaction. + If not provided, the current Scope will be used. + :param end_timestamp: Optional timestamp that should + be used as timestamp instead of the current time. + :param hub: The hub to use for this transaction. + This argument is DEPRECATED. Please use the `scope` + parameter, instead. + + :return: The event ID if the transaction was sent to Sentry, + otherwise None. + """ + if self.timestamp is not None: + # This transaction is already finished, ignore. + return None + + # For backwards compatibility, we must handle the case where `scope` + # or `hub` could both either be a `Scope` or a `Hub`. + scope: "Optional[sentry_sdk.Scope]" = self._get_scope_from_finish_args( + scope, hub + ) + + scope = scope or self.scope or sentry_sdk.get_current_scope() + client = sentry_sdk.get_client() + + if not client.is_active(): + # We have no active client and therefore nowhere to send this transaction. + return None + + if self._span_recorder is None: + # Explicit check against False needed because self.sampled might be None + if self.sampled is False: + logger.debug("Discarding transaction because sampled = False") + else: + logger.debug( + "Discarding transaction because it was not started with sentry_sdk.start_transaction" + ) + + # This is not entirely accurate because discards here are not + # exclusively based on sample rate but also traces sampler, but + # we handle this the same here. + if client.transport and has_tracing_enabled(client.options): + if client.monitor and client.monitor.downsample_factor > 0: + reason = "backpressure" + else: + reason = "sample_rate" + + client.transport.record_lost_event(reason, data_category="transaction") + + # Only one span (the transaction itself) is discarded, since we did not record any spans here. + client.transport.record_lost_event(reason, data_category="span") + return None + + if not self.name: + logger.warning( + "Transaction has no name, falling back to ``." + ) + self.name = "" + + super().finish(scope, end_timestamp) + + status_code = self._data.get(SPANDATA.HTTP_STATUS_CODE) + if ( + status_code is not None + and status_code in client.options["trace_ignore_status_codes"] + ): + logger.debug( + "[Tracing] Discarding {transaction_description} because the HTTP status code {status_code} is matched by trace_ignore_status_codes: {trace_ignore_status_codes}".format( + transaction_description=self._get_log_representation(), + status_code=self._data[SPANDATA.HTTP_STATUS_CODE], + trace_ignore_status_codes=client.options[ + "trace_ignore_status_codes" + ], + ) + ) + if client.transport: + client.transport.record_lost_event( + "event_processor", data_category="transaction" + ) + + num_spans = len(self._span_recorder.spans) + 1 + client.transport.record_lost_event( + "event_processor", data_category="span", quantity=num_spans + ) + + self.sampled = False + + if not self.sampled: + # At this point a `sampled = None` should have already been resolved + # to a concrete decision. + if self.sampled is None: + logger.warning("Discarding transaction without sampling decision.") + + return None + + finished_spans = [ + span.to_json() + for span in self._span_recorder.spans + if span.timestamp is not None + ] + + len_diff = len(self._span_recorder.spans) - len(finished_spans) + dropped_spans = len_diff + self._span_recorder.dropped_spans + + # we do this to break the circular reference of transaction -> span + # recorder -> span -> containing transaction (which is where we started) + # before either the spans or the transaction goes out of scope and has + # to be garbage collected + self._span_recorder = None + + contexts = {} + contexts.update(self._contexts) + contexts.update({"trace": self.get_trace_context()}) + profile_context = self.get_profile_context() + if profile_context is not None: + contexts.update({"profile": profile_context}) + + event: "Event" = { + "type": "transaction", + "transaction": self.name, + "transaction_info": {"source": self.source}, + "contexts": contexts, + "tags": self._tags, + "timestamp": self.timestamp, + "start_timestamp": self.start_timestamp, + "spans": finished_spans, + } + + if dropped_spans > 0: + event["_dropped_spans"] = dropped_spans + + if self._profile is not None and self._profile.valid(): + event["profile"] = self._profile + self._profile = None + + event["measurements"] = self._measurements + + return scope.capture_event(event) + + def set_measurement( + self, name: str, value: float, unit: "MeasurementUnit" = "" + ) -> None: + """ + .. deprecated:: 2.28.0 + This function is deprecated and will be removed in the next major release. + """ + + warnings.warn( + "`set_measurement()` is deprecated and will be removed in the next major version. Please use `set_data()` instead.", + DeprecationWarning, + stacklevel=2, + ) + self._measurements[name] = {"value": value, "unit": unit} + + def set_context(self, key: str, value: "dict[str, Any]") -> None: + """Sets a context. Transactions can have multiple contexts + and they should follow the format described in the "Contexts Interface" + documentation. + + :param key: The name of the context. + :param value: The information about the context. + """ + self._contexts[key] = value + + def set_http_status(self, http_status: int) -> None: + """Sets the status of the Transaction according to the given HTTP status. + + :param http_status: The HTTP status code.""" + super().set_http_status(http_status) + self.set_context("response", {"status_code": http_status}) + + def to_json(self) -> "Dict[str, Any]": + """Returns a JSON-compatible representation of the transaction.""" + rv = super().to_json() + + rv["name"] = self.name + rv["source"] = self.source + rv["sampled"] = self.sampled + + return rv + + def get_trace_context(self) -> "Any": + trace_context = super().get_trace_context() + + if self._data: + trace_context["data"] = self._data + + return trace_context + + def get_baggage(self) -> "Baggage": + """Returns the :py:class:`~sentry_sdk.tracing_utils.Baggage` + associated with the Transaction. + + The first time a new baggage with Sentry items is made, + it will be frozen.""" + if not self._baggage or self._baggage.mutable: + self._baggage = Baggage.populate_from_transaction(self) + + return self._baggage + + def _set_initial_sampling_decision( + self, sampling_context: "SamplingContext" + ) -> None: + """ + Sets the transaction's sampling decision, according to the following + precedence rules: + + 1. If a sampling decision is passed to `start_transaction` + (`start_transaction(name: "my transaction", sampled: True)`), that + decision will be used, regardless of anything else + + 2. If `traces_sampler` is defined, its decision will be used. It can + choose to keep or ignore any parent sampling decision, or use the + sampling context data to make its own decision or to choose a sample + rate for the transaction. + + 3. If `traces_sampler` is not defined, but there's a parent sampling + decision, the parent sampling decision will be used. + + 4. If `traces_sampler` is not defined and there's no parent sampling + decision, `traces_sample_rate` will be used. + """ + client = sentry_sdk.get_client() + + transaction_description = self._get_log_representation() + + # nothing to do if tracing is disabled + if not has_tracing_enabled(client.options): + self.sampled = False + return + + # if the user has forced a sampling decision by passing a `sampled` + # value when starting the transaction, go with that + if self.sampled is not None: + self.sample_rate = float(self.sampled) + return + + # we would have bailed already if neither `traces_sampler` nor + # `traces_sample_rate` were defined, so one of these should work; prefer + # the hook if so + sample_rate = ( + client.options["traces_sampler"](sampling_context) + if callable(client.options.get("traces_sampler")) + # default inheritance behavior + else ( + sampling_context["parent_sampled"] + if sampling_context["parent_sampled"] is not None + else client.options["traces_sample_rate"] + ) + ) + + # Since this is coming from the user (or from a function provided by the + # user), who knows what we might get. (The only valid values are + # booleans or numbers between 0 and 1.) + if not is_valid_sample_rate(sample_rate, source="Tracing"): + logger.warning( + "[Tracing] Discarding {transaction_description} because of invalid sample rate.".format( + transaction_description=transaction_description, + ) + ) + self.sampled = False + return + + self.sample_rate = float(sample_rate) + + if client.monitor: + self.sample_rate /= 2**client.monitor.downsample_factor + + # if the function returned 0 (or false), or if `traces_sample_rate` is + # 0, it's a sign the transaction should be dropped + if not self.sample_rate: + logger.debug( + "[Tracing] Discarding {transaction_description} because {reason}".format( + transaction_description=transaction_description, + reason=( + "traces_sampler returned 0 or False" + if callable(client.options.get("traces_sampler")) + else "traces_sample_rate is set to 0" + ), + ) + ) + self.sampled = False + return + + # Now we roll the dice. + self.sampled = self._sample_rand < self.sample_rate + + if self.sampled: + logger.debug( + "[Tracing] Starting {transaction_description}".format( + transaction_description=transaction_description, + ) + ) + else: + logger.debug( + "[Tracing] Discarding {transaction_description} because it's not included in the random sample (sampling rate = {sample_rate})".format( + transaction_description=transaction_description, + sample_rate=self.sample_rate, + ) + ) + + +class NoOpSpan(Span): + def __repr__(self) -> str: + return "<%s>" % self.__class__.__name__ + + @property + def containing_transaction(self) -> "Optional[Transaction]": + return None + + def start_child( + self, instrumenter: str = INSTRUMENTER.SENTRY, **kwargs: "Any" + ) -> "NoOpSpan": + return NoOpSpan() + + def to_traceparent(self) -> str: + return "" + + def to_baggage(self) -> "Optional[Baggage]": + return None + + def get_baggage(self) -> "Optional[Baggage]": + return None + + def iter_headers(self) -> "Iterator[Tuple[str, str]]": + return iter(()) + + def set_tag(self, key: str, value: "Any") -> None: + pass + + def set_data(self, key: str, value: "Any") -> None: + pass + + def update_data(self, data: "Dict[str, Any]") -> None: + pass + + def set_status(self, value: str) -> None: + pass + + def set_http_status(self, http_status: int) -> None: + pass + + def is_success(self) -> bool: + return True + + def to_json(self) -> "Dict[str, Any]": + return {} + + def get_trace_context(self) -> "Any": + return {} + + def get_profile_context(self) -> "Any": + return {} + + def finish( + self, + scope: "Optional[sentry_sdk.Scope]" = None, + end_timestamp: "Optional[Union[float, datetime]]" = None, + *, + hub: "Optional[sentry_sdk.Hub]" = None, + ) -> "Optional[str]": + """ + The `hub` parameter is deprecated. Please use the `scope` parameter, instead. + """ + pass + + def set_measurement( + self, name: str, value: float, unit: "MeasurementUnit" = "" + ) -> None: + pass + + def set_context(self, key: str, value: "dict[str, Any]") -> None: + pass + + def init_span_recorder(self, maxlen: int) -> None: + pass + + def _set_initial_sampling_decision( + self, sampling_context: "SamplingContext" + ) -> None: + pass + + +if TYPE_CHECKING: + + @overload + def trace( + func: None = None, + *, + op: "Optional[str]" = None, + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, + template: "SPANTEMPLATE" = SPANTEMPLATE.DEFAULT, + ) -> "Callable[[Callable[P, R]], Callable[P, R]]": + # Handles: @trace() and @trace(op="custom") + pass + + @overload + def trace(func: "Callable[P, R]") -> "Callable[P, R]": + # Handles: @trace + pass + + +def trace( + func: "Optional[Callable[P, R]]" = None, + *, + op: "Optional[str]" = None, + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, + template: "SPANTEMPLATE" = SPANTEMPLATE.DEFAULT, +) -> "Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]": + """ + Decorator to start a child span around a function call. + + This decorator automatically creates a new span when the decorated function + is called, and finishes the span when the function returns or raises an exception. + + :param func: The function to trace. When used as a decorator without parentheses, + this is the function being decorated. When used with parameters (e.g., + ``@trace(op="custom")``, this should be None. + :type func: Callable or None + + :param op: The operation name for the span. This is a high-level description + of what the span represents (e.g., "http.client", "db.query"). + You can use predefined constants from :py:class:`sentry_sdk.consts.OP` + or provide your own string. If not provided, a default operation will + be assigned based on the template. + :type op: str or None + + :param name: The human-readable name/description for the span. If not provided, + defaults to the function name. This provides more specific details about + what the span represents (e.g., "GET /api/users", "process_user_data"). + :type name: str or None + + :param attributes: A dictionary of key-value pairs to add as attributes to the span. + Attribute values must be strings, integers, floats, or booleans. These + attributes provide additional context about the span's execution. + :type attributes: dict[str, Any] or None + + :param template: The type of span to create. This determines what kind of + span instrumentation and data collection will be applied. Use predefined + constants from :py:class:`sentry_sdk.consts.SPANTEMPLATE`. + The default is `SPANTEMPLATE.DEFAULT` which is the right choice for most + use cases. + :type template: :py:class:`sentry_sdk.consts.SPANTEMPLATE` + + :returns: When used as ``@trace``, returns the decorated function. When used as + ``@trace(...)`` with parameters, returns a decorator function. + :rtype: Callable or decorator function + + Example:: + + import sentry_sdk + from sentry_sdk.consts import OP, SPANTEMPLATE + + # Simple usage with default values + @sentry_sdk.trace + def process_data(): + # Function implementation + pass + + # With custom parameters + @sentry_sdk.trace( + op=OP.DB_QUERY, + name="Get user data", + attributes={"postgres": True} + ) + def make_db_query(sql): + # Function implementation + pass + + # With a custom template + @sentry_sdk.trace(template=SPANTEMPLATE.AI_TOOL) + def calculate_interest_rate(amount, rate, years): + # Function implementation + pass + """ + from sentry_sdk.tracing_utils import create_span_decorator + + decorator = create_span_decorator( + op=op, + name=name, + attributes=attributes, + template=template, + ) + + if func: + return decorator(func) + else: + return decorator + + +# Circular imports + +from sentry_sdk.tracing_utils import ( + Baggage, + EnvironHeaders, + _generate_sample_rand, + extract_sentrytrace_data, + has_tracing_enabled, + maybe_create_breadcrumbs_from_span, +) diff --git a/lib/python3.12/site-packages/sentry_sdk/tracing_utils.py b/lib/python3.12/site-packages/sentry_sdk/tracing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c1d6c4453531e50b2bf9c77897c0ffdf94c52fbb --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/tracing_utils.py @@ -0,0 +1,1322 @@ +import contextlib +import functools +import inspect +import os +import re +import sys +from collections.abc import Mapping, MutableMapping +from datetime import timedelta +from random import Random +from urllib.parse import quote, unquote +import uuid + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS, SPANTEMPLATE +from sentry_sdk.utils import ( + capture_internal_exceptions, + filename_for_module, + logger, + match_regex_list, + qualname_from_function, + safe_repr, + to_string, + try_convert, + is_sentry_url, + _is_external_source, + _is_in_project_root, + _module_in_list, +) + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any + from typing import Dict + from typing import Generator + from typing import Optional + from typing import Union + from typing import Iterator + from typing import Tuple + + from types import FrameType + + +SENTRY_TRACE_REGEX = re.compile( + "^[ \t]*" # whitespace + "([0-9a-f]{32})?" # trace_id + "-?([0-9a-f]{16})?" # span_id + "-?([01])?" # sampled + "[ \t]*$" # whitespace +) + + +# This is a normal base64 regex, modified to reflect that fact that we strip the +# trailing = or == off +base64_stripped = ( + # any of the characters in the base64 "alphabet", in multiples of 4 + "([a-zA-Z0-9+/]{4})*" + # either nothing or 2 or 3 base64-alphabet characters (see + # https://en.wikipedia.org/wiki/Base64#Decoding_Base64_without_padding for + # why there's never only 1 extra character) + "([a-zA-Z0-9+/]{2,3})?" +) + + +class EnvironHeaders(Mapping): # type: ignore + def __init__( + self, + environ: "Mapping[str, str]", + prefix: str = "HTTP_", + ) -> None: + self.environ = environ + self.prefix = prefix + + def __getitem__(self, key: str) -> "Optional[Any]": + return self.environ[self.prefix + key.replace("-", "_").upper()] + + def __len__(self) -> int: + return sum(1 for _ in iter(self)) + + def __iter__(self) -> "Generator[str, None, None]": + for k in self.environ: + if not isinstance(k, str): + continue + + k = k.replace("-", "_").upper() + if not k.startswith(self.prefix): + continue + + yield k[len(self.prefix) :] + + +def has_tracing_enabled(options: "Optional[Dict[str, Any]]") -> bool: + """ + Returns True if either traces_sample_rate or traces_sampler is + defined and enable_tracing is set and not false. + """ + if options is None: + return False + + return bool( + options.get("enable_tracing") is not False + and ( + options.get("traces_sample_rate") is not None + or options.get("traces_sampler") is not None + ) + ) + + +def has_span_streaming_enabled(options: "Optional[dict[str, Any]]") -> bool: + if options is None: + return False + + return (options.get("_experiments") or {}).get("trace_lifecycle") == "stream" + + +@contextlib.contextmanager +def record_sql_queries( + cursor: "Any", + query: "Any", + params_list: "Any", + paramstyle: "Optional[str]", + executemany: bool, + record_cursor_repr: bool = False, + span_origin: str = "manual", +) -> "Generator[sentry_sdk.tracing.Span, None, None]": + # TODO: Bring back capturing of params by default + if sentry_sdk.get_client().options["_experiments"].get("record_sql_params", False): + if not params_list or params_list == [None]: + params_list = None + + if paramstyle == "pyformat": + paramstyle = "format" + else: + params_list = None + paramstyle = None + + query = _format_sql(cursor, query) + + data = {} + if params_list is not None: + data["db.params"] = params_list + if paramstyle is not None: + data["db.paramstyle"] = paramstyle + if executemany: + data["db.executemany"] = True + if record_cursor_repr and cursor is not None: + data["db.cursor"] = cursor + + with capture_internal_exceptions(): + sentry_sdk.add_breadcrumb(message=query, category="query", data=data) + + with sentry_sdk.start_span( + op=OP.DB, + name=query, + origin=span_origin, + ) as span: + for k, v in data.items(): + span.set_data(k, v) + yield span + + +def maybe_create_breadcrumbs_from_span( + scope: "sentry_sdk.Scope", span: "sentry_sdk.tracing.Span" +) -> None: + if span.op == OP.DB_REDIS: + scope.add_breadcrumb( + message=span.description, type="redis", category="redis", data=span._tags + ) + + elif span.op == OP.HTTP_CLIENT: + level = None + status_code = span._data.get(SPANDATA.HTTP_STATUS_CODE) + if status_code: + if 500 <= status_code <= 599: + level = "error" + elif 400 <= status_code <= 499: + level = "warning" + + if level: + scope.add_breadcrumb( + type="http", category="httplib", data=span._data, level=level + ) + else: + scope.add_breadcrumb(type="http", category="httplib", data=span._data) + + elif span.op == "subprocess": + scope.add_breadcrumb( + type="subprocess", + category="subprocess", + message=span.description, + data=span._data, + ) + + +def _get_frame_module_abs_path(frame: "FrameType") -> "Optional[str]": + try: + return frame.f_code.co_filename + except Exception: + return None + + +def _should_be_included( + is_sentry_sdk_frame: bool, + namespace: "Optional[str]", + in_app_include: "Optional[list[str]]", + in_app_exclude: "Optional[list[str]]", + abs_path: "Optional[str]", + project_root: "Optional[str]", +) -> bool: + # in_app_include takes precedence over in_app_exclude + should_be_included = _module_in_list(namespace, in_app_include) + should_be_excluded = _is_external_source(abs_path) or _module_in_list( + namespace, in_app_exclude + ) + return not is_sentry_sdk_frame and ( + should_be_included + or (_is_in_project_root(abs_path, project_root) and not should_be_excluded) + ) + + +def add_source( + span: "sentry_sdk.tracing.Span", + project_root: "Optional[str]", + in_app_include: "Optional[list[str]]", + in_app_exclude: "Optional[list[str]]", +) -> None: + """ + Adds OTel compatible source code information to the span + """ + # Find the correct frame + frame: "Union[FrameType, None]" = sys._getframe() + while frame is not None: + abs_path = _get_frame_module_abs_path(frame) + + try: + namespace: "Optional[str]" = frame.f_globals.get("__name__") + except Exception: + namespace = None + + is_sentry_sdk_frame = namespace is not None and namespace.startswith( + "sentry_sdk." + ) + + should_be_included = _should_be_included( + is_sentry_sdk_frame=is_sentry_sdk_frame, + namespace=namespace, + in_app_include=in_app_include, + in_app_exclude=in_app_exclude, + abs_path=abs_path, + project_root=project_root, + ) + if should_be_included: + break + + frame = frame.f_back + else: + frame = None + + # Set the data + if frame is not None: + try: + lineno = frame.f_lineno + except Exception: + lineno = None + if lineno is not None: + span.set_data(SPANDATA.CODE_LINENO, frame.f_lineno) + + try: + namespace = frame.f_globals.get("__name__") + except Exception: + namespace = None + if namespace is not None: + span.set_data(SPANDATA.CODE_NAMESPACE, namespace) + + filepath = _get_frame_module_abs_path(frame) + if filepath is not None: + if namespace is not None: + in_app_path = filename_for_module(namespace, filepath) + elif project_root is not None and filepath.startswith(project_root): + in_app_path = filepath.replace(project_root, "").lstrip(os.sep) + else: + in_app_path = filepath + span.set_data(SPANDATA.CODE_FILEPATH, in_app_path) + + try: + code_function = frame.f_code.co_name + except Exception: + code_function = None + + if code_function is not None: + span.set_data(SPANDATA.CODE_FUNCTION, frame.f_code.co_name) + + +def add_query_source(span: "sentry_sdk.tracing.Span") -> None: + """ + Adds OTel compatible source code information to a database query span + """ + client = sentry_sdk.get_client() + if not client.is_active(): + return + + if span.timestamp is None or span.start_timestamp is None: + return + + should_add_query_source = client.options.get("enable_db_query_source", True) + if not should_add_query_source: + return + + duration = span.timestamp - span.start_timestamp + threshold = client.options.get("db_query_source_threshold_ms", 0) + slow_query = duration / timedelta(milliseconds=1) > threshold + + if not slow_query: + return + + add_source( + span=span, + project_root=client.options["project_root"], + in_app_include=client.options.get("in_app_include"), + in_app_exclude=client.options.get("in_app_exclude"), + ) + + +def add_http_request_source(span: "sentry_sdk.tracing.Span") -> None: + """ + Adds OTel compatible source code information to a span for an outgoing HTTP request + """ + client = sentry_sdk.get_client() + if not client.is_active(): + return + + if span.timestamp is None or span.start_timestamp is None: + return + + should_add_request_source = client.options.get("enable_http_request_source", True) + if not should_add_request_source: + return + + duration = span.timestamp - span.start_timestamp + threshold = client.options.get("http_request_source_threshold_ms", 0) + slow_query = duration / timedelta(milliseconds=1) > threshold + + if not slow_query: + return + + add_source( + span=span, + project_root=client.options["project_root"], + in_app_include=client.options.get("in_app_include"), + in_app_exclude=client.options.get("in_app_exclude"), + ) + + +def extract_sentrytrace_data( + header: "Optional[str]", +) -> "Optional[Dict[str, Union[str, bool, None]]]": + """ + Given a `sentry-trace` header string, return a dictionary of data. + """ + if not header: + return None + + if header.startswith("00-") and header.endswith("-00"): + header = header[3:-3] + + match = SENTRY_TRACE_REGEX.match(header) + if not match: + return None + + trace_id, parent_span_id, sampled_str = match.groups() + parent_sampled = None + + if trace_id: + trace_id = "{:032x}".format(int(trace_id, 16)) + if parent_span_id: + parent_span_id = "{:016x}".format(int(parent_span_id, 16)) + if sampled_str: + parent_sampled = sampled_str != "0" + + return { + "trace_id": trace_id, + "parent_span_id": parent_span_id, + "parent_sampled": parent_sampled, + } + + +def _format_sql(cursor: "Any", sql: str) -> "Optional[str]": + real_sql = None + + # If we're using psycopg2, it could be that we're + # looking at a query that uses Composed objects. Use psycopg2's mogrify + # function to format the query. We lose per-parameter trimming but gain + # accuracy in formatting. + try: + if hasattr(cursor, "mogrify"): + real_sql = cursor.mogrify(sql) + if isinstance(real_sql, bytes): + real_sql = real_sql.decode(cursor.connection.encoding) + except Exception: + real_sql = None + + return real_sql or to_string(sql) + + +class PropagationContext: + """ + The PropagationContext represents the data of a trace in Sentry. + """ + + __slots__ = ( + "_trace_id", + "_span_id", + "parent_span_id", + "parent_sampled", + "baggage", + ) + + def __init__( + self, + trace_id: "Optional[str]" = None, + span_id: "Optional[str]" = None, + parent_span_id: "Optional[str]" = None, + parent_sampled: "Optional[bool]" = None, + dynamic_sampling_context: "Optional[Dict[str, str]]" = None, + baggage: "Optional[Baggage]" = None, + ) -> None: + self._trace_id = trace_id + """The trace id of the Sentry trace.""" + + self._span_id = span_id + """The span id of the currently executing span.""" + + self.parent_span_id = parent_span_id + """The id of the parent span that started this span. + The parent span could also be a span in an upstream service.""" + + self.parent_sampled = parent_sampled + """Boolean indicator if the parent span was sampled. + Important when the parent span originated in an upstream service, + because we want to sample the whole trace, or nothing from the trace.""" + + self.baggage = baggage + """Parsed baggage header that is used for dynamic sampling decisions.""" + + """DEPRECATED this only exists for backwards compat of constructor.""" + if baggage is None and dynamic_sampling_context is not None: + self.baggage = Baggage(dynamic_sampling_context) + + @classmethod + def from_incoming_data( + cls, incoming_data: "Dict[str, Any]" + ) -> "PropagationContext": + propagation_context = PropagationContext() + normalized_data = normalize_incoming_data(incoming_data) + + sentry_trace_header = normalized_data.get(SENTRY_TRACE_HEADER_NAME) + sentrytrace_data = extract_sentrytrace_data(sentry_trace_header) + + # nothing to propagate if no sentry-trace + if sentrytrace_data is None: + return propagation_context + + baggage_header = normalized_data.get(BAGGAGE_HEADER_NAME) + baggage = ( + Baggage.from_incoming_header(baggage_header) if baggage_header else None + ) + + if not _should_continue_trace(baggage): + return propagation_context + + propagation_context.update(sentrytrace_data) + if baggage: + propagation_context.baggage = baggage + + propagation_context._fill_sample_rand() + + return propagation_context + + @property + def trace_id(self) -> str: + """The trace id of the Sentry trace.""" + if not self._trace_id: + # New trace, don't fill in sample_rand + self._trace_id = uuid.uuid4().hex + + return self._trace_id + + @trace_id.setter + def trace_id(self, value: str) -> None: + self._trace_id = value + + @property + def span_id(self) -> str: + """The span id of the currently executed span.""" + if not self._span_id: + self._span_id = uuid.uuid4().hex[16:] + + return self._span_id + + @span_id.setter + def span_id(self, value: str) -> None: + self._span_id = value + + @property + def dynamic_sampling_context(self) -> "Optional[Dict[str, Any]]": + return self.get_baggage().dynamic_sampling_context() + + def to_traceparent(self) -> str: + return f"{self.trace_id}-{self.span_id}" + + def get_baggage(self) -> "Baggage": + if self.baggage is None: + self.baggage = Baggage.populate_from_propagation_context(self) + return self.baggage + + def iter_headers(self) -> "Iterator[Tuple[str, str]]": + """ + Creates a generator which returns the propagation_context's ``sentry-trace`` and ``baggage`` headers. + """ + yield SENTRY_TRACE_HEADER_NAME, self.to_traceparent() + + baggage = self.get_baggage().serialize() + if baggage: + yield BAGGAGE_HEADER_NAME, baggage + + def update(self, other_dict: "Dict[str, Any]") -> None: + """ + Updates the PropagationContext with data from the given dictionary. + """ + for key, value in other_dict.items(): + try: + setattr(self, key, value) + except AttributeError: + pass + + def __repr__(self) -> str: + return "".format( + self._trace_id, + self._span_id, + self.parent_span_id, + self.parent_sampled, + self.baggage, + ) + + def _fill_sample_rand(self) -> None: + """ + Ensure that there is a valid sample_rand value in the baggage. + + If there is a valid sample_rand value in the baggage, we keep it. + Otherwise, we generate a sample_rand value according to the following: + + - If we have a parent_sampled value and a sample_rate in the DSC, we compute + a sample_rand value randomly in the range: + - [0, sample_rate) if parent_sampled is True, + - or, in the range [sample_rate, 1) if parent_sampled is False. + + - If either parent_sampled or sample_rate is missing, we generate a random + value in the range [0, 1). + + The sample_rand is deterministically generated from the trace_id, if present. + + This function does nothing if there is no baggage. + """ + if self.baggage is None: + return + + sample_rand = try_convert(float, self.baggage.sentry_items.get("sample_rand")) + if sample_rand is not None and 0 <= sample_rand < 1: + # sample_rand is present and valid, so don't overwrite it + return + + # Get the sample rate and compute the transformation that will map the random value + # to the desired range: [0, 1), [0, sample_rate), or [sample_rate, 1). + sample_rate = try_convert(float, self.baggage.sentry_items.get("sample_rate")) + lower, upper = _sample_rand_range(self.parent_sampled, sample_rate) + + try: + sample_rand = _generate_sample_rand(self.trace_id, interval=(lower, upper)) + except ValueError: + # ValueError is raised if the interval is invalid, i.e. lower >= upper. + # lower >= upper might happen if the incoming trace's sampled flag + # and sample_rate are inconsistent, e.g. sample_rate=0.0 but sampled=True. + # We cannot generate a sensible sample_rand value in this case. + logger.debug( + f"Could not backfill sample_rand, since parent_sampled={self.parent_sampled} " + f"and sample_rate={sample_rate}." + ) + return + + self.baggage.sentry_items["sample_rand"] = f"{sample_rand:.6f}" # noqa: E231 + + def _sample_rand(self) -> "Optional[str]": + """Convenience method to get the sample_rand value from the baggage.""" + if self.baggage is None: + return None + + return self.baggage.sentry_items.get("sample_rand") + + +class Baggage: + """ + The W3C Baggage header information (see https://www.w3.org/TR/baggage/). + + Before mutating a `Baggage` object, calling code must check that `mutable` is `True`. + Mutating a `Baggage` object that has `mutable` set to `False` is not allowed, but + it is the caller's responsibility to enforce this restriction. + """ + + __slots__ = ("sentry_items", "third_party_items", "mutable") + + SENTRY_PREFIX = "sentry-" + SENTRY_PREFIX_REGEX = re.compile("^sentry-") + + def __init__( + self, + sentry_items: "Dict[str, str]", + third_party_items: str = "", + mutable: bool = True, + ): + self.sentry_items = sentry_items + self.third_party_items = third_party_items + self.mutable = mutable + + @classmethod + def from_incoming_header( + cls, + header: "Optional[str]", + *, + _sample_rand: "Optional[str]" = None, + ) -> "Baggage": + """ + freeze if incoming header already has sentry baggage + """ + sentry_items = {} + third_party_items = "" + mutable = True + + if header: + for item in header.split(","): + if "=" not in item: + continue + + with capture_internal_exceptions(): + item = item.strip() + key, val = item.split("=") + if Baggage.SENTRY_PREFIX_REGEX.match(key): + baggage_key = unquote(key.split("-")[1]) + sentry_items[baggage_key] = unquote(val) + mutable = False + else: + third_party_items += ("," if third_party_items else "") + item + + if _sample_rand is not None: + sentry_items["sample_rand"] = str(_sample_rand) + mutable = False + + return Baggage(sentry_items, third_party_items, mutable) + + @classmethod + def from_options(cls, scope: "sentry_sdk.scope.Scope") -> "Optional[Baggage]": + """ + Deprecated: use populate_from_propagation_context + """ + if scope._propagation_context is None: + return Baggage({}) + + return Baggage.populate_from_propagation_context(scope._propagation_context) + + @classmethod + def populate_from_propagation_context( + cls, propagation_context: "PropagationContext" + ) -> "Baggage": + sentry_items: "Dict[str, str]" = {} + third_party_items = "" + mutable = False + + client = sentry_sdk.get_client() + + if not client.is_active(): + return Baggage(sentry_items) + + options = client.options + + sentry_items["trace_id"] = propagation_context.trace_id + + if options.get("environment"): + sentry_items["environment"] = options["environment"] + + if options.get("release"): + sentry_items["release"] = options["release"] + + if client.parsed_dsn: + sentry_items["public_key"] = client.parsed_dsn.public_key + if client.parsed_dsn.org_id: + sentry_items["org_id"] = client.parsed_dsn.org_id + + if options.get("traces_sample_rate"): + sentry_items["sample_rate"] = str(options["traces_sample_rate"]) + + return Baggage(sentry_items, third_party_items, mutable) + + @classmethod + def populate_from_transaction( + cls, transaction: "sentry_sdk.tracing.Transaction" + ) -> "Baggage": + """ + Populate fresh baggage entry with sentry_items and make it immutable + if this is the head SDK which originates traces. + """ + client = sentry_sdk.get_client() + sentry_items: "Dict[str, str]" = {} + + if not client.is_active(): + return Baggage(sentry_items) + + options = client.options or {} + + sentry_items["trace_id"] = transaction.trace_id + sentry_items["sample_rand"] = f"{transaction._sample_rand:.6f}" # noqa: E231 + + if options.get("environment"): + sentry_items["environment"] = options["environment"] + + if options.get("release"): + sentry_items["release"] = options["release"] + + if client.parsed_dsn: + sentry_items["public_key"] = client.parsed_dsn.public_key + if client.parsed_dsn.org_id: + sentry_items["org_id"] = client.parsed_dsn.org_id + + if ( + transaction.name + and transaction.source not in LOW_QUALITY_TRANSACTION_SOURCES + ): + sentry_items["transaction"] = transaction.name + + if transaction.sample_rate is not None: + sentry_items["sample_rate"] = str(transaction.sample_rate) + + if transaction.sampled is not None: + sentry_items["sampled"] = "true" if transaction.sampled else "false" + + # there's an existing baggage but it was mutable, + # which is why we are creating this new baggage. + # However, if by chance the user put some sentry items in there, give them precedence. + if transaction._baggage and transaction._baggage.sentry_items: + sentry_items.update(transaction._baggage.sentry_items) + + return Baggage(sentry_items, mutable=False) + + def freeze(self) -> None: + self.mutable = False + + def dynamic_sampling_context(self) -> "Dict[str, str]": + header = {} + + for key, item in self.sentry_items.items(): + header[key] = item + + return header + + def serialize(self, include_third_party: bool = False) -> str: + items = [] + + for key, val in self.sentry_items.items(): + with capture_internal_exceptions(): + item = Baggage.SENTRY_PREFIX + quote(key) + "=" + quote(str(val)) + items.append(item) + + if include_third_party: + items.append(self.third_party_items) + + return ",".join(items) + + @staticmethod + def strip_sentry_baggage(header: str) -> str: + """Remove Sentry baggage from the given header. + + Given a Baggage header, return a new Baggage header with all Sentry baggage items removed. + """ + return ",".join( + ( + item + for item in header.split(",") + if not Baggage.SENTRY_PREFIX_REGEX.match(item.strip()) + ) + ) + + def _sample_rand(self) -> "Optional[float]": + """Convenience method to get the sample_rand value from the sentry_items. + + We validate the value and parse it as a float before returning it. The value is considered + valid if it is a float in the range [0, 1). + """ + sample_rand = try_convert(float, self.sentry_items.get("sample_rand")) + + if sample_rand is not None and 0.0 <= sample_rand < 1.0: + return sample_rand + + return None + + def __repr__(self) -> str: + return f'' + + +def should_propagate_trace(client: "sentry_sdk.client.BaseClient", url: str) -> bool: + """ + Returns True if url matches trace_propagation_targets configured in the given client. Otherwise, returns False. + """ + trace_propagation_targets = client.options["trace_propagation_targets"] + + if is_sentry_url(client, url): + return False + + return match_regex_list(url, trace_propagation_targets, substring_matching=True) + + +def normalize_incoming_data(incoming_data: "Dict[str, Any]") -> "Dict[str, Any]": + """ + Normalizes incoming data so the keys are all lowercase with dashes instead of underscores and stripped from known prefixes. + """ + data = {} + for key, value in incoming_data.items(): + if key.startswith("HTTP_"): + key = key[5:] + + key = key.replace("_", "-").lower() + data[key] = value + + return data + + +def create_span_decorator( + op: "Optional[Union[str, OP]]" = None, + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, + template: "SPANTEMPLATE" = SPANTEMPLATE.DEFAULT, +) -> "Any": + """ + Create a span decorator that can wrap both sync and async functions. + + :param op: The operation type for the span. + :type op: str or :py:class:`sentry_sdk.consts.OP` or None + :param name: The name of the span. + :type name: str or None + :param attributes: Additional attributes to set on the span. + :type attributes: dict or None + :param template: The type of span to create. This determines what kind of + span instrumentation and data collection will be applied. Use predefined + constants from :py:class:`sentry_sdk.consts.SPANTEMPLATE`. + The default is `SPANTEMPLATE.DEFAULT` which is the right choice for most + use cases. + :type template: :py:class:`sentry_sdk.consts.SPANTEMPLATE` + """ + from sentry_sdk.scope import should_send_default_pii + + def span_decorator(f: "Any") -> "Any": + """ + Decorator to create a span for the given function. + """ + + @functools.wraps(f) + async def async_wrapper(*args: "Any", **kwargs: "Any") -> "Any": + current_span = get_current_span() + + if current_span is None: + logger.debug( + "Cannot create a child span for %s. " + "Please start a Sentry transaction before calling this function.", + qualname_from_function(f), + ) + return await f(*args, **kwargs) + + span_op = op or _get_span_op(template) + function_name = name or qualname_from_function(f) or "" + span_name = _get_span_name(template, function_name, kwargs) + send_pii = should_send_default_pii() + + with current_span.start_child( + op=span_op, + name=span_name, + ) as span: + span.update_data(attributes or {}) + _set_input_attributes( + span, template, send_pii, function_name, f, args, kwargs + ) + + result = await f(*args, **kwargs) + + _set_output_attributes(span, template, send_pii, result) + + return result + + try: + async_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + except Exception: + pass + + @functools.wraps(f) + def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any": + current_span = get_current_span() + + if current_span is None: + logger.debug( + "Cannot create a child span for %s. " + "Please start a Sentry transaction before calling this function.", + qualname_from_function(f), + ) + return f(*args, **kwargs) + + span_op = op or _get_span_op(template) + function_name = name or qualname_from_function(f) or "" + span_name = _get_span_name(template, function_name, kwargs) + send_pii = should_send_default_pii() + + with current_span.start_child( + op=span_op, + name=span_name, + ) as span: + span.update_data(attributes or {}) + _set_input_attributes( + span, template, send_pii, function_name, f, args, kwargs + ) + + result = f(*args, **kwargs) + + _set_output_attributes(span, template, send_pii, result) + + return result + + try: + sync_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + except Exception: + pass + + if inspect.iscoroutinefunction(f): + return async_wrapper + else: + return sync_wrapper + + return span_decorator + + +def get_current_span(scope: "Optional[sentry_sdk.Scope]" = None) -> "Optional[Span]": + """ + Returns the currently active span if there is one running, otherwise `None` + """ + scope = scope or sentry_sdk.get_current_scope() + current_span = scope.span + return current_span + + +def set_span_errored(span: "Optional[Span]" = None) -> None: + """ + Set the status of the current or given span to INTERNAL_ERROR. + Also sets the status of the transaction (root span) to INTERNAL_ERROR. + """ + span = span or get_current_span() + if span is not None: + span.set_status(SPANSTATUS.INTERNAL_ERROR) + if span.containing_transaction is not None: + span.containing_transaction.set_status(SPANSTATUS.INTERNAL_ERROR) + + +def _generate_sample_rand( + trace_id: "Optional[str]", + *, + interval: "tuple[float, float]" = (0.0, 1.0), +) -> float: + """Generate a sample_rand value from a trace ID. + + The generated value will be pseudorandomly chosen from the provided + interval. Specifically, given (lower, upper) = interval, the generated + value will be in the range [lower, upper). The value has 6-digit precision, + so when printing with .6f, the value will never be rounded up. + + The pseudorandom number generator is seeded with the trace ID. + """ + lower, upper = interval + if not lower < upper: # using `if lower >= upper` would handle NaNs incorrectly + raise ValueError("Invalid interval: lower must be less than upper") + + rng = Random(trace_id) + lower_scaled = int(lower * 1_000_000) + upper_scaled = int(upper * 1_000_000) + try: + sample_rand_scaled = rng.randrange(lower_scaled, upper_scaled) + except ValueError: + # In some corner cases it might happen that the range is too small + # In that case, just take the lower bound + sample_rand_scaled = lower_scaled + + return sample_rand_scaled / 1_000_000 + + +def _sample_rand_range( + parent_sampled: "Optional[bool]", sample_rate: "Optional[float]" +) -> "tuple[float, float]": + """ + Compute the lower (inclusive) and upper (exclusive) bounds of the range of values + that a generated sample_rand value must fall into, given the parent_sampled and + sample_rate values. + """ + if parent_sampled is None or sample_rate is None: + return 0.0, 1.0 + elif parent_sampled is True: + return 0.0, sample_rate + else: # parent_sampled is False + return sample_rate, 1.0 + + +def _get_value(source: "Any", key: str) -> "Optional[Any]": + """ + Gets a value from a source object. The source can be a dict or an object. + It is checked for dictionary keys and object attributes. + """ + value = None + if isinstance(source, dict): + value = source.get(key) + else: + if hasattr(source, key): + try: + value = getattr(source, key) + except Exception: + value = None + return value + + +def _get_span_name( + template: "Union[str, SPANTEMPLATE]", + name: str, + kwargs: "Optional[dict[str, Any]]" = None, +) -> str: + """ + Get the name of the span based on the template and the name. + """ + span_name = name + + if template == SPANTEMPLATE.AI_CHAT: + model = None + if kwargs: + for key in ("model", "model_name"): + if kwargs.get(key) and isinstance(kwargs[key], str): + model = kwargs[key] + break + + span_name = f"chat {model}" if model else "chat" + + elif template == SPANTEMPLATE.AI_AGENT: + span_name = f"invoke_agent {name}" + + elif template == SPANTEMPLATE.AI_TOOL: + span_name = f"execute_tool {name}" + + return span_name + + +def _get_span_op(template: "Union[str, SPANTEMPLATE]") -> str: + """ + Get the operation of the span based on the template. + """ + mapping: "dict[Union[str, SPANTEMPLATE], Union[str, OP]]" = { + SPANTEMPLATE.AI_CHAT: OP.GEN_AI_CHAT, + SPANTEMPLATE.AI_AGENT: OP.GEN_AI_INVOKE_AGENT, + SPANTEMPLATE.AI_TOOL: OP.GEN_AI_EXECUTE_TOOL, + } + op = mapping.get(template, OP.FUNCTION) + + return str(op) + + +def _get_input_attributes( + template: "Union[str, SPANTEMPLATE]", + send_pii: bool, + args: "tuple[Any, ...]", + kwargs: "dict[str, Any]", +) -> "dict[str, Any]": + """ + Get input attributes for the given span template. + """ + attributes: "dict[str, Any]" = {} + + if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]: + mapping = { + "model": (SPANDATA.GEN_AI_REQUEST_MODEL, str), + "model_name": (SPANDATA.GEN_AI_REQUEST_MODEL, str), + "agent": (SPANDATA.GEN_AI_AGENT_NAME, str), + "agent_name": (SPANDATA.GEN_AI_AGENT_NAME, str), + "max_tokens": (SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, int), + "frequency_penalty": (SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, float), + "presence_penalty": (SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, float), + "temperature": (SPANDATA.GEN_AI_REQUEST_TEMPERATURE, float), + "top_p": (SPANDATA.GEN_AI_REQUEST_TOP_P, float), + "top_k": (SPANDATA.GEN_AI_REQUEST_TOP_K, int), + } + + def _set_from_key(key: str, value: "Any") -> None: + if key in mapping: + (attribute, data_type) = mapping[key] + if value is not None and isinstance(value, data_type): + attributes[attribute] = value + + for key, value in list(kwargs.items()): + if key == "prompt" and isinstance(value, str): + attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append( + {"role": "user", "content": value} + ) + continue + + if key == "system_prompt" and isinstance(value, str): + attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append( + {"role": "system", "content": value} + ) + continue + + _set_from_key(key, value) + + if template == SPANTEMPLATE.AI_TOOL and send_pii: + attributes[SPANDATA.GEN_AI_TOOL_INPUT] = safe_repr( + {"args": args, "kwargs": kwargs} + ) + + # Coerce to string + if SPANDATA.GEN_AI_REQUEST_MESSAGES in attributes: + attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES] = safe_repr( + attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) + + return attributes + + +def _get_usage_attributes(usage: "Any") -> "dict[str, Any]": + """ + Get usage attributes. + """ + attributes = {} + + def _set_from_keys(attribute: str, keys: "tuple[str, ...]") -> None: + for key in keys: + value = _get_value(usage, key) + if value is not None and isinstance(value, int): + attributes[attribute] = value + + _set_from_keys( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, + ("prompt_tokens", "input_tokens"), + ) + _set_from_keys( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + ("completion_tokens", "output_tokens"), + ) + _set_from_keys( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, + ("total_tokens",), + ) + + return attributes + + +def _get_output_attributes( + template: "Union[str, SPANTEMPLATE]", send_pii: bool, result: "Any" +) -> "dict[str, Any]": + """ + Get output attributes for the given span template. + """ + attributes: "dict[str, Any]" = {} + + if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]: + with capture_internal_exceptions(): + # Usage from result, result.usage, and result.metadata.usage + usage_candidates = [result] + + usage = _get_value(result, "usage") + usage_candidates.append(usage) + + meta = _get_value(result, "metadata") + usage = _get_value(meta, "usage") + usage_candidates.append(usage) + + for usage_candidate in usage_candidates: + if usage_candidate is not None: + attributes.update(_get_usage_attributes(usage_candidate)) + + # Response model + model_name = _get_value(result, "model") + if model_name is not None and isinstance(model_name, str): + attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = model_name + + model_name = _get_value(result, "model_name") + if model_name is not None and isinstance(model_name, str): + attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = model_name + + # Tool output + if template == SPANTEMPLATE.AI_TOOL and send_pii: + attributes[SPANDATA.GEN_AI_TOOL_OUTPUT] = safe_repr(result) + + return attributes + + +def _set_input_attributes( + span: "Span", + template: "Union[str, SPANTEMPLATE]", + send_pii: bool, + name: str, + f: "Any", + args: "tuple[Any, ...]", + kwargs: "dict[str, Any]", +) -> None: + """ + Set span input attributes based on the given span template. + + :param span: The span to set attributes on. + :param template: The template to use to set attributes on the span. + :param send_pii: Whether to send PII data. + :param f: The wrapped function. + :param args: The arguments to the wrapped function. + :param kwargs: The keyword arguments to the wrapped function. + """ + attributes: "dict[str, Any]" = {} + + if template == SPANTEMPLATE.AI_AGENT: + attributes = { + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + SPANDATA.GEN_AI_AGENT_NAME: name, + } + elif template == SPANTEMPLATE.AI_CHAT: + attributes = { + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + } + elif template == SPANTEMPLATE.AI_TOOL: + attributes = { + SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", + SPANDATA.GEN_AI_TOOL_NAME: name, + } + + docstring = f.__doc__ + if docstring is not None: + attributes[SPANDATA.GEN_AI_TOOL_DESCRIPTION] = docstring + + attributes.update(_get_input_attributes(template, send_pii, args, kwargs)) + span.update_data(attributes or {}) + + +def _set_output_attributes( + span: "Span", template: "Union[str, SPANTEMPLATE]", send_pii: bool, result: "Any" +) -> None: + """ + Set span output attributes based on the given span template. + + :param span: The span to set attributes on. + :param template: The template to use to set attributes on the span. + :param send_pii: Whether to send PII data. + :param result: The result of the wrapped function. + """ + span.update_data(_get_output_attributes(template, send_pii, result) or {}) + + +def _should_continue_trace(baggage: "Optional[Baggage]") -> bool: + """ + Check if we should continue the incoming trace according to the strict_trace_continuation spec. + https://develop.sentry.dev/sdk/telemetry/traces/#stricttracecontinuation + """ + + client = sentry_sdk.get_client() + parsed_dsn = client.parsed_dsn + client_org_id = parsed_dsn.org_id if parsed_dsn else None + baggage_org_id = baggage.sentry_items.get("org_id") if baggage else None + + if ( + client_org_id is not None + and baggage_org_id is not None + and client_org_id != baggage_org_id + ): + logger.debug( + f"Starting a new trace because org IDs don't match (incoming baggage org_id: {baggage_org_id}, SDK org_id: {client_org_id})" + ) + return False + + strict_trace_continuation: bool = client.options.get( + "strict_trace_continuation", False + ) + if strict_trace_continuation: + if (baggage_org_id is not None and client_org_id is None) or ( + baggage_org_id is None and client_org_id is not None + ): + logger.debug( + f"Starting a new trace because strict trace continuation is enabled and one org ID is missing (incoming baggage org_id: {baggage_org_id}, SDK org_id: {client_org_id})" + ) + return False + + return True + + +def add_sentry_baggage_to_headers( + headers: "MutableMapping[str, str]", sentry_baggage: str +) -> None: + """Add the Sentry baggage to the headers. + + This function directly mutates the provided headers. The provided sentry_baggage + is appended to the existing baggage. If the baggage already contains Sentry items, + they are stripped out first. + """ + existing_baggage = headers.get(BAGGAGE_HEADER_NAME, "") + stripped_existing_baggage = Baggage.strip_sentry_baggage(existing_baggage) + + separator = "," if len(stripped_existing_baggage) > 0 else "" + + headers[BAGGAGE_HEADER_NAME] = ( + stripped_existing_baggage + separator + sentry_baggage + ) + + +# Circular imports +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + LOW_QUALITY_TRANSACTION_SOURCES, + SENTRY_TRACE_HEADER_NAME, +) + +if TYPE_CHECKING: + from sentry_sdk.tracing import Span diff --git a/lib/python3.12/site-packages/sentry_sdk/transport.py b/lib/python3.12/site-packages/sentry_sdk/transport.py new file mode 100644 index 0000000000000000000000000000000000000000..dcfe55406b645424a0c802b45c1588ef3f17119d --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/transport.py @@ -0,0 +1,888 @@ +from abc import ABC, abstractmethod +import io +import os +import gzip +import socket +import ssl +import time +import warnings +from datetime import datetime, timedelta, timezone +from collections import defaultdict +from urllib.request import getproxies + +try: + import brotli # type: ignore +except ImportError: + brotli = None + +import urllib3 +import certifi + +import sentry_sdk +from sentry_sdk.consts import EndpointType +from sentry_sdk.utils import Dsn, logger, capture_internal_exceptions +from sentry_sdk.worker import BackgroundWorker +from sentry_sdk.envelope import Envelope, Item, PayloadRef + +from typing import TYPE_CHECKING, cast, List, Dict + +if TYPE_CHECKING: + from typing import Any + from typing import Callable + from typing import DefaultDict + from typing import Iterable + from typing import Mapping + from typing import Optional + from typing import Self + from typing import Tuple + from typing import Type + from typing import Union + + from urllib3.poolmanager import PoolManager + from urllib3.poolmanager import ProxyManager + + from sentry_sdk._types import Event, EventDataCategory + +KEEP_ALIVE_SOCKET_OPTIONS = [] +for option in [ + (socket.SOL_SOCKET, lambda: getattr(socket, "SO_KEEPALIVE"), 1), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPIDLE"), 45), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPINTVL"), 10), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPCNT"), 6), # noqa: B009 +]: + try: + KEEP_ALIVE_SOCKET_OPTIONS.append((option[0], option[1](), option[2])) + except AttributeError: + # a specific option might not be available on specific systems, + # e.g. TCP_KEEPIDLE doesn't exist on macOS + pass + + +class Transport(ABC): + """Baseclass for all transports. + + A transport is used to send an event to sentry. + """ + + parsed_dsn: "Optional[Dsn]" = None + + def __init__(self: "Self", options: "Optional[Dict[str, Any]]" = None) -> None: + self.options = options + if options and options["dsn"] is not None and options["dsn"]: + self.parsed_dsn = Dsn(options["dsn"], options.get("org_id")) + else: + self.parsed_dsn = None + + def capture_event(self: "Self", event: "Event") -> None: + """ + DEPRECATED: Please use capture_envelope instead. + + This gets invoked with the event dictionary when an event should + be sent to sentry. + """ + + warnings.warn( + "capture_event is deprecated, please use capture_envelope instead!", + DeprecationWarning, + stacklevel=2, + ) + + envelope = Envelope() + envelope.add_event(event) + self.capture_envelope(envelope) + + @abstractmethod + def capture_envelope(self: "Self", envelope: "Envelope") -> None: + """ + Send an envelope to Sentry. + + Envelopes are a data container format that can hold any type of data + submitted to Sentry. We use it to send all event data (including errors, + transactions, crons check-ins, etc.) to Sentry. + """ + pass + + def flush( + self: "Self", + timeout: float, + callback: "Optional[Any]" = None, + ) -> None: + """ + Wait `timeout` seconds for the current events to be sent out. + + The default implementation is a no-op, since this method may only be relevant to some transports. + Subclasses should override this method if necessary. + """ + return None + + def kill(self: "Self") -> None: + """ + Forcefully kills the transport. + + The default implementation is a no-op, since this method may only be relevant to some transports. + Subclasses should override this method if necessary. + """ + return None + + def record_lost_event( + self, + reason: str, + data_category: "Optional[EventDataCategory]" = None, + item: "Optional[Item]" = None, + *, + quantity: int = 1, + ) -> None: + """This increments a counter for event loss by reason and + data category by the given positive-int quantity (default 1). + + If an item is provided, the data category and quantity are + extracted from the item, and the values passed for + data_category and quantity are ignored. + + When recording a lost transaction via data_category="transaction", + the calling code should also record the lost spans via this method. + When recording lost spans, `quantity` should be set to the number + of contained spans, plus one for the transaction itself. When + passing an Item containing a transaction via the `item` parameter, + this method automatically records the lost spans. + """ + return None + + def is_healthy(self: "Self") -> bool: + return True + + +def _parse_rate_limits( + header: str, now: "Optional[datetime]" = None +) -> "Iterable[Tuple[Optional[EventDataCategory], datetime]]": + if now is None: + now = datetime.now(timezone.utc) + + for limit in header.split(","): + try: + parameters = limit.strip().split(":") + retry_after_val, categories = parameters[:2] + + retry_after = now + timedelta(seconds=int(retry_after_val)) + for category in categories and categories.split(";") or (None,): + yield category, retry_after # type: ignore + except (LookupError, ValueError): + continue + + +class BaseHttpTransport(Transport): + """The base HTTP transport.""" + + TIMEOUT = 30 # seconds + + def __init__(self: "Self", options: "Dict[str, Any]") -> None: + from sentry_sdk.consts import VERSION + + Transport.__init__(self, options) + assert self.parsed_dsn is not None + self.options: "Dict[str, Any]" = options + self._worker = BackgroundWorker(queue_size=options["transport_queue_size"]) + self._auth = self.parsed_dsn.to_auth("sentry.python/%s" % VERSION) + self._disabled_until: "Dict[Optional[EventDataCategory], datetime]" = {} + # We only use this Retry() class for the `get_retry_after` method it exposes + self._retry = urllib3.util.Retry() + self._discarded_events: "DefaultDict[Tuple[EventDataCategory, str], int]" = ( + defaultdict(int) + ) + self._last_client_report_sent = time.time() + + self._pool = self._make_pool() + + # Backwards compatibility for deprecated `self.hub_class` attribute + self._hub_cls = sentry_sdk.Hub + + experiments = options.get("_experiments", {}) + compression_level = experiments.get( + "transport_compression_level", + experiments.get("transport_zlib_compression_level"), + ) + compression_algo = experiments.get( + "transport_compression_algo", + ( + "gzip" + # if only compression level is set, assume gzip for backwards compatibility + # if we don't have brotli available, fallback to gzip + if compression_level is not None or brotli is None + else "br" + ), + ) + + if compression_algo == "br" and brotli is None: + logger.warning( + "You asked for brotli compression without the Brotli module, falling back to gzip -9" + ) + compression_algo = "gzip" + compression_level = None + + if compression_algo not in ("br", "gzip"): + logger.warning( + "Unknown compression algo %s, disabling compression", compression_algo + ) + self._compression_level = 0 + self._compression_algo = None + else: + self._compression_algo = compression_algo + + if compression_level is not None: + self._compression_level = compression_level + elif self._compression_algo == "gzip": + self._compression_level = 9 + elif self._compression_algo == "br": + self._compression_level = 4 + + def record_lost_event( + self, + reason: str, + data_category: "Optional[EventDataCategory]" = None, + item: "Optional[Item]" = None, + *, + quantity: int = 1, + ) -> None: + if not self.options["send_client_reports"]: + return + + if item is not None: + data_category = item.data_category + quantity = 1 # If an item is provided, we always count it as 1 (except for attachments, handled below). + + if data_category == "transaction": + # Also record the lost spans + event = item.get_transaction_event() or {} + + # +1 for the transaction itself + span_count = ( + len(cast(List[Dict[str, object]], event.get("spans") or [])) + 1 + ) + self.record_lost_event(reason, "span", quantity=span_count) + + elif data_category == "log_item" and item: + # Also record size of lost logs in bytes + bytes_size = len(item.get_bytes()) + self.record_lost_event(reason, "log_byte", quantity=bytes_size) + + elif data_category == "attachment": + # quantity of 0 is actually 1 as we do not want to count + # empty attachments as actually empty. + quantity = len(item.get_bytes()) or 1 + + elif data_category is None: + raise TypeError("data category not provided") + + self._discarded_events[data_category, reason] += quantity + + def _get_header_value( + self: "Self", response: "Any", header: str + ) -> "Optional[str]": + return response.headers.get(header) + + def _update_rate_limits( + self: "Self", response: "Union[urllib3.BaseHTTPResponse, httpcore.Response]" + ) -> None: + # new sentries with more rate limit insights. We honor this header + # no matter of the status code to update our internal rate limits. + header = self._get_header_value(response, "x-sentry-rate-limits") + if header: + logger.warning("Rate-limited via x-sentry-rate-limits") + self._disabled_until.update(_parse_rate_limits(header)) + + # old sentries only communicate global rate limit hits via the + # retry-after header on 429. This header can also be emitted on new + # sentries if a proxy in front wants to globally slow things down. + elif response.status == 429: + logger.warning("Rate-limited via 429") + retry_after_value = self._get_header_value(response, "Retry-After") + retry_after = ( + self._retry.parse_retry_after(retry_after_value) + if retry_after_value is not None + else None + ) or 60 + self._disabled_until[None] = datetime.now(timezone.utc) + timedelta( + seconds=retry_after + ) + + def _send_request( + self: "Self", + body: bytes, + headers: "Dict[str, str]", + endpoint_type: "EndpointType" = EndpointType.ENVELOPE, + envelope: "Optional[Envelope]" = None, + ) -> None: + def record_loss(reason: str) -> None: + if envelope is None: + self.record_lost_event(reason, data_category="error") + else: + for item in envelope.items: + self.record_lost_event(reason, item=item) + + headers.update( + { + "User-Agent": str(self._auth.client), + "X-Sentry-Auth": str(self._auth.to_header()), + } + ) + try: + response = self._request( + "POST", + endpoint_type, + body, + headers, + ) + except Exception: + self.on_dropped_event("network") + record_loss("network_error") + raise + + try: + self._update_rate_limits(response) + + if response.status == 413: + size_exceeded_message = ( + "HTTP 413: Event dropped due to exceeded envelope size limit" + ) + response_message = getattr( + response, "data", getattr(response, "content", None) + ) + if response_message is not None: + size_exceeded_message += f" (body: {response_message})" + + logger.error(size_exceeded_message) + self.on_dropped_event("status_413") + record_loss("send_error") + + elif response.status == 429: + # if we hit a 429. Something was rate limited but we already + # acted on this in `self._update_rate_limits`. Note that we + # do not want to record event loss here as we will have recorded + # an outcome in relay already. + self.on_dropped_event("status_429") + pass + + elif response.status >= 300 or response.status < 200: + logger.error( + "Unexpected status code: %s (body: %s)", + response.status, + getattr(response, "data", getattr(response, "content", None)), + ) + self.on_dropped_event("status_{}".format(response.status)) + record_loss("network_error") + finally: + response.close() + + def on_dropped_event(self: "Self", _reason: str) -> None: + return None + + def _fetch_pending_client_report( + self: "Self", force: bool = False, interval: int = 60 + ) -> "Optional[Item]": + if not self.options["send_client_reports"]: + return None + + if not (force or self._last_client_report_sent < time.time() - interval): + return None + + discarded_events = self._discarded_events + self._discarded_events = defaultdict(int) + self._last_client_report_sent = time.time() + + if not discarded_events: + return None + + return Item( + PayloadRef( + json={ + "timestamp": time.time(), + "discarded_events": [ + {"reason": reason, "category": category, "quantity": quantity} + for ( + (category, reason), + quantity, + ) in discarded_events.items() + ], + } + ), + type="client_report", + ) + + def _flush_client_reports(self: "Self", force: bool = False) -> None: + client_report = self._fetch_pending_client_report(force=force, interval=60) + if client_report is not None: + self.capture_envelope(Envelope(items=[client_report])) + + def _check_disabled(self, category: str) -> bool: + def _disabled(bucket: "Any") -> bool: + ts = self._disabled_until.get(bucket) + return ts is not None and ts > datetime.now(timezone.utc) + + return _disabled(category) or _disabled(None) + + def _is_rate_limited(self: "Self") -> bool: + return any( + ts > datetime.now(timezone.utc) for ts in self._disabled_until.values() + ) + + def _is_worker_full(self: "Self") -> bool: + return self._worker.full() + + def is_healthy(self: "Self") -> bool: + return not (self._is_worker_full() or self._is_rate_limited()) + + def _send_envelope(self: "Self", envelope: "Envelope") -> None: + # remove all items from the envelope which are over quota + new_items = [] + for item in envelope.items: + if self._check_disabled(item.data_category): + if item.data_category in ("transaction", "error", "default", "statsd"): + self.on_dropped_event("self_rate_limits") + self.record_lost_event("ratelimit_backoff", item=item) + else: + new_items.append(item) + + # Since we're modifying the envelope here make a copy so that others + # that hold references do not see their envelope modified. + envelope = Envelope(headers=envelope.headers, items=new_items) + + if not envelope.items: + return None + + # since we're already in the business of sending out an envelope here + # check if we have one pending for the stats session envelopes so we + # can attach it to this enveloped scheduled for sending. This will + # currently typically attach the client report to the most recent + # session update. + client_report_item = self._fetch_pending_client_report(interval=30) + if client_report_item is not None: + envelope.items.append(client_report_item) + + content_encoding, body = self._serialize_envelope(envelope) + + assert self.parsed_dsn is not None + logger.debug( + "Sending envelope [%s] project:%s host:%s", + envelope.description, + self.parsed_dsn.project_id, + self.parsed_dsn.host, + ) + + headers = { + "Content-Type": "application/x-sentry-envelope", + } + if content_encoding: + headers["Content-Encoding"] = content_encoding + + self._send_request( + body.getvalue(), + headers=headers, + endpoint_type=EndpointType.ENVELOPE, + envelope=envelope, + ) + return None + + def _serialize_envelope( + self: "Self", envelope: "Envelope" + ) -> "tuple[Optional[str], io.BytesIO]": + content_encoding = None + body = io.BytesIO() + if self._compression_level == 0 or self._compression_algo is None: + envelope.serialize_into(body) + else: + content_encoding = self._compression_algo + if self._compression_algo == "br" and brotli is not None: + body.write( + brotli.compress( + envelope.serialize(), quality=self._compression_level + ) + ) + else: # assume gzip as we sanitize the algo value in init + with gzip.GzipFile( + fileobj=body, mode="w", compresslevel=self._compression_level + ) as f: + envelope.serialize_into(f) + + return content_encoding, body + + def _get_pool_options(self: "Self") -> "Dict[str, Any]": + raise NotImplementedError() + + def _in_no_proxy(self: "Self", parsed_dsn: "Dsn") -> bool: + no_proxy = getproxies().get("no") + if not no_proxy: + return False + for host in no_proxy.split(","): + host = host.strip() + if parsed_dsn.host.endswith(host) or parsed_dsn.netloc.endswith(host): + return True + return False + + def _make_pool( + self: "Self", + ) -> "Union[PoolManager, ProxyManager, httpcore.SOCKSProxy, httpcore.HTTPProxy, httpcore.ConnectionPool]": + raise NotImplementedError() + + def _request( + self: "Self", + method: str, + endpoint_type: "EndpointType", + body: "Any", + headers: "Mapping[str, str]", + ) -> "Union[urllib3.BaseHTTPResponse, httpcore.Response]": + raise NotImplementedError() + + def capture_envelope( + self, + envelope: "Envelope", + ) -> None: + def send_envelope_wrapper() -> None: + with capture_internal_exceptions(): + self._send_envelope(envelope) + self._flush_client_reports() + + if not self._worker.submit(send_envelope_wrapper): + self.on_dropped_event("full_queue") + for item in envelope.items: + self.record_lost_event("queue_overflow", item=item) + + def flush( + self: "Self", + timeout: float, + callback: "Optional[Callable[[int, float], None]]" = None, + ) -> None: + logger.debug("Flushing HTTP transport") + + if timeout > 0: + self._worker.submit(lambda: self._flush_client_reports(force=True)) + self._worker.flush(timeout, callback) + + def kill(self: "Self") -> None: + logger.debug("Killing HTTP transport") + self._worker.kill() + + @staticmethod + def _warn_hub_cls() -> None: + """Convenience method to warn users about the deprecation of the `hub_cls` attribute.""" + warnings.warn( + "The `hub_cls` attribute is deprecated and will be removed in a future release.", + DeprecationWarning, + stacklevel=3, + ) + + @property + def hub_cls(self: "Self") -> "type[sentry_sdk.Hub]": + """DEPRECATED: This attribute is deprecated and will be removed in a future release.""" + HttpTransport._warn_hub_cls() + return self._hub_cls + + @hub_cls.setter + def hub_cls(self: "Self", value: "type[sentry_sdk.Hub]") -> None: + """DEPRECATED: This attribute is deprecated and will be removed in a future release.""" + HttpTransport._warn_hub_cls() + self._hub_cls = value + + +class HttpTransport(BaseHttpTransport): + if TYPE_CHECKING: + _pool: "Union[PoolManager, ProxyManager]" + + def _get_pool_options(self: "Self") -> "Dict[str, Any]": + num_pools = self.options.get("_experiments", {}).get("transport_num_pools") + options = { + "num_pools": 2 if num_pools is None else int(num_pools), + "cert_reqs": "CERT_REQUIRED", + "timeout": urllib3.Timeout(total=self.TIMEOUT), + } + + socket_options: "Optional[List[Tuple[int, int, int | bytes]]]" = None + + if self.options["socket_options"] is not None: + socket_options = self.options["socket_options"] + + if self.options["keep_alive"]: + if socket_options is None: + socket_options = [] + + used_options = {(o[0], o[1]) for o in socket_options} + for default_option in KEEP_ALIVE_SOCKET_OPTIONS: + if (default_option[0], default_option[1]) not in used_options: + socket_options.append(default_option) + + if socket_options is not None: + options["socket_options"] = socket_options + + options["ca_certs"] = ( + self.options["ca_certs"] # User-provided bundle from the SDK init + or os.environ.get("SSL_CERT_FILE") + or os.environ.get("REQUESTS_CA_BUNDLE") + or certifi.where() + ) + + options["cert_file"] = self.options["cert_file"] or os.environ.get( + "CLIENT_CERT_FILE" + ) + options["key_file"] = self.options["key_file"] or os.environ.get( + "CLIENT_KEY_FILE" + ) + + return options + + def _make_pool(self: "Self") -> "Union[PoolManager, ProxyManager]": + if self.parsed_dsn is None: + raise ValueError("Cannot create HTTP-based transport without valid DSN") + + proxy = None + no_proxy = self._in_no_proxy(self.parsed_dsn) + + # try HTTPS first + https_proxy = self.options["https_proxy"] + if self.parsed_dsn.scheme == "https" and (https_proxy != ""): + proxy = https_proxy or (not no_proxy and getproxies().get("https")) + + # maybe fallback to HTTP proxy + http_proxy = self.options["http_proxy"] + if not proxy and (http_proxy != ""): + proxy = http_proxy or (not no_proxy and getproxies().get("http")) + + opts = self._get_pool_options() + + if proxy: + proxy_headers = self.options["proxy_headers"] + if proxy_headers: + opts["proxy_headers"] = proxy_headers + + if proxy.startswith("socks"): + use_socks_proxy = True + try: + # Check if PySocks dependency is available + from urllib3.contrib.socks import SOCKSProxyManager + except ImportError: + use_socks_proxy = False + logger.warning( + "You have configured a SOCKS proxy (%s) but support for SOCKS proxies is not installed. Disabling proxy support. Please add `PySocks` (or `urllib3` with the `[socks]` extra) to your dependencies.", + proxy, + ) + + if use_socks_proxy: + return SOCKSProxyManager(proxy, **opts) + else: + return urllib3.PoolManager(**opts) + else: + return urllib3.ProxyManager(proxy, **opts) + else: + return urllib3.PoolManager(**opts) + + def _request( + self: "Self", + method: str, + endpoint_type: "EndpointType", + body: "Any", + headers: "Mapping[str, str]", + ) -> "urllib3.BaseHTTPResponse": + return self._pool.request( + method, + self._auth.get_api_url(endpoint_type), + body=body, + headers=headers, + ) + + +try: + import httpcore + import h2 # noqa: F401 +except ImportError: + # Sorry, no Http2Transport for you + class Http2Transport(HttpTransport): + def __init__(self: "Self", options: "Dict[str, Any]") -> None: + super().__init__(options) + logger.warning( + "You tried to use HTTP2Transport but don't have httpcore[http2] installed. Falling back to HTTPTransport." + ) + +else: + + class Http2Transport(BaseHttpTransport): # type: ignore + """The HTTP2 transport based on httpcore.""" + + TIMEOUT = 15 + + if TYPE_CHECKING: + _pool: """Union[ + httpcore.SOCKSProxy, httpcore.HTTPProxy, httpcore.ConnectionPool + ]""" + + def _get_header_value( + self: "Self", response: "httpcore.Response", header: str + ) -> "Optional[str]": + return next( + ( + val.decode("ascii") + for key, val in response.headers + if key.decode("ascii").lower() == header + ), + None, + ) + + def _request( + self: "Self", + method: str, + endpoint_type: "EndpointType", + body: "Any", + headers: "Mapping[str, str]", + ) -> "httpcore.Response": + response = self._pool.request( + method, + self._auth.get_api_url(endpoint_type), + content=body, + headers=headers, # type: ignore + extensions={ + "timeout": { + "pool": self.TIMEOUT, + "connect": self.TIMEOUT, + "write": self.TIMEOUT, + "read": self.TIMEOUT, + } + }, + ) + return response + + def _get_pool_options(self: "Self") -> "Dict[str, Any]": + options: "Dict[str, Any]" = { + "http2": self.parsed_dsn is not None + and self.parsed_dsn.scheme == "https", + "retries": 3, + } + + socket_options = ( + self.options["socket_options"] + if self.options["socket_options"] is not None + else [] + ) + + used_options = {(o[0], o[1]) for o in socket_options} + for default_option in KEEP_ALIVE_SOCKET_OPTIONS: + if (default_option[0], default_option[1]) not in used_options: + socket_options.append(default_option) + + options["socket_options"] = socket_options + + ssl_context = ssl.create_default_context() + ssl_context.load_verify_locations( + self.options["ca_certs"] # User-provided bundle from the SDK init + or os.environ.get("SSL_CERT_FILE") + or os.environ.get("REQUESTS_CA_BUNDLE") + or certifi.where() + ) + cert_file = self.options["cert_file"] or os.environ.get("CLIENT_CERT_FILE") + key_file = self.options["key_file"] or os.environ.get("CLIENT_KEY_FILE") + if cert_file is not None: + ssl_context.load_cert_chain(cert_file, key_file) + + options["ssl_context"] = ssl_context + + return options + + def _make_pool( + self: "Self", + ) -> "Union[httpcore.SOCKSProxy, httpcore.HTTPProxy, httpcore.ConnectionPool]": + if self.parsed_dsn is None: + raise ValueError("Cannot create HTTP-based transport without valid DSN") + proxy = None + no_proxy = self._in_no_proxy(self.parsed_dsn) + + # try HTTPS first + https_proxy = self.options["https_proxy"] + if self.parsed_dsn.scheme == "https" and (https_proxy != ""): + proxy = https_proxy or (not no_proxy and getproxies().get("https")) + + # maybe fallback to HTTP proxy + http_proxy = self.options["http_proxy"] + if not proxy and (http_proxy != ""): + proxy = http_proxy or (not no_proxy and getproxies().get("http")) + + opts = self._get_pool_options() + + if proxy: + proxy_headers = self.options["proxy_headers"] + if proxy_headers: + opts["proxy_headers"] = proxy_headers + + if proxy.startswith("socks"): + try: + if "socket_options" in opts: + socket_options = opts.pop("socket_options") + if socket_options: + logger.warning( + "You have defined socket_options but using a SOCKS proxy which doesn't support these. We'll ignore socket_options." + ) + return httpcore.SOCKSProxy(proxy_url=proxy, **opts) + except RuntimeError: + logger.warning( + "You have configured a SOCKS proxy (%s) but support for SOCKS proxies is not installed. Disabling proxy support.", + proxy, + ) + else: + return httpcore.HTTPProxy(proxy_url=proxy, **opts) + + return httpcore.ConnectionPool(**opts) + + +class _FunctionTransport(Transport): + """ + DEPRECATED: Users wishing to provide a custom transport should subclass + the Transport class, rather than providing a function. + """ + + def __init__( + self, + func: "Callable[[Event], None]", + ) -> None: + Transport.__init__(self) + self._func = func + + def capture_event( + self, + event: "Event", + ) -> None: + self._func(event) + return None + + def capture_envelope(self, envelope: "Envelope") -> None: + # Since function transports expect to be called with an event, we need + # to iterate over the envelope and call the function for each event, via + # the deprecated capture_event method. + event = envelope.get_event() + if event is not None: + self.capture_event(event) + + +def make_transport(options: "Dict[str, Any]") -> "Optional[Transport]": + ref_transport = options["transport"] + + use_http2_transport = options.get("_experiments", {}).get("transport_http2", False) + + # By default, we use the http transport class + transport_cls: "Type[Transport]" = ( + Http2Transport if use_http2_transport else HttpTransport + ) + + if isinstance(ref_transport, Transport): + return ref_transport + elif isinstance(ref_transport, type) and issubclass(ref_transport, Transport): + transport_cls = ref_transport + elif callable(ref_transport): + warnings.warn( + "Function transports are deprecated and will be removed in a future release." + "Please provide a Transport instance or subclass, instead.", + DeprecationWarning, + stacklevel=2, + ) + return _FunctionTransport(ref_transport) + + # if a transport class is given only instantiate it if the dsn is not + # empty or None + if options["dsn"]: + return transport_cls(options) + + return None diff --git a/lib/python3.12/site-packages/sentry_sdk/types.py b/lib/python3.12/site-packages/sentry_sdk/types.py new file mode 100644 index 0000000000000000000000000000000000000000..8b28166462ba34052438c3ecc2efd977efa8f8c8 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/types.py @@ -0,0 +1,52 @@ +""" +This module contains type definitions for the Sentry SDK's public API. +The types are re-exported from the internal module `sentry_sdk._types`. + +Disclaimer: Since types are a form of documentation, type definitions +may change in minor releases. Removing a type would be considered a +breaking change, and so we will only remove type definitions in major +releases. +""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + # Re-export types to make them available in the public API + from sentry_sdk._types import ( + Breadcrumb, + BreadcrumbHint, + Event, + EventDataCategory, + Hint, + Log, + MonitorConfig, + SamplingContext, + Metric, + ) +else: + from typing import Any + + # The lines below allow the types to be imported from outside `if TYPE_CHECKING` + # guards. The types in this module are only intended to be used for type hints. + Breadcrumb = Any + BreadcrumbHint = Any + Event = Any + EventDataCategory = Any + Hint = Any + Log = Any + MonitorConfig = Any + SamplingContext = Any + Metric = Any + + +__all__ = ( + "Breadcrumb", + "BreadcrumbHint", + "Event", + "EventDataCategory", + "Hint", + "Log", + "MonitorConfig", + "SamplingContext", + "Metric", +) diff --git a/lib/python3.12/site-packages/sentry_sdk/utils.py b/lib/python3.12/site-packages/sentry_sdk/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..79ef5a2cefd0aff18ffae6301e267c57b5e055b2 --- /dev/null +++ b/lib/python3.12/site-packages/sentry_sdk/utils.py @@ -0,0 +1,2109 @@ +import base64 +import json +import linecache +import logging +import math +import os +import copy +import random +import re +import subprocess +import sys +import threading +import time +from collections import namedtuple +from datetime import datetime, timezone +from decimal import Decimal +from functools import partial, partialmethod, wraps +from numbers import Real +from urllib.parse import parse_qs, unquote, urlencode, urlsplit, urlunsplit + +try: + # Python 3.11 + from builtins import BaseExceptionGroup +except ImportError: + # Python 3.10 and below + BaseExceptionGroup = None # type: ignore + +from typing import TYPE_CHECKING + +import sentry_sdk +from sentry_sdk._compat import PY37 +from sentry_sdk._types import SENSITIVE_DATA_SUBSTITUTE, Annotated, AnnotatedValue +from sentry_sdk.consts import ( + DEFAULT_ADD_FULL_STACK, + DEFAULT_MAX_STACK_FRAMES, + DEFAULT_MAX_VALUE_LENGTH, + EndpointType, +) + +if TYPE_CHECKING: + from types import FrameType, TracebackType + from typing import ( + Any, + Callable, + ContextManager, + Dict, + Iterator, + List, + Literal, + NoReturn, + Optional, + ParamSpec, + Set, + Tuple, + Type, + TypeVar, + Union, + cast, + overload, + ) + + from gevent.hub import Hub + + from sentry_sdk._types import ( + AttributeValue, + SerializedAttributeValue, + Event, + ExcInfo, + Hint, + Log, + Metric, + ) + + P = ParamSpec("P") + R = TypeVar("R") + + +epoch = datetime(1970, 1, 1) + +# The logger is created here but initialized in the debug support module +logger = logging.getLogger("sentry_sdk.errors") + +_installed_modules = None + +BASE64_ALPHABET = re.compile(r"^[a-zA-Z0-9/+=]*$") + +FALSY_ENV_VALUES = frozenset(("false", "f", "n", "no", "off", "0")) +TRUTHY_ENV_VALUES = frozenset(("true", "t", "y", "yes", "on", "1")) + +MAX_STACK_FRAMES = 2000 +"""Maximum number of stack frames to send to Sentry. + +If we have more than this number of stack frames, we will stop processing +the stacktrace to avoid getting stuck in a long-lasting loop. This value +exceeds the default sys.getrecursionlimit() of 1000, so users will only +be affected by this limit if they have a custom recursion limit. +""" + + +def env_to_bool(value: "Any", *, strict: "Optional[bool]" = False) -> "bool | None": + """Casts an ENV variable value to boolean using the constants defined above. + In strict mode, it may return None if the value doesn't match any of the predefined values. + """ + normalized = str(value).lower() if value is not None else None + + if normalized in FALSY_ENV_VALUES: + return False + + if normalized in TRUTHY_ENV_VALUES: + return True + + return None if strict else bool(value) + + +def json_dumps(data: "Any") -> bytes: + """Serialize data into a compact JSON representation encoded as UTF-8.""" + return json.dumps(data, allow_nan=False, separators=(",", ":")).encode("utf-8") + + +def get_git_revision() -> "Optional[str]": + try: + with open(os.path.devnull, "w+") as null: + # prevent command prompt windows from popping up on windows + startupinfo = None + if sys.platform == "win32" or sys.platform == "cygwin": + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + + revision = ( + subprocess.Popen( + ["git", "rev-parse", "HEAD"], + startupinfo=startupinfo, + stdout=subprocess.PIPE, + stderr=null, + stdin=null, + ) + .communicate()[0] + .strip() + .decode("utf-8") + ) + except (OSError, IOError, FileNotFoundError): + return None + + return revision + + +def get_default_release() -> "Optional[str]": + """Try to guess a default release.""" + release = os.environ.get("SENTRY_RELEASE") + if release: + return release + + release = get_git_revision() + if release: + return release + + for var in ( + "HEROKU_SLUG_COMMIT", + "SOURCE_VERSION", + "CODEBUILD_RESOLVED_SOURCE_VERSION", + "CIRCLE_SHA1", + "GAE_DEPLOYMENT_ID", + "K_REVISION", + ): + release = os.environ.get(var) + if release: + return release + return None + + +def get_sdk_name(installed_integrations: "List[str]") -> str: + """Return the SDK name including the name of the used web framework.""" + + # Note: I can not use for example sentry_sdk.integrations.django.DjangoIntegration.identifier + # here because if django is not installed the integration is not accessible. + framework_integrations = [ + "django", + "flask", + "fastapi", + "bottle", + "falcon", + "quart", + "sanic", + "starlette", + "litestar", + "starlite", + "chalice", + "serverless", + "pyramid", + "tornado", + "aiohttp", + "aws_lambda", + "gcp", + "beam", + "asgi", + "wsgi", + ] + + for integration in framework_integrations: + if integration in installed_integrations: + return "sentry.python.{}".format(integration) + + return "sentry.python" + + +class CaptureInternalException: + __slots__ = () + + def __enter__(self) -> "ContextManager[Any]": + return self + + def __exit__( + self, + ty: "Optional[Type[BaseException]]", + value: "Optional[BaseException]", + tb: "Optional[TracebackType]", + ) -> bool: + if ty is not None and value is not None: + capture_internal_exception((ty, value, tb)) + + return True + + +_CAPTURE_INTERNAL_EXCEPTION = CaptureInternalException() + + +def capture_internal_exceptions() -> "ContextManager[Any]": + return _CAPTURE_INTERNAL_EXCEPTION + + +def capture_internal_exception(exc_info: "ExcInfo") -> None: + """ + Capture an exception that is likely caused by a bug in the SDK + itself. + + These exceptions do not end up in Sentry and are just logged instead. + """ + if sentry_sdk.get_client().is_active(): + logger.error("Internal error in sentry_sdk", exc_info=exc_info) + + +def to_timestamp(value: "datetime") -> float: + return (value - epoch).total_seconds() + + +def format_timestamp(value: "datetime") -> str: + """Formats a timestamp in RFC 3339 format. + + Any datetime objects with a non-UTC timezone are converted to UTC, so that all timestamps are formatted in UTC. + """ + utctime = value.astimezone(timezone.utc) + + # We use this custom formatting rather than isoformat for backwards compatibility (we have used this format for + # several years now), and isoformat is slightly different. + return utctime.strftime("%Y-%m-%dT%H:%M:%S.%fZ") + + +ISO_TZ_SEPARATORS = frozenset(("+", "-")) + + +def datetime_from_isoformat(value: str) -> "datetime": + try: + result = datetime.fromisoformat(value) + except (AttributeError, ValueError): + # py 3.6 + timestamp_format = ( + "%Y-%m-%dT%H:%M:%S.%f" if "." in value else "%Y-%m-%dT%H:%M:%S" + ) + if value.endswith("Z"): + value = value[:-1] + "+0000" + + if value[-6] in ISO_TZ_SEPARATORS: + timestamp_format += "%z" + value = value[:-3] + value[-2:] + elif value[-5] in ISO_TZ_SEPARATORS: + timestamp_format += "%z" + + result = datetime.strptime(value, timestamp_format) + return result.astimezone(timezone.utc) + + +def event_hint_with_exc_info( + exc_info: "Optional[ExcInfo]" = None, +) -> "Dict[str, Optional[ExcInfo]]": + """Creates a hint with the exc info filled in.""" + if exc_info is None: + exc_info = sys.exc_info() + else: + exc_info = exc_info_from_error(exc_info) + if exc_info[0] is None: + exc_info = None + return {"exc_info": exc_info} + + +class BadDsn(ValueError): + """Raised on invalid DSNs.""" + + +class Dsn: + """Represents a DSN.""" + + ORG_ID_REGEX = re.compile(r"^o(\d+)\.") + + def __init__( + self, value: "Union[Dsn, str]", org_id: "Optional[str]" = None + ) -> None: + if isinstance(value, Dsn): + self.__dict__ = dict(value.__dict__) + return + parts = urlsplit(str(value)) + + if parts.scheme not in ("http", "https"): + raise BadDsn("Unsupported scheme %r" % parts.scheme) + self.scheme = parts.scheme + + if parts.hostname is None: + raise BadDsn("Missing hostname") + + self.host = parts.hostname + + if org_id is not None: + self.org_id: "Optional[str]" = org_id + else: + org_id_match = Dsn.ORG_ID_REGEX.match(self.host) + self.org_id = org_id_match.group(1) if org_id_match else None + + if parts.port is None: + self.port: int = self.scheme == "https" and 443 or 80 + else: + self.port = parts.port + + if not parts.username: + raise BadDsn("Missing public key") + + self.public_key = parts.username + self.secret_key = parts.password + + path = parts.path.rsplit("/", 1) + + try: + self.project_id = str(int(path.pop())) + except (ValueError, TypeError): + raise BadDsn("Invalid project in DSN (%r)" % (parts.path or "")[1:]) + + self.path = "/".join(path) + "/" + + @property + def netloc(self) -> str: + """The netloc part of a DSN.""" + rv = self.host + if (self.scheme, self.port) not in (("http", 80), ("https", 443)): + rv = "%s:%s" % (rv, self.port) + return rv + + def to_auth(self, client: "Optional[Any]" = None) -> "Auth": + """Returns the auth info object for this dsn.""" + return Auth( + scheme=self.scheme, + host=self.netloc, + path=self.path, + project_id=self.project_id, + public_key=self.public_key, + secret_key=self.secret_key, + client=client, + ) + + def __str__(self) -> str: + return "%s://%s%s@%s%s%s" % ( + self.scheme, + self.public_key, + self.secret_key and "@" + self.secret_key or "", + self.netloc, + self.path, + self.project_id, + ) + + +class Auth: + """Helper object that represents the auth info.""" + + def __init__( + self, + scheme: str, + host: str, + project_id: str, + public_key: str, + secret_key: "Optional[str]" = None, + version: int = 7, + client: "Optional[Any]" = None, + path: str = "/", + ) -> None: + self.scheme = scheme + self.host = host + self.path = path + self.project_id = project_id + self.public_key = public_key + self.secret_key = secret_key + self.version = version + self.client = client + + def get_api_url( + self, + type: "EndpointType" = EndpointType.ENVELOPE, + ) -> str: + """Returns the API url for storing events.""" + return "%s://%s%sapi/%s/%s/" % ( + self.scheme, + self.host, + self.path, + self.project_id, + type.value, + ) + + def to_header(self) -> str: + """Returns the auth header a string.""" + rv = [("sentry_key", self.public_key), ("sentry_version", self.version)] + if self.client is not None: + rv.append(("sentry_client", self.client)) + if self.secret_key is not None: + rv.append(("sentry_secret", self.secret_key)) + return "Sentry " + ", ".join("%s=%s" % (key, value) for key, value in rv) + + +def get_type_name(cls: "Optional[type]") -> "Optional[str]": + return getattr(cls, "__qualname__", None) or getattr(cls, "__name__", None) + + +def get_type_module(cls: "Optional[type]") -> "Optional[str]": + mod = getattr(cls, "__module__", None) + if mod not in (None, "builtins", "__builtins__"): + return mod + return None + + +def should_hide_frame(frame: "FrameType") -> bool: + try: + mod = frame.f_globals["__name__"] + if mod.startswith("sentry_sdk."): + return True + except (AttributeError, KeyError): + pass + + for flag_name in "__traceback_hide__", "__tracebackhide__": + try: + if frame.f_locals[flag_name]: + return True + except Exception: + pass + + return False + + +def iter_stacks(tb: "Optional[TracebackType]") -> "Iterator[TracebackType]": + tb_: "Optional[TracebackType]" = tb + while tb_ is not None: + if not should_hide_frame(tb_.tb_frame): + yield tb_ + tb_ = tb_.tb_next + + +def get_lines_from_file( + filename: str, + lineno: int, + max_length: "Optional[int]" = None, + loader: "Optional[Any]" = None, + module: "Optional[str]" = None, +) -> "Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]": + context_lines = 5 + source = None + if loader is not None and hasattr(loader, "get_source"): + try: + source_str: "Optional[str]" = loader.get_source(module) + except (ImportError, IOError): + source_str = None + if source_str is not None: + source = source_str.splitlines() + + if source is None: + try: + source = linecache.getlines(filename) + except (OSError, IOError): + return [], None, [] + + if not source: + return [], None, [] + + lower_bound = max(0, lineno - context_lines) + upper_bound = min(lineno + 1 + context_lines, len(source)) + + try: + pre_context = [ + strip_string(line.strip("\r\n"), max_length=max_length) + for line in source[lower_bound:lineno] + ] + context_line = strip_string(source[lineno].strip("\r\n"), max_length=max_length) + post_context = [ + strip_string(line.strip("\r\n"), max_length=max_length) + for line in source[(lineno + 1) : upper_bound] + ] + return pre_context, context_line, post_context + except IndexError: + # the file may have changed since it was loaded into memory + return [], None, [] + + +def get_source_context( + frame: "FrameType", + tb_lineno: "Optional[int]", + max_value_length: "Optional[int]" = None, +) -> "Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]": + try: + abs_path: "Optional[str]" = frame.f_code.co_filename + except Exception: + abs_path = None + try: + module = frame.f_globals["__name__"] + except Exception: + return [], None, [] + try: + loader = frame.f_globals["__loader__"] + except Exception: + loader = None + + if tb_lineno is not None and abs_path: + lineno = tb_lineno - 1 + return get_lines_from_file( + abs_path, lineno, max_value_length, loader=loader, module=module + ) + + return [], None, [] + + +def safe_str(value: "Any") -> str: + try: + return str(value) + except Exception: + return safe_repr(value) + + +def safe_repr(value: "Any") -> str: + try: + return repr(value) + except Exception: + return "" + + +def filename_for_module( + module: "Optional[str]", abs_path: "Optional[str]" +) -> "Optional[str]": + if not abs_path or not module: + return abs_path + + try: + if abs_path.endswith(".pyc"): + abs_path = abs_path[:-1] + + base_module = module.split(".", 1)[0] + if base_module == module: + return os.path.basename(abs_path) + + base_module_path = sys.modules[base_module].__file__ + if not base_module_path: + return abs_path + + return abs_path.split(base_module_path.rsplit(os.sep, 2)[0], 1)[-1].lstrip( + os.sep + ) + except Exception: + return abs_path + + +def serialize_frame( + frame: "FrameType", + tb_lineno: "Optional[int]" = None, + include_local_variables: bool = True, + include_source_context: bool = True, + max_value_length: "Optional[int]" = None, + custom_repr: "Optional[Callable[..., Optional[str]]]" = None, +) -> "Dict[str, Any]": + f_code = getattr(frame, "f_code", None) + if not f_code: + abs_path = None + function = None + else: + abs_path = frame.f_code.co_filename + function = frame.f_code.co_name + try: + module = frame.f_globals["__name__"] + except Exception: + module = None + + if tb_lineno is None: + tb_lineno = frame.f_lineno + + try: + os_abs_path = os.path.abspath(abs_path) if abs_path else None + except Exception: + os_abs_path = None + + rv: "Dict[str, Any]" = { + "filename": filename_for_module(module, abs_path) or None, + "abs_path": os_abs_path, + "function": function or "", + "module": module, + "lineno": tb_lineno, + } + + if include_source_context: + rv["pre_context"], rv["context_line"], rv["post_context"] = get_source_context( + frame, tb_lineno, max_value_length + ) + + if include_local_variables: + from sentry_sdk.serializer import serialize + + rv["vars"] = serialize( + dict(frame.f_locals), is_vars=True, custom_repr=custom_repr + ) + + return rv + + +def current_stacktrace( + include_local_variables: bool = True, + include_source_context: bool = True, + max_value_length: "Optional[int]" = None, +) -> "Dict[str, Any]": + __tracebackhide__ = True + frames = [] + + f: "Optional[FrameType]" = sys._getframe() + while f is not None: + if not should_hide_frame(f): + frames.append( + serialize_frame( + f, + include_local_variables=include_local_variables, + include_source_context=include_source_context, + max_value_length=max_value_length, + ) + ) + f = f.f_back + + frames.reverse() + + return {"frames": frames} + + +def get_errno(exc_value: BaseException) -> "Optional[Any]": + return getattr(exc_value, "errno", None) + + +def get_error_message(exc_value: "Optional[BaseException]") -> str: + message: str = safe_str( + getattr(exc_value, "message", "") + or getattr(exc_value, "detail", "") + or safe_str(exc_value) + ) + + # __notes__ should be a list of strings when notes are added + # via add_note, but can be anything else if __notes__ is set + # directly. We only support strings in __notes__, since that + # is the correct use. + notes: object = getattr(exc_value, "__notes__", None) + if isinstance(notes, list) and len(notes) > 0: + message += "\n" + "\n".join(note for note in notes if isinstance(note, str)) + + return message + + +def single_exception_from_error_tuple( + exc_type: "Optional[type]", + exc_value: "Optional[BaseException]", + tb: "Optional[TracebackType]", + client_options: "Optional[Dict[str, Any]]" = None, + mechanism: "Optional[Dict[str, Any]]" = None, + exception_id: "Optional[int]" = None, + parent_id: "Optional[int]" = None, + source: "Optional[str]" = None, + full_stack: "Optional[list[dict[str, Any]]]" = None, +) -> "Dict[str, Any]": + """ + Creates a dict that goes into the events `exception.values` list and is ingestible by Sentry. + + See the Exception Interface documentation for more details: + https://develop.sentry.dev/sdk/event-payloads/exception/ + """ + exception_value: "Dict[str, Any]" = {} + exception_value["mechanism"] = ( + mechanism.copy() if mechanism else {"type": "generic", "handled": True} + ) + if exception_id is not None: + exception_value["mechanism"]["exception_id"] = exception_id + + if exc_value is not None: + errno = get_errno(exc_value) + else: + errno = None + + if errno is not None: + exception_value["mechanism"].setdefault("meta", {}).setdefault( + "errno", {} + ).setdefault("number", errno) + + if source is not None: + exception_value["mechanism"]["source"] = source + + is_root_exception = exception_id == 0 + if not is_root_exception and parent_id is not None: + exception_value["mechanism"]["parent_id"] = parent_id + exception_value["mechanism"]["type"] = "chained" + + if is_root_exception and "type" not in exception_value["mechanism"]: + exception_value["mechanism"]["type"] = "generic" + + is_exception_group = BaseExceptionGroup is not None and isinstance( + exc_value, BaseExceptionGroup + ) + if is_exception_group: + exception_value["mechanism"]["is_exception_group"] = True + + exception_value["module"] = get_type_module(exc_type) + exception_value["type"] = get_type_name(exc_type) + exception_value["value"] = get_error_message(exc_value) + + if client_options is None: + include_local_variables = True + include_source_context = True + max_value_length = DEFAULT_MAX_VALUE_LENGTH # fallback + custom_repr = None + else: + include_local_variables = client_options["include_local_variables"] + include_source_context = client_options["include_source_context"] + max_value_length = client_options["max_value_length"] + custom_repr = client_options.get("custom_repr") + + frames: "List[Dict[str, Any]]" = [ + serialize_frame( + tb.tb_frame, + tb_lineno=tb.tb_lineno, + include_local_variables=include_local_variables, + include_source_context=include_source_context, + max_value_length=max_value_length, + custom_repr=custom_repr, + ) + # Process at most MAX_STACK_FRAMES + 1 frames, to avoid hanging on + # processing a super-long stacktrace. + for tb, _ in zip(iter_stacks(tb), range(MAX_STACK_FRAMES + 1)) + ] + + if len(frames) > MAX_STACK_FRAMES: + # If we have more frames than the limit, we remove the stacktrace completely. + # We don't trim the stacktrace here because we have not processed the whole + # thing (see above, we stop at MAX_STACK_FRAMES + 1). Normally, Relay would + # intelligently trim by removing frames in the middle of the stacktrace, but + # since we don't have the whole stacktrace, we can't do that. Instead, we + # drop the entire stacktrace. + exception_value["stacktrace"] = AnnotatedValue.removed_because_over_size_limit( + value=None + ) + + elif frames: + if not full_stack: + new_frames = frames + else: + new_frames = merge_stack_frames(frames, full_stack, client_options) + + exception_value["stacktrace"] = {"frames": new_frames} + + return exception_value + + +HAS_CHAINED_EXCEPTIONS = hasattr(Exception, "__suppress_context__") + +if HAS_CHAINED_EXCEPTIONS: + + def walk_exception_chain(exc_info: "ExcInfo") -> "Iterator[ExcInfo]": + exc_type, exc_value, tb = exc_info + + seen_exceptions = [] + seen_exception_ids: "Set[int]" = set() + + while ( + exc_type is not None + and exc_value is not None + and id(exc_value) not in seen_exception_ids + ): + yield exc_type, exc_value, tb + + # Avoid hashing random types we don't know anything + # about. Use the list to keep a ref so that the `id` is + # not used for another object. + seen_exceptions.append(exc_value) + seen_exception_ids.add(id(exc_value)) + + if exc_value.__suppress_context__: + cause = exc_value.__cause__ + else: + cause = exc_value.__context__ + if cause is None: + break + exc_type = type(cause) + exc_value = cause + tb = getattr(cause, "__traceback__", None) + +else: + + def walk_exception_chain(exc_info: "ExcInfo") -> "Iterator[ExcInfo]": + yield exc_info + + +def exceptions_from_error( + exc_type: "Optional[type]", + exc_value: "Optional[BaseException]", + tb: "Optional[TracebackType]", + client_options: "Optional[Dict[str, Any]]" = None, + mechanism: "Optional[Dict[str, Any]]" = None, + exception_id: int = 0, + parent_id: int = 0, + source: "Optional[str]" = None, + full_stack: "Optional[list[dict[str, Any]]]" = None, +) -> "Tuple[int, List[Dict[str, Any]]]": + """ + Creates the list of exceptions. + This can include chained exceptions and exceptions from an ExceptionGroup. + + See the Exception Interface documentation for more details: + https://develop.sentry.dev/sdk/event-payloads/exception/ + """ + + parent = single_exception_from_error_tuple( + exc_type=exc_type, + exc_value=exc_value, + tb=tb, + client_options=client_options, + mechanism=mechanism, + exception_id=exception_id, + parent_id=parent_id, + source=source, + full_stack=full_stack, + ) + exceptions = [parent] + + parent_id = exception_id + exception_id += 1 + + should_supress_context = ( + hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__ # type: ignore + ) + if should_supress_context: + # Add direct cause. + # The field `__cause__` is set when raised with the exception (using the `from` keyword). + exception_has_cause = ( + exc_value + and hasattr(exc_value, "__cause__") + and exc_value.__cause__ is not None + ) + if exception_has_cause: + cause = exc_value.__cause__ # type: ignore + (exception_id, child_exceptions) = exceptions_from_error( + exc_type=type(cause), + exc_value=cause, + tb=getattr(cause, "__traceback__", None), + client_options=client_options, + mechanism=mechanism, + exception_id=exception_id, + source="__cause__", + full_stack=full_stack, + ) + exceptions.extend(child_exceptions) + + else: + # Add indirect cause. + # The field `__context__` is assigned if another exception occurs while handling the exception. + exception_has_content = ( + exc_value + and hasattr(exc_value, "__context__") + and exc_value.__context__ is not None + ) + if exception_has_content: + context = exc_value.__context__ # type: ignore + (exception_id, child_exceptions) = exceptions_from_error( + exc_type=type(context), + exc_value=context, + tb=getattr(context, "__traceback__", None), + client_options=client_options, + mechanism=mechanism, + exception_id=exception_id, + source="__context__", + full_stack=full_stack, + ) + exceptions.extend(child_exceptions) + + # Add exceptions from an ExceptionGroup. + is_exception_group = exc_value and hasattr(exc_value, "exceptions") + if is_exception_group: + for idx, e in enumerate(exc_value.exceptions): # type: ignore + (exception_id, child_exceptions) = exceptions_from_error( + exc_type=type(e), + exc_value=e, + tb=getattr(e, "__traceback__", None), + client_options=client_options, + mechanism=mechanism, + exception_id=exception_id, + parent_id=parent_id, + source="exceptions[%s]" % idx, + full_stack=full_stack, + ) + exceptions.extend(child_exceptions) + + return (exception_id, exceptions) + + +def exceptions_from_error_tuple( + exc_info: "ExcInfo", + client_options: "Optional[Dict[str, Any]]" = None, + mechanism: "Optional[Dict[str, Any]]" = None, + full_stack: "Optional[list[dict[str, Any]]]" = None, +) -> "List[Dict[str, Any]]": + exc_type, exc_value, tb = exc_info + + is_exception_group = BaseExceptionGroup is not None and isinstance( + exc_value, BaseExceptionGroup + ) + + if is_exception_group: + (_, exceptions) = exceptions_from_error( + exc_type=exc_type, + exc_value=exc_value, + tb=tb, + client_options=client_options, + mechanism=mechanism, + exception_id=0, + parent_id=0, + full_stack=full_stack, + ) + + else: + exceptions = [] + for exc_type, exc_value, tb in walk_exception_chain(exc_info): + exceptions.append( + single_exception_from_error_tuple( + exc_type=exc_type, + exc_value=exc_value, + tb=tb, + client_options=client_options, + mechanism=mechanism, + full_stack=full_stack, + ) + ) + + exceptions.reverse() + + return exceptions + + +def to_string(value: str) -> str: + try: + return str(value) + except UnicodeDecodeError: + return repr(value)[1:-1] + + +def iter_event_stacktraces(event: "Event") -> "Iterator[Annotated[Dict[str, Any]]]": + if "stacktrace" in event: + yield event["stacktrace"] + if "threads" in event: + for thread in event["threads"].get("values") or (): + if "stacktrace" in thread: + yield thread["stacktrace"] + if "exception" in event: + for exception in event["exception"].get("values") or (): + if isinstance(exception, dict) and "stacktrace" in exception: + yield exception["stacktrace"] + + +def iter_event_frames(event: "Event") -> "Iterator[Dict[str, Any]]": + for stacktrace in iter_event_stacktraces(event): + if isinstance(stacktrace, AnnotatedValue): + stacktrace = stacktrace.value or {} + + for frame in stacktrace.get("frames") or (): + yield frame + + +def handle_in_app( + event: "Event", + in_app_exclude: "Optional[List[str]]" = None, + in_app_include: "Optional[List[str]]" = None, + project_root: "Optional[str]" = None, +) -> "Event": + for stacktrace in iter_event_stacktraces(event): + if isinstance(stacktrace, AnnotatedValue): + stacktrace = stacktrace.value or {} + + set_in_app_in_frames( + stacktrace.get("frames"), + in_app_exclude=in_app_exclude, + in_app_include=in_app_include, + project_root=project_root, + ) + + return event + + +def set_in_app_in_frames( + frames: "Any", + in_app_exclude: "Optional[List[str]]", + in_app_include: "Optional[List[str]]", + project_root: "Optional[str]" = None, +) -> "Optional[Any]": + if not frames: + return None + + for frame in frames: + # if frame has already been marked as in_app, skip it + current_in_app = frame.get("in_app") + if current_in_app is not None: + continue + + module = frame.get("module") + + # check if module in frame is in the list of modules to include + if _module_in_list(module, in_app_include): + frame["in_app"] = True + continue + + # check if module in frame is in the list of modules to exclude + if _module_in_list(module, in_app_exclude): + frame["in_app"] = False + continue + + # if frame has no abs_path, skip further checks + abs_path = frame.get("abs_path") + if abs_path is None: + continue + + if _is_external_source(abs_path): + frame["in_app"] = False + continue + + if _is_in_project_root(abs_path, project_root): + frame["in_app"] = True + continue + + return frames + + +def exc_info_from_error(error: "Union[BaseException, ExcInfo]") -> "ExcInfo": + if isinstance(error, tuple) and len(error) == 3: + exc_type, exc_value, tb = error + elif isinstance(error, BaseException): + tb = getattr(error, "__traceback__", None) + if tb is not None: + exc_type = type(error) + exc_value = error + else: + exc_type, exc_value, tb = sys.exc_info() + if exc_value is not error: + tb = None + exc_value = error + exc_type = type(error) + + else: + raise ValueError("Expected Exception object to report, got %s!" % type(error)) + + exc_info = (exc_type, exc_value, tb) + + if TYPE_CHECKING: + # This cast is safe because exc_type and exc_value are either both + # None or both not None. + exc_info = cast(ExcInfo, exc_info) + + return exc_info + + +def merge_stack_frames( + frames: "List[Dict[str, Any]]", + full_stack: "List[Dict[str, Any]]", + client_options: "Optional[Dict[str, Any]]", +) -> "List[Dict[str, Any]]": + """ + Add the missing frames from full_stack to frames and return the merged list. + """ + frame_ids = { + ( + frame["abs_path"], + frame["context_line"], + frame["lineno"], + frame["function"], + ) + for frame in frames + } + + new_frames = [ + stackframe + for stackframe in full_stack + if ( + stackframe["abs_path"], + stackframe["context_line"], + stackframe["lineno"], + stackframe["function"], + ) + not in frame_ids + ] + new_frames.extend(frames) + + # Limit the number of frames + max_stack_frames = ( + client_options.get("max_stack_frames", DEFAULT_MAX_STACK_FRAMES) + if client_options + else None + ) + if max_stack_frames is not None: + new_frames = new_frames[len(new_frames) - max_stack_frames :] + + return new_frames + + +def event_from_exception( + exc_info: "Union[BaseException, ExcInfo]", + client_options: "Optional[Dict[str, Any]]" = None, + mechanism: "Optional[Dict[str, Any]]" = None, +) -> "Tuple[Event, Dict[str, Any]]": + exc_info = exc_info_from_error(exc_info) + hint = event_hint_with_exc_info(exc_info) + + if client_options and client_options.get("add_full_stack", DEFAULT_ADD_FULL_STACK): + full_stack = current_stacktrace( + include_local_variables=client_options["include_local_variables"], + max_value_length=client_options["max_value_length"], + )["frames"] + else: + full_stack = None + + return ( + { + "level": "error", + "exception": { + "values": exceptions_from_error_tuple( + exc_info, client_options, mechanism, full_stack + ) + }, + }, + hint, + ) + + +def _module_in_list(name: "Optional[str]", items: "Optional[List[str]]") -> bool: + if name is None: + return False + + if not items: + return False + + for item in items: + if item == name or name.startswith(item + "."): + return True + + return False + + +def _is_external_source(abs_path: "Optional[str]") -> bool: + # check if frame is in 'site-packages' or 'dist-packages' + if abs_path is None: + return False + + external_source = ( + re.search(r"[\\/](?:dist|site)-packages[\\/]", abs_path) is not None + ) + return external_source + + +def _is_in_project_root( + abs_path: "Optional[str]", project_root: "Optional[str]" +) -> bool: + if abs_path is None or project_root is None: + return False + + # check if path is in the project root + if abs_path.startswith(project_root): + return True + + return False + + +def _truncate_by_bytes(string: str, max_bytes: int) -> str: + """ + Truncate a UTF-8-encodable string to the last full codepoint so that it fits in max_bytes. + """ + truncated = string.encode("utf-8")[: max_bytes - 3].decode("utf-8", errors="ignore") + + return truncated + "..." + + +def _get_size_in_bytes(value: str) -> "Optional[int]": + try: + return len(value.encode("utf-8")) + except (UnicodeEncodeError, UnicodeDecodeError): + return None + + +def strip_string( + value: str, max_length: "Optional[int]" = None +) -> "Union[AnnotatedValue, str]": + if not value: + return value + + if max_length is None: + max_length = DEFAULT_MAX_VALUE_LENGTH + + byte_size = _get_size_in_bytes(value) + text_size = len(value) + + if byte_size is not None and byte_size > max_length: + # truncate to max_length bytes, preserving code points + truncated_value = _truncate_by_bytes(value, max_length) + elif text_size is not None and text_size > max_length: + # fallback to truncating by string length + truncated_value = value[: max_length - 3] + "..." + else: + return value + + return AnnotatedValue( + value=truncated_value, + metadata={ + "len": byte_size or text_size, + "rem": [["!limit", "x", max_length - 3, max_length]], + }, + ) + + +def parse_version(version: str) -> "Optional[Tuple[int, ...]]": + """ + Parses a version string into a tuple of integers. + This uses the parsing loging from PEP 440: + https://peps.python.org/pep-0440/#appendix-b-parsing-version-strings-with-regular-expressions + """ + VERSION_PATTERN = r""" # noqa: N806 + v? + (?: + (?:(?P[0-9]+)!)? # epoch + (?P[0-9]+(?:\.[0-9]+)*) # release segment + (?P
                                              # pre-release
+                [-_\.]?
+                (?P(a|b|c|rc|alpha|beta|pre|preview))
+                [-_\.]?
+                (?P[0-9]+)?
+            )?
+            (?P                                         # post release
+                (?:-(?P[0-9]+))
+                |
+                (?:
+                    [-_\.]?
+                    (?Ppost|rev|r)
+                    [-_\.]?
+                    (?P[0-9]+)?
+                )
+            )?
+            (?P                                          # dev release
+                [-_\.]?
+                (?Pdev)
+                [-_\.]?
+                (?P[0-9]+)?
+            )?
+        )
+        (?:\+(?P[a-z0-9]+(?:[-_\.][a-z0-9]+)*))?       # local version
+    """
+
+    pattern = re.compile(
+        r"^\s*" + VERSION_PATTERN + r"\s*$",
+        re.VERBOSE | re.IGNORECASE,
+    )
+
+    try:
+        release = pattern.match(version).groupdict()["release"]  # type: ignore
+        release_tuple: "Tuple[int, ...]" = tuple(map(int, release.split(".")[:3]))
+    except (TypeError, ValueError, AttributeError):
+        return None
+
+    return release_tuple
+
+
+def _is_contextvars_broken() -> bool:
+    """
+    Returns whether gevent/eventlet have patched the stdlib in a way where thread locals are now more "correct" than contextvars.
+    """
+    try:
+        import gevent
+        from gevent.monkey import is_object_patched
+
+        # Get the MAJOR and MINOR version numbers of Gevent
+        version_tuple = tuple(
+            [int(part) for part in re.split(r"a|b|rc|\.", gevent.__version__)[:2]]
+        )
+        if is_object_patched("threading", "local"):
+            # Gevent 20.9.0 depends on Greenlet 0.4.17 which natively handles switching
+            # context vars when greenlets are switched, so, Gevent 20.9.0+ is all fine.
+            # Ref: https://github.com/gevent/gevent/blob/83c9e2ae5b0834b8f84233760aabe82c3ba065b4/src/gevent/monkey.py#L604-L609
+            # Gevent 20.5, that doesn't depend on Greenlet 0.4.17 with native support
+            # for contextvars, is able to patch both thread locals and contextvars, in
+            # that case, check if contextvars are effectively patched.
+            if (
+                # Gevent 20.9.0+
+                (sys.version_info >= (3, 7) and version_tuple >= (20, 9))
+                # Gevent 20.5.0+ or Python < 3.7
+                or (is_object_patched("contextvars", "ContextVar"))
+            ):
+                return False
+
+            return True
+    except ImportError:
+        pass
+
+    try:
+        import greenlet
+        from eventlet.patcher import is_monkey_patched  # type: ignore
+
+        greenlet_version = parse_version(greenlet.__version__)
+
+        if greenlet_version is None:
+            logger.error(
+                "Internal error in Sentry SDK: Could not parse Greenlet version from greenlet.__version__."
+            )
+            return False
+
+        if is_monkey_patched("thread") and greenlet_version < (0, 5):
+            return True
+    except ImportError:
+        pass
+
+    return False
+
+
+def _make_threadlocal_contextvars(local: type) -> type:
+    class ContextVar:
+        # Super-limited impl of ContextVar
+
+        def __init__(self, name: str, default: "Any" = None) -> None:
+            self._name = name
+            self._default = default
+            self._local = local()
+            self._original_local = local()
+
+        def get(self, default: "Any" = None) -> "Any":
+            return getattr(self._local, "value", default or self._default)
+
+        def set(self, value: "Any") -> "Any":
+            token = str(random.getrandbits(64))
+            original_value = self.get()
+            setattr(self._original_local, token, original_value)
+            self._local.value = value
+            return token
+
+        def reset(self, token: "Any") -> None:
+            self._local.value = getattr(self._original_local, token)
+            # delete the original value (this way it works in Python 3.6+)
+            del self._original_local.__dict__[token]
+
+    return ContextVar
+
+
+def _get_contextvars() -> "Tuple[bool, type]":
+    """
+    Figure out the "right" contextvars installation to use. Returns a
+    `contextvars.ContextVar`-like class with a limited API.
+
+    See https://docs.sentry.io/platforms/python/contextvars/ for more information.
+    """
+    if not _is_contextvars_broken():
+        # aiocontextvars is a PyPI package that ensures that the contextvars
+        # backport (also a PyPI package) works with asyncio under Python 3.6
+        #
+        # Import it if available.
+        if sys.version_info < (3, 7):
+            # `aiocontextvars` is absolutely required for functional
+            # contextvars on Python 3.6.
+            try:
+                from aiocontextvars import ContextVar
+
+                return True, ContextVar
+            except ImportError:
+                pass
+        else:
+            # On Python 3.7 contextvars are functional.
+            try:
+                from contextvars import ContextVar
+
+                return True, ContextVar
+            except ImportError:
+                pass
+
+    # Fall back to basic thread-local usage.
+
+    from threading import local
+
+    return False, _make_threadlocal_contextvars(local)
+
+
+HAS_REAL_CONTEXTVARS, ContextVar = _get_contextvars()
+
+CONTEXTVARS_ERROR_MESSAGE = """
+
+With asyncio/ASGI applications, the Sentry SDK requires a functional
+installation of `contextvars` to avoid leaking scope/context data across
+requests.
+
+Please refer to https://docs.sentry.io/platforms/python/contextvars/ for more information.
+"""
+
+
+def qualname_from_function(func: "Callable[..., Any]") -> "Optional[str]":
+    """Return the qualified name of func. Works with regular function, lambda, partial and partialmethod."""
+    func_qualname: "Optional[str]" = None
+
+    # Python 2
+    try:
+        return "%s.%s.%s" % (
+            func.im_class.__module__,  # type: ignore
+            func.im_class.__name__,  # type: ignore
+            func.__name__,
+        )
+    except Exception:
+        pass
+
+    prefix, suffix = "", ""
+
+    if isinstance(func, partial) and hasattr(func.func, "__name__"):
+        prefix, suffix = "partial()"
+        func = func.func
+    else:
+        # The _partialmethod attribute of methods wrapped with partialmethod() was renamed to __partialmethod__ in CPython 3.13:
+        # https://github.com/python/cpython/pull/16600
+        partial_method = getattr(func, "_partialmethod", None) or getattr(
+            func, "__partialmethod__", None
+        )
+        if isinstance(partial_method, partialmethod):
+            prefix, suffix = "partialmethod()"
+            func = partial_method.func
+
+    if hasattr(func, "__qualname__"):
+        func_qualname = func.__qualname__
+    elif hasattr(func, "__name__"):  # Python 2.7 has no __qualname__
+        func_qualname = func.__name__
+
+    # Python 3: methods, functions, classes
+    if func_qualname is not None:
+        if hasattr(func, "__module__") and isinstance(func.__module__, str):
+            func_qualname = func.__module__ + "." + func_qualname
+        func_qualname = prefix + func_qualname + suffix
+
+    return func_qualname
+
+
+def transaction_from_function(func: "Callable[..., Any]") -> "Optional[str]":
+    return qualname_from_function(func)
+
+
+disable_capture_event = ContextVar("disable_capture_event")
+
+
+class ServerlessTimeoutWarning(Exception):  # noqa: N818
+    """Raised when a serverless method is about to reach its timeout."""
+
+    pass
+
+
+class TimeoutThread(threading.Thread):
+    """Creates a Thread which runs (sleeps) for a time duration equal to
+    waiting_time and raises a custom ServerlessTimeout exception.
+    """
+
+    def __init__(
+        self,
+        waiting_time: float,
+        configured_timeout: int,
+        isolation_scope: "Optional[sentry_sdk.Scope]" = None,
+        current_scope: "Optional[sentry_sdk.Scope]" = None,
+    ) -> None:
+        threading.Thread.__init__(self)
+        self.waiting_time = waiting_time
+        self.configured_timeout = configured_timeout
+
+        self.isolation_scope = isolation_scope
+        self.current_scope = current_scope
+
+        self._stop_event = threading.Event()
+
+    def stop(self) -> None:
+        self._stop_event.set()
+
+    def _capture_exception(self) -> "ExcInfo":
+        exc_info = sys.exc_info()
+
+        client = sentry_sdk.get_client()
+        event, hint = event_from_exception(
+            exc_info,
+            client_options=client.options,
+            mechanism={"type": "threading", "handled": False},
+        )
+        sentry_sdk.capture_event(event, hint=hint)
+
+        return exc_info
+
+    def run(self) -> None:
+        self._stop_event.wait(self.waiting_time)
+
+        if self._stop_event.is_set():
+            return
+
+        integer_configured_timeout = int(self.configured_timeout)
+
+        # Setting up the exact integer value of configured time(in seconds)
+        if integer_configured_timeout < self.configured_timeout:
+            integer_configured_timeout = integer_configured_timeout + 1
+
+        # Raising Exception after timeout duration is reached
+        if self.isolation_scope is not None and self.current_scope is not None:
+            with sentry_sdk.scope.use_isolation_scope(self.isolation_scope):
+                with sentry_sdk.scope.use_scope(self.current_scope):
+                    try:
+                        raise ServerlessTimeoutWarning(
+                            "WARNING : Function is expected to get timed out. Configured timeout duration = {} seconds.".format(
+                                integer_configured_timeout
+                            )
+                        )
+                    except Exception:
+                        reraise(*self._capture_exception())
+
+        raise ServerlessTimeoutWarning(
+            "WARNING : Function is expected to get timed out. Configured timeout duration = {} seconds.".format(
+                integer_configured_timeout
+            )
+        )
+
+
+def to_base64(original: str) -> "Optional[str]":
+    """
+    Convert a string to base64, via UTF-8. Returns None on invalid input.
+    """
+    base64_string = None
+
+    try:
+        utf8_bytes = original.encode("UTF-8")
+        base64_bytes = base64.b64encode(utf8_bytes)
+        base64_string = base64_bytes.decode("UTF-8")
+    except Exception as err:
+        logger.warning("Unable to encode {orig} to base64:".format(orig=original), err)
+
+    return base64_string
+
+
+def from_base64(base64_string: str) -> "Optional[str]":
+    """
+    Convert a string from base64, via UTF-8. Returns None on invalid input.
+    """
+    utf8_string = None
+
+    try:
+        only_valid_chars = BASE64_ALPHABET.match(base64_string)
+        assert only_valid_chars
+
+        base64_bytes = base64_string.encode("UTF-8")
+        utf8_bytes = base64.b64decode(base64_bytes)
+        utf8_string = utf8_bytes.decode("UTF-8")
+    except Exception as err:
+        logger.warning(
+            "Unable to decode {b64} from base64:".format(b64=base64_string), err
+        )
+
+    return utf8_string
+
+
+Components = namedtuple("Components", ["scheme", "netloc", "path", "query", "fragment"])
+
+
+def sanitize_url(
+    url: str,
+    remove_authority: bool = True,
+    remove_query_values: bool = True,
+    split: bool = False,
+) -> "Union[str, Components]":
+    """
+    Removes the authority and query parameter values from a given URL.
+    """
+    parsed_url = urlsplit(url)
+    query_params = parse_qs(parsed_url.query, keep_blank_values=True)
+
+    # strip username:password (netloc can be usr:pwd@example.com)
+    if remove_authority:
+        netloc_parts = parsed_url.netloc.split("@")
+        if len(netloc_parts) > 1:
+            netloc = "%s:%s@%s" % (
+                SENSITIVE_DATA_SUBSTITUTE,
+                SENSITIVE_DATA_SUBSTITUTE,
+                netloc_parts[-1],
+            )
+        else:
+            netloc = parsed_url.netloc
+    else:
+        netloc = parsed_url.netloc
+
+    # strip values from query string
+    if remove_query_values:
+        query_string = unquote(
+            urlencode({key: SENSITIVE_DATA_SUBSTITUTE for key in query_params})
+        )
+    else:
+        query_string = parsed_url.query
+
+    components = Components(
+        scheme=parsed_url.scheme,
+        netloc=netloc,
+        query=query_string,
+        path=parsed_url.path,
+        fragment=parsed_url.fragment,
+    )
+
+    if split:
+        return components
+    else:
+        return urlunsplit(components)
+
+
+ParsedUrl = namedtuple("ParsedUrl", ["url", "query", "fragment"])
+
+
+def parse_url(url: str, sanitize: bool = True) -> "ParsedUrl":
+    """
+    Splits a URL into a url (including path), query and fragment. If sanitize is True, the query
+    parameters will be sanitized to remove sensitive data. The autority (username and password)
+    in the URL will always be removed.
+    """
+    parsed_url = sanitize_url(
+        url, remove_authority=True, remove_query_values=sanitize, split=True
+    )
+
+    base_url = urlunsplit(
+        Components(
+            scheme=parsed_url.scheme,  # type: ignore
+            netloc=parsed_url.netloc,  # type: ignore
+            query="",
+            path=parsed_url.path,  # type: ignore
+            fragment="",
+        )
+    )
+
+    return ParsedUrl(
+        url=base_url,
+        query=parsed_url.query,  # type: ignore
+        fragment=parsed_url.fragment,  # type: ignore
+    )
+
+
+def is_valid_sample_rate(rate: "Any", source: str) -> bool:
+    """
+    Checks the given sample rate to make sure it is valid type and value (a
+    boolean or a number between 0 and 1, inclusive).
+    """
+
+    # both booleans and NaN are instances of Real, so a) checking for Real
+    # checks for the possibility of a boolean also, and b) we have to check
+    # separately for NaN and Decimal does not derive from Real so need to check that too
+    if not isinstance(rate, (Real, Decimal)) or math.isnan(rate):
+        logger.warning(
+            "{source} Given sample rate is invalid. Sample rate must be a boolean or a number between 0 and 1. Got {rate} of type {type}.".format(
+                source=source, rate=rate, type=type(rate)
+            )
+        )
+        return False
+
+    # in case rate is a boolean, it will get cast to 1 if it's True and 0 if it's False
+    rate = float(rate)
+    if rate < 0 or rate > 1:
+        logger.warning(
+            "{source} Given sample rate is invalid. Sample rate must be between 0 and 1. Got {rate}.".format(
+                source=source, rate=rate
+            )
+        )
+        return False
+
+    return True
+
+
+def match_regex_list(
+    item: str,
+    regex_list: "Optional[List[str]]" = None,
+    substring_matching: bool = False,
+) -> bool:
+    if regex_list is None:
+        return False
+
+    for item_matcher in regex_list:
+        if not substring_matching and item_matcher[-1] != "$":
+            item_matcher += "$"
+
+        matched = re.search(item_matcher, item)
+        if matched:
+            return True
+
+    return False
+
+
+def is_sentry_url(client: "sentry_sdk.client.BaseClient", url: str) -> bool:
+    """
+    Determines whether the given URL matches the Sentry DSN.
+    """
+    return (
+        client is not None
+        and client.transport is not None
+        and client.transport.parsed_dsn is not None
+        and client.transport.parsed_dsn.netloc in url
+    )
+
+
+def _generate_installed_modules() -> "Iterator[Tuple[str, str]]":
+    try:
+        from importlib import metadata
+
+        yielded = set()
+        for dist in metadata.distributions():
+            name = dist.metadata.get("Name", None)  # type: ignore[attr-defined]
+            # `metadata` values may be `None`, see:
+            # https://github.com/python/cpython/issues/91216
+            # and
+            # https://github.com/python/importlib_metadata/issues/371
+            if name is not None:
+                normalized_name = _normalize_module_name(name)
+                if dist.version is not None and normalized_name not in yielded:
+                    yield normalized_name, dist.version
+                    yielded.add(normalized_name)
+
+    except ImportError:
+        # < py3.8
+        try:
+            import pkg_resources
+        except ImportError:
+            return
+
+        for info in pkg_resources.working_set:
+            yield _normalize_module_name(info.key), info.version
+
+
+def _normalize_module_name(name: str) -> str:
+    return name.lower()
+
+
+def _replace_hyphens_dots_and_underscores_with_dashes(name: str) -> str:
+    # https://peps.python.org/pep-0503/#normalized-names
+    return re.sub(r"[-_.]+", "-", name)
+
+
+def _get_installed_modules() -> "Dict[str, str]":
+    global _installed_modules
+    if _installed_modules is None:
+        _installed_modules = dict(_generate_installed_modules())
+    return _installed_modules
+
+
+def package_version(package: str) -> "Optional[Tuple[int, ...]]":
+    normalized_package = _normalize_module_name(
+        _replace_hyphens_dots_and_underscores_with_dashes(package)
+    )
+
+    installed_packages = {
+        _replace_hyphens_dots_and_underscores_with_dashes(module): v
+        for module, v in _get_installed_modules().items()
+    }
+    version = installed_packages.get(normalized_package)
+    if version is None:
+        return None
+
+    return parse_version(version)
+
+
+def reraise(
+    tp: "Optional[Type[BaseException]]",
+    value: "Optional[BaseException]",
+    tb: "Optional[Any]" = None,
+) -> "NoReturn":
+    assert value is not None
+    if value.__traceback__ is not tb:
+        raise value.with_traceback(tb)
+    raise value
+
+
+def _no_op(*_a: "Any", **_k: "Any") -> None:
+    """No-op function for ensure_integration_enabled."""
+    pass
+
+
+if TYPE_CHECKING:
+
+    @overload
+    def ensure_integration_enabled(
+        integration: "type[sentry_sdk.integrations.Integration]",
+        original_function: "Callable[P, R]",
+    ) -> "Callable[[Callable[P, R]], Callable[P, R]]": ...
+
+    @overload
+    def ensure_integration_enabled(
+        integration: "type[sentry_sdk.integrations.Integration]",
+    ) -> "Callable[[Callable[P, None]], Callable[P, None]]": ...
+
+
+def ensure_integration_enabled(
+    integration: "type[sentry_sdk.integrations.Integration]",
+    original_function: "Union[Callable[P, R], Callable[P, None]]" = _no_op,
+) -> "Callable[[Callable[P, R]], Callable[P, R]]":
+    """
+    Ensures a given integration is enabled prior to calling a Sentry-patched function.
+
+    The function takes as its parameters the integration that must be enabled and the original
+    function that the SDK is patching. The function returns a function that takes the
+    decorated (Sentry-patched) function as its parameter, and returns a function that, when
+    called, checks whether the given integration is enabled. If the integration is enabled, the
+    function calls the decorated, Sentry-patched function. If the integration is not enabled,
+    the original function is called.
+
+    The function also takes care of preserving the original function's signature and docstring.
+
+    Example usage:
+
+    ```python
+    @ensure_integration_enabled(MyIntegration, my_function)
+    def patch_my_function():
+        with sentry_sdk.start_transaction(...):
+            return my_function()
+    ```
+    """
+    if TYPE_CHECKING:
+        # Type hint to ensure the default function has the right typing. The overloads
+        # ensure the default _no_op function is only used when R is None.
+        original_function = cast(Callable[P, R], original_function)
+
+    def patcher(sentry_patched_function: "Callable[P, R]") -> "Callable[P, R]":
+        def runner(*args: "P.args", **kwargs: "P.kwargs") -> "R":
+            if sentry_sdk.get_client().get_integration(integration) is None:
+                return original_function(*args, **kwargs)
+
+            return sentry_patched_function(*args, **kwargs)
+
+        if original_function is _no_op:
+            return wraps(sentry_patched_function)(runner)
+
+        return wraps(original_function)(runner)
+
+    return patcher
+
+
+if PY37:
+
+    def nanosecond_time() -> int:
+        return time.perf_counter_ns()
+
+else:
+
+    def nanosecond_time() -> int:
+        return int(time.perf_counter() * 1e9)
+
+
+def now() -> float:
+    return time.perf_counter()
+
+
+try:
+    from gevent import get_hub as get_gevent_hub
+    from gevent.monkey import is_module_patched
+except ImportError:
+    # it's not great that the signatures are different, get_hub can't return None
+    # consider adding an if TYPE_CHECKING to change the signature to Optional[Hub]
+    def get_gevent_hub() -> "Optional[Hub]":  # type: ignore[misc]
+        return None
+
+    def is_module_patched(mod_name: str) -> bool:
+        # unable to import from gevent means no modules have been patched
+        return False
+
+
+def is_gevent() -> bool:
+    return is_module_patched("threading") or is_module_patched("_thread")
+
+
+def get_current_thread_meta(
+    thread: "Optional[threading.Thread]" = None,
+) -> "Tuple[Optional[int], Optional[str]]":
+    """
+    Try to get the id of the current thread, with various fall backs.
+    """
+
+    # if a thread is specified, that takes priority
+    if thread is not None:
+        try:
+            thread_id = thread.ident
+            thread_name = thread.name
+            if thread_id is not None:
+                return thread_id, thread_name
+        except AttributeError:
+            pass
+
+    # if the app is using gevent, we should look at the gevent hub first
+    # as the id there differs from what the threading module reports
+    if is_gevent():
+        gevent_hub = get_gevent_hub()
+        if gevent_hub is not None:
+            try:
+                # this is undocumented, so wrap it in try except to be safe
+                return gevent_hub.thread_ident, None
+            except AttributeError:
+                pass
+
+    # use the current thread's id if possible
+    try:
+        thread = threading.current_thread()
+        thread_id = thread.ident
+        thread_name = thread.name
+        if thread_id is not None:
+            return thread_id, thread_name
+    except AttributeError:
+        pass
+
+    # if we can't get the current thread id, fall back to the main thread id
+    try:
+        thread = threading.main_thread()
+        thread_id = thread.ident
+        thread_name = thread.name
+        if thread_id is not None:
+            return thread_id, thread_name
+    except AttributeError:
+        pass
+
+    # we've tried everything, time to give up
+    return None, None
+
+
+def should_be_treated_as_error(ty: "Any", value: "Any") -> bool:
+    if ty == SystemExit and hasattr(value, "code") and value.code in (0, None):
+        # https://docs.python.org/3/library/exceptions.html#SystemExit
+        return False
+
+    return True
+
+
+if TYPE_CHECKING:
+    T = TypeVar("T")
+
+
+def try_convert(convert_func: "Callable[[Any], T]", value: "Any") -> "Optional[T]":
+    """
+    Attempt to convert from an unknown type to a specific type, using the
+    given function. Return None if the conversion fails, i.e. if the function
+    raises an exception.
+    """
+    try:
+        if isinstance(value, convert_func):  # type: ignore
+            return value
+    except TypeError:
+        pass
+
+    try:
+        return convert_func(value)
+    except Exception:
+        return None
+
+
+def safe_serialize(data: "Any") -> str:
+    """Safely serialize to a readable string."""
+
+    def serialize_item(
+        item: "Any",
+    ) -> "Union[str, dict[Any, Any], list[Any], tuple[Any, ...]]":
+        if callable(item):
+            try:
+                module = getattr(item, "__module__", None)
+                qualname = getattr(item, "__qualname__", None)
+                name = getattr(item, "__name__", "anonymous")
+
+                if module and qualname:
+                    full_path = f"{module}.{qualname}"
+                elif module and name:
+                    full_path = f"{module}.{name}"
+                else:
+                    full_path = name
+
+                return f""
+            except Exception:
+                return f""
+        elif isinstance(item, dict):
+            return {k: serialize_item(v) for k, v in item.items()}
+        elif isinstance(item, (list, tuple)):
+            return [serialize_item(x) for x in item]
+        elif hasattr(item, "__dict__"):
+            try:
+                attrs = {
+                    k: serialize_item(v)
+                    for k, v in vars(item).items()
+                    if not k.startswith("_")
+                }
+                return f"<{type(item).__name__} {attrs}>"
+            except Exception:
+                return repr(item)
+        else:
+            return item
+
+    try:
+        serialized = serialize_item(data)
+        return json.dumps(serialized, default=str)
+    except Exception:
+        return str(data)
+
+
+def has_logs_enabled(options: "Optional[dict[str, Any]]") -> bool:
+    if options is None:
+        return False
+
+    return bool(
+        options.get("enable_logs", False)
+        or options["_experiments"].get("enable_logs", False)
+    )
+
+
+def get_before_send_log(
+    options: "Optional[dict[str, Any]]",
+) -> "Optional[Callable[[Log, Hint], Optional[Log]]]":
+    if options is None:
+        return None
+
+    return options.get("before_send_log") or options["_experiments"].get(
+        "before_send_log"
+    )
+
+
+def has_metrics_enabled(options: "Optional[dict[str, Any]]") -> bool:
+    if options is None:
+        return False
+
+    return bool(options.get("enable_metrics", True))
+
+
+def get_before_send_metric(
+    options: "Optional[dict[str, Any]]",
+) -> "Optional[Callable[[Metric, Hint], Optional[Metric]]]":
+    if options is None:
+        return None
+
+    return options.get("before_send_metric") or options["_experiments"].get(
+        "before_send_metric"
+    )
+
+
+def format_attribute(val: "Any") -> "AttributeValue":
+    """
+    Turn unsupported attribute value types into an AttributeValue.
+
+    We do this as soon as a user-provided attribute is set, to prevent spans,
+    logs, metrics and similar from having live references to various objects.
+
+    Note: This is not the final attribute value format. Before they're sent,
+    they're serialized further into the actual format the protocol expects:
+    https://develop.sentry.dev/sdk/telemetry/attributes/
+    """
+    if isinstance(val, (bool, int, float, str)):
+        return val
+
+    # Only lists of elements of a single type are supported
+    list_types: 'dict[type, Literal["string[]", "integer[]", "double[]", "boolean[]"]]' = {
+        str: "string[]",
+        int: "integer[]",
+        float: "double[]",
+        bool: "boolean[]",
+    }
+
+    if isinstance(val, (list, tuple)) and not val:
+        return []
+    elif isinstance(val, list):
+        ty = type(val[0])
+        if ty in list_types and all(type(v) is ty for v in val):
+            return copy.deepcopy(val)
+    elif isinstance(val, tuple):
+        ty = type(val[0])
+        if ty in list_types and all(type(v) is ty for v in val):
+            return list(val)
+
+    return safe_repr(val)
+
+
+def serialize_attribute(val: "AttributeValue") -> "SerializedAttributeValue":
+    """Serialize attribute value to the transport format."""
+    if isinstance(val, bool):
+        return {"value": val, "type": "boolean"}
+    if isinstance(val, int):
+        return {"value": val, "type": "integer"}
+    if isinstance(val, float):
+        return {"value": val, "type": "double"}
+    if isinstance(val, str):
+        return {"value": val, "type": "string"}
+
+    if isinstance(val, list):
+        if not val:
+            return {"value": [], "type": "array"}
+
+        # Only lists of elements of a single type are supported
+        ty = type(val[0])
+        if ty in (int, str, bool, float) and all(type(v) is ty for v in val):
+            return {"value": val, "type": "array"}
+
+    # Coerce to string if we don't know what to do with the value. This should
+    # never happen as we pre-format early in format_attribute, but let's be safe.
+    return {"value": safe_repr(val), "type": "string"}
diff --git a/lib/python3.12/site-packages/sentry_sdk/worker.py b/lib/python3.12/site-packages/sentry_sdk/worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d85a653d66f2bcef37cdf32b1876a2504b043aa
--- /dev/null
+++ b/lib/python3.12/site-packages/sentry_sdk/worker.py
@@ -0,0 +1,130 @@
+import os
+import threading
+
+from time import sleep, time
+from sentry_sdk._queue import Queue, FullError
+from sentry_sdk.utils import logger
+from sentry_sdk.consts import DEFAULT_QUEUE_SIZE
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Optional
+    from typing import Callable
+
+
+_TERMINATOR = object()
+
+
+class BackgroundWorker:
+    def __init__(self, queue_size: int = DEFAULT_QUEUE_SIZE) -> None:
+        self._queue: "Queue" = Queue(queue_size)
+        self._lock = threading.Lock()
+        self._thread: "Optional[threading.Thread]" = None
+        self._thread_for_pid: "Optional[int]" = None
+
+    @property
+    def is_alive(self) -> bool:
+        if self._thread_for_pid != os.getpid():
+            return False
+        if not self._thread:
+            return False
+        return self._thread.is_alive()
+
+    def _ensure_thread(self) -> None:
+        if not self.is_alive:
+            self.start()
+
+    def _timed_queue_join(self, timeout: float) -> bool:
+        deadline = time() + timeout
+        queue = self._queue
+
+        queue.all_tasks_done.acquire()
+
+        try:
+            while queue.unfinished_tasks:
+                delay = deadline - time()
+                if delay <= 0:
+                    return False
+                queue.all_tasks_done.wait(timeout=delay)
+
+            return True
+        finally:
+            queue.all_tasks_done.release()
+
+    def start(self) -> None:
+        with self._lock:
+            if not self.is_alive:
+                self._thread = threading.Thread(
+                    target=self._target, name="sentry-sdk.BackgroundWorker"
+                )
+                self._thread.daemon = True
+                try:
+                    self._thread.start()
+                    self._thread_for_pid = os.getpid()
+                except RuntimeError:
+                    # At this point we can no longer start because the interpreter
+                    # is already shutting down.  Sadly at this point we can no longer
+                    # send out events.
+                    self._thread = None
+
+    def kill(self) -> None:
+        """
+        Kill worker thread. Returns immediately. Not useful for
+        waiting on shutdown for events, use `flush` for that.
+        """
+        logger.debug("background worker got kill request")
+        with self._lock:
+            if self._thread:
+                try:
+                    self._queue.put_nowait(_TERMINATOR)
+                except FullError:
+                    logger.debug("background worker queue full, kill failed")
+
+                self._thread = None
+                self._thread_for_pid = None
+
+    def flush(self, timeout: float, callback: "Optional[Any]" = None) -> None:
+        logger.debug("background worker got flush request")
+        with self._lock:
+            if self.is_alive and timeout > 0.0:
+                self._wait_flush(timeout, callback)
+        logger.debug("background worker flushed")
+
+    def full(self) -> bool:
+        return self._queue.full()
+
+    def _wait_flush(self, timeout: float, callback: "Optional[Any]") -> None:
+        initial_timeout = min(0.1, timeout)
+        if not self._timed_queue_join(initial_timeout):
+            pending = self._queue.qsize() + 1
+            logger.debug("%d event(s) pending on flush", pending)
+            if callback is not None:
+                callback(pending, timeout)
+
+            if not self._timed_queue_join(timeout - initial_timeout):
+                pending = self._queue.qsize() + 1
+                logger.error("flush timed out, dropped %s events", pending)
+
+    def submit(self, callback: "Callable[[], None]") -> bool:
+        self._ensure_thread()
+        try:
+            self._queue.put_nowait(callback)
+            return True
+        except FullError:
+            return False
+
+    def _target(self) -> None:
+        while True:
+            callback = self._queue.get()
+            try:
+                if callback is _TERMINATOR:
+                    break
+                try:
+                    callback()
+                except Exception:
+                    logger.error("Failed processing job", exc_info=True)
+            finally:
+                self._queue.task_done()
+            sleep(0)
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/INSTALLER b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/INSTALLER
new file mode 100644
index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/INSTALLER
@@ -0,0 +1 @@
+pip
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/METADATA b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..372fcbd14b5f8772f072adcecbe66a02a6d3b459
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/METADATA
@@ -0,0 +1,171 @@
+Metadata-Version: 2.4
+Name: tiktoken
+Version: 0.12.0
+Summary: tiktoken is a fast BPE tokeniser for use with OpenAI's models
+Author: Shantanu Jain
+Author-email: shantanu@openai.com
+License: MIT License
+        
+        Copyright (c) 2022 OpenAI, Shantanu Jain
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: homepage, https://github.com/openai/tiktoken
+Project-URL: repository, https://github.com/openai/tiktoken
+Project-URL: changelog, https://github.com/openai/tiktoken/blob/main/CHANGELOG.md
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: regex>=2022.1.18
+Requires-Dist: requests>=2.26.0
+Provides-Extra: blobfile
+Requires-Dist: blobfile>=2; extra == "blobfile"
+Dynamic: license-file
+
+# ⏳ tiktoken
+
+tiktoken is a fast [BPE](https://en.wikipedia.org/wiki/Byte_pair_encoding) tokeniser for use with
+OpenAI's models.
+
+```python
+import tiktoken
+enc = tiktoken.get_encoding("o200k_base")
+assert enc.decode(enc.encode("hello world")) == "hello world"
+
+# To get the tokeniser corresponding to a specific model in the OpenAI API:
+enc = tiktoken.encoding_for_model("gpt-4o")
+```
+
+The open source version of `tiktoken` can be installed from [PyPI](https://pypi.org/project/tiktoken):
+```
+pip install tiktoken
+```
+
+The tokeniser API is documented in `tiktoken/core.py`.
+
+Example code using `tiktoken` can be found in the
+[OpenAI Cookbook](https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb).
+
+
+## Performance
+
+`tiktoken` is between 3-6x faster than a comparable open source tokeniser:
+
+![image](https://raw.githubusercontent.com/openai/tiktoken/main/perf.svg)
+
+Performance measured on 1GB of text using the GPT-2 tokeniser, using `GPT2TokenizerFast` from
+`tokenizers==0.13.2`, `transformers==4.24.0` and `tiktoken==0.2.0`.
+
+
+## Getting help
+
+Please post questions in the [issue tracker](https://github.com/openai/tiktoken/issues).
+
+If you work at OpenAI, make sure to check the internal documentation or feel free to contact
+@shantanu.
+
+## What is BPE anyway?
+
+Language models don't see text like you and I, instead they see a sequence of numbers (known as tokens).
+Byte pair encoding (BPE) is a way of converting text into tokens. It has a couple desirable
+properties:
+1) It's reversible and lossless, so you can convert tokens back into the original text
+2) It works on arbitrary text, even text that is not in the tokeniser's training data
+3) It compresses the text: the token sequence is shorter than the bytes corresponding to the
+   original text. On average, in practice, each token corresponds to about 4 bytes.
+4) It attempts to let the model see common subwords. For instance, "ing" is a common subword in
+   English, so BPE encodings will often split "encoding" into tokens like "encod" and "ing"
+   (instead of e.g. "enc" and "oding"). Because the model will then see the "ing" token again and
+   again in different contexts, it helps models generalise and better understand grammar.
+
+`tiktoken` contains an educational submodule that is friendlier if you want to learn more about
+the details of BPE, including code that helps visualise the BPE procedure:
+```python
+from tiktoken._educational import *
+
+# Train a BPE tokeniser on a small amount of text
+enc = train_simple_encoding()
+
+# Visualise how the GPT-4 encoder encodes text
+enc = SimpleBytePairEncoding.from_tiktoken("cl100k_base")
+enc.encode("hello world aaaaaaaaaaaa")
+```
+
+
+## Extending tiktoken
+
+You may wish to extend `tiktoken` to support new encodings. There are two ways to do this.
+
+
+**Create your `Encoding` object exactly the way you want and simply pass it around.**
+
+```python
+cl100k_base = tiktoken.get_encoding("cl100k_base")
+
+# In production, load the arguments directly instead of accessing private attributes
+# See openai_public.py for examples of arguments for specific encodings
+enc = tiktoken.Encoding(
+    # If you're changing the set of special tokens, make sure to use a different name
+    # It should be clear from the name what behaviour to expect.
+    name="cl100k_im",
+    pat_str=cl100k_base._pat_str,
+    mergeable_ranks=cl100k_base._mergeable_ranks,
+    special_tokens={
+        **cl100k_base._special_tokens,
+        "<|im_start|>": 100264,
+        "<|im_end|>": 100265,
+    }
+)
+```
+
+**Use the `tiktoken_ext` plugin mechanism to register your `Encoding` objects with `tiktoken`.**
+
+This is only useful if you need `tiktoken.get_encoding` to find your encoding, otherwise prefer
+option 1.
+
+To do this, you'll need to create a namespace package under `tiktoken_ext`.
+
+Layout your project like this, making sure to omit the `tiktoken_ext/__init__.py` file:
+```
+my_tiktoken_extension
+├── tiktoken_ext
+│   └── my_encodings.py
+└── setup.py
+```
+
+`my_encodings.py` should be a module that contains a variable named `ENCODING_CONSTRUCTORS`.
+This is a dictionary from an encoding name to a function that takes no arguments and returns
+arguments that can be passed to `tiktoken.Encoding` to construct that encoding. For an example, see
+`tiktoken_ext/openai_public.py`. For precise details, see `tiktoken/registry.py`.
+
+Your `setup.py` should look something like this:
+```python
+from setuptools import setup, find_namespace_packages
+
+setup(
+    name="my_tiktoken_extension",
+    packages=find_namespace_packages(include=['tiktoken_ext*']),
+    install_requires=["tiktoken"],
+    ...
+)
+```
+
+Then simply `pip install ./my_tiktoken_extension` and you should be able to use your
+custom encodings! Make sure **not** to use an editable install.
+
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/RECORD b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/RECORD
new file mode 100644
index 0000000000000000000000000000000000000000..628766c30f3eb645b441edb15a666814144c2f3a
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/RECORD
@@ -0,0 +1,22 @@
+tiktoken-0.12.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+tiktoken-0.12.0.dist-info/METADATA,sha256=07KObsjwnEfLTQ-oRD0vhrE3Zw-oNMJqTlBbEeJxtZ8,6688
+tiktoken-0.12.0.dist-info/RECORD,,
+tiktoken-0.12.0.dist-info/WHEEL,sha256=VXvNKn6nFeCM45GEUrNLJOO_J_e-cNJphGt9rWFxyE0,113
+tiktoken-0.12.0.dist-info/licenses/LICENSE,sha256=QYy0mbQ2Eo1lPXmUEzOlQ3t74uqSE9zC8E0V1dLFHYY,1078
+tiktoken-0.12.0.dist-info/top_level.txt,sha256=54G5MceQnuD7EXvp7jzGxDDapA1iOwsh77jhCN9WKkc,22
+tiktoken/__init__.py,sha256=eHlkakibO43-11JFQJUgpC8z2v4ID1r3l3LXjMyEwKc,346
+tiktoken/__pycache__/__init__.cpython-312.pyc,,
+tiktoken/__pycache__/_educational.cpython-312.pyc,,
+tiktoken/__pycache__/core.cpython-312.pyc,,
+tiktoken/__pycache__/load.cpython-312.pyc,,
+tiktoken/__pycache__/model.cpython-312.pyc,,
+tiktoken/__pycache__/registry.cpython-312.pyc,,
+tiktoken/_educational.py,sha256=TUFOp8Q91WjrTvGKhCNEyrhtva82UlenXfhPy9zS7VQ,8229
+tiktoken/_tiktoken.cpython-312-x86_64-linux-gnu.so,sha256=qCn0iO_VQ7YJKD6D8jDXz6WHH64mSUK59dicngi37S8,3525056
+tiktoken/core.py,sha256=TCwORlettZl-da55Ysp52TlLk18nKD6e62Q_0ZFA404,17458
+tiktoken/load.py,sha256=dhTOiVIInbhiQ_zmtOZDshKvqSKzXyNOJJWPmJ0S9RU,5919
+tiktoken/model.py,sha256=d57kixsksIv6VESndVjvmGBRj8LrSFGAwUCV5xZtxRk,4061
+tiktoken/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tiktoken/registry.py,sha256=7fktZbJ1Kcm8sVyWgEfIy-ZxfUvcXupLUNXKPfSGwQU,3256
+tiktoken_ext/__pycache__/openai_public.cpython-312.pyc,,
+tiktoken_ext/openai_public.py,sha256=lUOSc45g0Pttyh2tgIcu_EfI4nM7q-y78KI5cO1mwss,5613
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/WHEEL b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/WHEEL
new file mode 100644
index 0000000000000000000000000000000000000000..d3bb8ca34441955086c902160632acd0302490d5
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/WHEEL
@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: false
+Tag: cp312-cp312-manylinux_2_28_x86_64
+
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/licenses/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..83ed1036f70d4f419307e8a044a35e163cc35201
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/licenses/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 OpenAI, Shantanu Jain
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/top_level.txt b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/top_level.txt
new file mode 100644
index 0000000000000000000000000000000000000000..859880ea40d957171e49537207249fdca13946cb
--- /dev/null
+++ b/lib/python3.12/site-packages/tiktoken-0.12.0.dist-info/top_level.txt
@@ -0,0 +1,2 @@
+tiktoken
+tiktoken_ext
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/INSTALLER b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/INSTALLER
new file mode 100644
index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/INSTALLER
@@ -0,0 +1 @@
+pip
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/LICENSE b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..1bec23eaf1dd562ae3d3216420b1b1bbfbd39cbc
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/LICENSE
@@ -0,0 +1,25 @@
+BSD 2-Clause License
+
+Copyright (c) 2017 Facebook Inc. (Soumith Chintala), 
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/METADATA b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..88b509dd116a4c3396d57eca6af5d3b6bd126fa3
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/METADATA
@@ -0,0 +1,124 @@
+Metadata-Version: 2.2
+Name: torchaudio
+Version: 2.6.0
+Summary: An audio package for PyTorch
+Home-page: https://github.com/pytorch/audio
+Author: Soumith Chintala, David Pollack, Sean Naren, Peter Goldsborough, Moto Hira, Caroline Chen, Jeff Hwang, Zhaoheng Ni, Xiaohui Zhang
+Author-email: soumith@pytorch.org
+Maintainer: Moto Hira, Caroline Chen, Jeff Hwang, Zhaoheng Ni, Xiaohui Zhang
+Maintainer-email: moto@meta.com
+Classifier: Environment :: Plugins
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX
+Classifier: Programming Language :: C++
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Multimedia :: Sound/Audio
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch==2.6.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: maintainer
+Dynamic: maintainer-email
+Dynamic: requires-dist
+Dynamic: summary
+
+torchaudio: an audio library for PyTorch
+========================================
+
+[![Documentation](https://img.shields.io/badge/dynamic/json.svg?label=docs&url=https%3A%2F%2Fpypi.org%2Fpypi%2Ftorchaudio%2Fjson&query=%24.info.version&colorB=brightgreen&prefix=v)](https://pytorch.org/audio/main/)
+[![Anaconda Badge](https://anaconda.org/pytorch/torchaudio/badges/downloads.svg)](https://anaconda.org/pytorch/torchaudio)
+[![Anaconda-Server Badge](https://anaconda.org/pytorch/torchaudio/badges/platforms.svg)](https://anaconda.org/pytorch/torchaudio)
+
+![TorchAudio Logo](docs/source/_static/img/logo.png)
+
+The aim of torchaudio is to apply [PyTorch](https://github.com/pytorch/pytorch) to
+the audio domain. By supporting PyTorch, torchaudio follows the same philosophy
+of providing strong GPU acceleration, having a focus on trainable features through
+the autograd system, and having consistent style (tensor names and dimension names).
+Therefore, it is primarily a machine learning library and not a general signal
+processing library. The benefits of PyTorch can be seen in torchaudio through
+having all the computations be through PyTorch operations which makes it easy
+to use and feel like a natural extension.
+
+- [Support audio I/O (Load files, Save files)](http://pytorch.org/audio/main/)
+  - Load a variety of audio formats, such as `wav`, `mp3`, `ogg`, `flac`, `opus`, `sphere`, into a torch Tensor using SoX
+  - [Kaldi (ark/scp)](http://pytorch.org/audio/main/kaldi_io.html)
+- [Dataloaders for common audio datasets](http://pytorch.org/audio/main/datasets.html)
+- Audio and speech processing functions
+  - [forced_align](https://pytorch.org/audio/main/generated/torchaudio.functional.forced_align.html)
+- Common audio transforms
+  - [Spectrogram, AmplitudeToDB, MelScale, MelSpectrogram, MFCC, MuLawEncoding, MuLawDecoding, Resample](http://pytorch.org/audio/main/transforms.html)
+- Compliance interfaces: Run code using PyTorch that align with other libraries
+  - [Kaldi: spectrogram, fbank, mfcc](https://pytorch.org/audio/main/compliance.kaldi.html)
+
+Installation
+------------
+
+Please refer to https://pytorch.org/audio/main/installation.html for installation and build process of TorchAudio.
+
+
+API Reference
+-------------
+
+API Reference is located here: http://pytorch.org/audio/main/
+
+Contributing Guidelines
+-----------------------
+
+Please refer to [CONTRIBUTING.md](./CONTRIBUTING.md)
+
+Citation
+--------
+
+If you find this package useful, please cite as:
+
+```bibtex
+@article{yang2021torchaudio,
+  title={TorchAudio: Building Blocks for Audio and Speech Processing},
+  author={Yao-Yuan Yang and Moto Hira and Zhaoheng Ni and Anjali Chourdia and Artyom Astafurov and Caroline Chen and Ching-Feng Yeh and Christian Puhrsch and David Pollack and Dmitriy Genzel and Donny Greenberg and Edward Z. Yang and Jason Lian and Jay Mahadeokar and Jeff Hwang and Ji Chen and Peter Goldsborough and Prabhat Roy and Sean Narenthiran and Shinji Watanabe and Soumith Chintala and Vincent Quenneville-Bélair and Yangyang Shi},
+  journal={arXiv preprint arXiv:2110.15018},
+  year={2021}
+}
+```
+
+```bibtex
+@misc{hwang2023torchaudio,
+      title={TorchAudio 2.1: Advancing speech recognition, self-supervised learning, and audio processing components for PyTorch}, 
+      author={Jeff Hwang and Moto Hira and Caroline Chen and Xiaohui Zhang and Zhaoheng Ni and Guangzhi Sun and Pingchuan Ma and Ruizhe Huang and Vineel Pratap and Yuekai Zhang and Anurag Kumar and Chin-Yun Yu and Chuang Zhu and Chunxi Liu and Jacob Kahn and Mirco Ravanelli and Peng Sun and Shinji Watanabe and Yangyang Shi and Yumeng Tao and Robin Scheibler and Samuele Cornell and Sean Kim and Stavros Petridis},
+      year={2023},
+      eprint={2310.17864},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS}
+}
+```
+
+Disclaimer on Datasets
+----------------------
+
+This is a utility library that downloads and prepares public datasets. We do not host or distribute these datasets, vouch for their quality or fairness, or claim that you have license to use the dataset. It is your responsibility to determine whether you have permission to use the dataset under the dataset's license.
+
+If you're a dataset owner and wish to update any part of it (description, citation, etc.), or do not want your dataset to be included in this library, please get in touch through a GitHub issue. Thanks for your contribution to the ML community!
+
+Pre-trained Model License
+-------------------------
+
+The pre-trained models provided in this library may have their own licenses or terms and conditions derived from the dataset used for training. It is your responsibility to determine whether you have permission to use the models for your use case.
+
+For instance, SquimSubjective model is released under the Creative Commons Attribution Non Commercial 4.0 International (CC-BY-NC 4.0) license. See [the link](https://zenodo.org/record/4660670#.ZBtWPOxuerN) for additional details.
+
+Other pre-trained models that have different license are noted in documentation. Please checkout the [documentation page](https://pytorch.org/audio/main/).
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/RECORD b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/RECORD
new file mode 100644
index 0000000000000000000000000000000000000000..eef95d1f20dcf3cd596310fd68f1118c37538815
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/RECORD
@@ -0,0 +1,281 @@
+torchaudio-2.6.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+torchaudio-2.6.0.dist-info/LICENSE,sha256=k6WIYahYzBCOa2uDPgjnbosqZjOeSoAHyKWowf-cQNY,1338
+torchaudio-2.6.0.dist-info/METADATA,sha256=6Mt_Z5EGlz4hs3HGs2ekPhnD1m_6FX7iXryUcguI_a8,6618
+torchaudio-2.6.0.dist-info/RECORD,,
+torchaudio-2.6.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchaudio-2.6.0.dist-info/WHEEL,sha256=aDUArjEbIuDQzblE24IvdIsBzRhJEjiXmuNUYQ-gpaA,104
+torchaudio-2.6.0.dist-info/top_level.txt,sha256=GT0MktEbHKoLnvd-6ii7_dhJVvshupOujk840BcHU4U,17
+torchaudio/__init__.py,sha256=VSnZ6s4e5clAj7f7aCXBZt9amskeXg1j19txAQBQ2Iw,892
+torchaudio/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/__pycache__/kaldi_io.cpython-312.pyc,,
+torchaudio/__pycache__/version.cpython-312.pyc,,
+torchaudio/_backend/__init__.py,sha256=6zMYGajHaaCXUE_U7HuGLp0fqcviYAjBZdFDI4E7C-0,1631
+torchaudio/_backend/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/backend.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/common.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/ffmpeg.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/soundfile.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/soundfile_backend.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/sox.cpython-312.pyc,,
+torchaudio/_backend/__pycache__/utils.cpython-312.pyc,,
+torchaudio/_backend/backend.py,sha256=hSrfZcj5FMzx5ZpwubN-LLMvBFb7ENyw7HvT_6pVYVU,1565
+torchaudio/_backend/common.py,sha256=55Y0r0MsdW6gvTOT_Zy60UGFXc60DfdJ7uvycJKK3is,1783
+torchaudio/_backend/ffmpeg.py,sha256=oL_whDjkPtHzo6HJLiEPlGHdrOqzjlu81g-vlaNkRBA,11294
+torchaudio/_backend/soundfile.py,sha256=n0Epw0J9rBb89xVJWTXaDfb96YFz0-i2xarXIdDd-Cw,1703
+torchaudio/_backend/soundfile_backend.py,sha256=qJHEEXU1egCkPJ2Y9uJWFvVhW3AqDZ7z7P7mkJjJJWM,17376
+torchaudio/_backend/sox.py,sha256=p_y9bXKz_6Hto5LORGrHXVbXNS7nsvWUc9iucbN-tCA,3360
+torchaudio/_backend/utils.py,sha256=Q_RgMaeKFvwOoVdWfdwnL0CmpQli_tmi4wPQ2RRHyRA,13299
+torchaudio/_extension/__init__.py,sha256=lQPB8K7VSxWmtTEiMmF-u7WVq1O10_t5nEghkjCf4Ks,2202
+torchaudio/_extension/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/_extension/__pycache__/utils.cpython-312.pyc,,
+torchaudio/_extension/utils.py,sha256=4FTD6xwcSLqVJ3Kmpx5cvJp1oAUKmWwRjwuxpcbrmzw,6258
+torchaudio/_internal/__init__.py,sha256=gjU8g9HhVd9hHrHXJM0xOlZL6cT8ktO60MN8RHI6ZbA,241
+torchaudio/_internal/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/_internal/__pycache__/module_utils.cpython-312.pyc,,
+torchaudio/_internal/module_utils.py,sha256=SJr-RS6hs6uJkIVx_WZwsFPKUjtuG6lsfw3uI0UItDE,3562
+torchaudio/backend/__init__.py,sha256=AL8njOL5hDhIGq5tjRxfFzZXxQdTGlz5gs9g4RToblY,281
+torchaudio/backend/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/backend/__pycache__/_no_backend.cpython-312.pyc,,
+torchaudio/backend/__pycache__/_sox_io_backend.cpython-312.pyc,,
+torchaudio/backend/__pycache__/common.cpython-312.pyc,,
+torchaudio/backend/__pycache__/no_backend.cpython-312.pyc,,
+torchaudio/backend/__pycache__/soundfile_backend.cpython-312.pyc,,
+torchaudio/backend/__pycache__/sox_io_backend.cpython-312.pyc,,
+torchaudio/backend/_no_backend.py,sha256=9Ss3b4BMFao5Kfdqh6S8JSLUoYCodbPgNQCiDHbNhDQ,757
+torchaudio/backend/_sox_io_backend.py,sha256=PnH-ClsiOy0ekOTY1RKB-cL6xrTPtrzmuXGX3ugATps,11456
+torchaudio/backend/common.py,sha256=T_iYc4u_EzfIh7zbG_xW052fyJMXUXEpPfDOaAQ6sAY,443
+torchaudio/backend/no_backend.py,sha256=PBEQ9vFG5uVurktjxRAiEqSuVJxImnMyPQlt0reRpP0,469
+torchaudio/backend/soundfile_backend.py,sha256=2Tyh5yAn7LQqKzeqW-rx4o2QbmmUrocmh3iYPnuAds0,499
+torchaudio/backend/sox_io_backend.py,sha256=XsAB5HkRbI9-W2nXx-yUMUPJP2Ca5sd09TLywrQ2N-E,477
+torchaudio/compliance/__init__.py,sha256=hhNObUS0c-fS-VMudM7zl3-CvupvCDmESlikntSMn5g,48
+torchaudio/compliance/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/compliance/__pycache__/kaldi.cpython-312.pyc,,
+torchaudio/compliance/kaldi.py,sha256=XL6hpYTd6nSPb2imIdeU4TM06I2fqh1AmG968y8ZbSk,36666
+torchaudio/datasets/__init__.py,sha256=taRr3duDaEK1Pfzj9N1dFuZpXfy8e4uFItcJiRLAQwQ,1171
+torchaudio/datasets/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/cmuarctic.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/cmudict.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/commonvoice.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/dr_vctk.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/fluentcommands.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/gtzan.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/iemocap.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/librilight_limited.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/librimix.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/librispeech.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/librispeech_biasing.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/libritts.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/ljspeech.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/musdb_hq.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/quesst14.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/snips.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/speechcommands.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/tedlium.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/utils.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/vctk.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/voxceleb1.cpython-312.pyc,,
+torchaudio/datasets/__pycache__/yesno.cpython-312.pyc,,
+torchaudio/datasets/cmuarctic.py,sha256=KAhTHUJ3g5RSlmsU5mCTcvutOCm3Oqcd3643u3HNqIg,7097
+torchaudio/datasets/cmudict.py,sha256=9OEpNDYpyqeEyinAnyGIU8FampDj7ziSOHRwJLIlq2M,5990
+torchaudio/datasets/commonvoice.py,sha256=9khedUCmdEkCKPU6_r8VWz6I2VdJokatuziZ6BxJMZs,2763
+torchaudio/datasets/dr_vctk.py,sha256=Km4-tKllAgnOKCuq66YRWhTlNWmC7D0Xz3dAttRRGSo,4377
+torchaudio/datasets/fluentcommands.py,sha256=u3tkO4-AAaTWdbRQi6lIvad4x2plZgXM39KljGtmRsw,3245
+torchaudio/datasets/gtzan.py,sha256=I5dRP_QGuQ1joXWRwZwtvpwi22uZTb8QZm9Mr2W55Mg,24357
+torchaudio/datasets/iemocap.py,sha256=X_WCoXOzRqcWRRRoUtY0AlD9SJcUUOACIcgbV0irt48,4930
+torchaudio/datasets/librilight_limited.py,sha256=fAwpX0hEMze5aV57BP7rjBLwRiZa3Aje_NXi_3o16wA,4179
+torchaudio/datasets/librimix.py,sha256=VtKOhf6VJc1ysWCvUvh0SbtjOkXJChmBM_BhoSkg_2A,5116
+torchaudio/datasets/librispeech.py,sha256=zkzJFWchWs4AktYAI-ghmWH4ZeJ84C0uDo9E1_pTgSI,6308
+torchaudio/datasets/librispeech_biasing.py,sha256=d-02tyrXI-CSGbXBFYFcnM_yT8WSGABHfpNiFxyadL0,6958
+torchaudio/datasets/libritts.py,sha256=EtWOoCDz7_qGLZF5YcZfnHaLxH4Y8QJCnopafLiqFno,5870
+torchaudio/datasets/ljspeech.py,sha256=92NeLQsC1iKpqfiMkKKbcJDpaYdZKVdVEBQJze1wmxY,3494
+torchaudio/datasets/musdb_hq.py,sha256=TYKjpat6JKr9bkFqUecu7_hRdshRfQP2UbknaYR3Q0U,5075
+torchaudio/datasets/quesst14.py,sha256=QyGd4fMS820ATbP8YgBtu7bSSK09pw5RZklsPJ8Jf0Y,4455
+torchaudio/datasets/snips.py,sha256=WaYUknGFM3rnLklOj5ZYHSX5mhlf_Ce4p3LBZdA9yJc,5008
+torchaudio/datasets/speechcommands.py,sha256=cLSgiVYlQjEOuYPpFeAtcXSGirraH4IMoP8p9WIvUoY,7481
+torchaudio/datasets/tedlium.py,sha256=a8Hf2QvOki7_chgXcMAFMk-piTjodktfnc3HRbUVJkU,8698
+torchaudio/datasets/utils.py,sha256=QaI02lOcesy6Dnvlof4BeTDIbiOqUcoVEPxL5_T8vwU,1689
+torchaudio/datasets/vctk.py,sha256=twR_n8LyQcT8A_HrJoMx3RkaVrRXXZAnIVU1d0E0npQ,5699
+torchaudio/datasets/voxceleb1.py,sha256=9vU0ftB4-2usO8ZiEUKR_IQTEdHhA0M8l9scXCNehnw,11725
+torchaudio/datasets/yesno.py,sha256=4sgfMeSxz8HaRDk6A2UIFP-20q29MwEO_r8DoEtfbvE,3026
+torchaudio/functional/__init__.py,sha256=l-gC2WyY5COabU0lhkUS8EnwOYdEYR_6234OyoAIgnU,2357
+torchaudio/functional/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/functional/__pycache__/_alignment.cpython-312.pyc,,
+torchaudio/functional/__pycache__/filtering.cpython-312.pyc,,
+torchaudio/functional/__pycache__/functional.cpython-312.pyc,,
+torchaudio/functional/_alignment.py,sha256=wmDeohWvuoYORYDeIRxnYUhUqv1uCUkaCZYLEK_ryUg,4695
+torchaudio/functional/filtering.py,sha256=EdYtv2z893Qi58BHIR1VGDfRaGCo0sIKl4k98-vwPkg,61554
+torchaudio/functional/functional.py,sha256=c8qr3mmPXLi40N4NCLcpHvQeUIuoNtbszksjtruC15g,96006
+torchaudio/io/__init__.py,sha256=8nd6s_xuBh5iVzIvQ-qNlforukZzuCx36DyvCmHK748,297
+torchaudio/io/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/io/__pycache__/_effector.cpython-312.pyc,,
+torchaudio/io/__pycache__/_playback.cpython-312.pyc,,
+torchaudio/io/_effector.py,sha256=APDrIU2biwFsSVmhrXjelmc4ndcmb0JL-H189Zp689g,11870
+torchaudio/io/_playback.py,sha256=70IxGrGPkI1h4rz8_04SFCGsbbGZkTiUdRhbPOMLLgQ,2321
+torchaudio/kaldi_io.py,sha256=TwS2YgSLlJwOXjNNsHBuXyxhKeKKpptVHLBV7QYZCas,5073
+torchaudio/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchaudio/lib/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/lib/_torchaudio.so,sha256=AlinjlumUmVSFsGSMSsaERdSXL1EBvQvEBQFBumDJtI,134128
+torchaudio/lib/_torchaudio_sox.so,sha256=C6vXOy-_I98j0vIKPKmSipBPeVrsRpgCZv47PjzrKN0,261976
+torchaudio/lib/libctc_prefix_decoder.so,sha256=S5AfYy3KhQDs-A6HVjW71PeRFvh7sJZkPunr_B8i1To,4943448
+torchaudio/lib/libtorchaudio.so,sha256=QexwIXeMOZQIbsOTcU4xyuY2BfBVdw0XQL5W1cIGpWk,2447512
+torchaudio/lib/libtorchaudio_sox.so,sha256=kef-ms0zy2TM5qylaSjXRBut5gxKY4CncmFe8uLZ8vA,151736
+torchaudio/lib/pybind11_prefixctc.so,sha256=IN7RtBjwhNr7ToXGRa1aH-GZMdWkP4GE8Kcv2ZqL7_4,214040
+torchaudio/models/__init__.py,sha256=BNMNGuwpJAFRsdtwHYQ6slGClkrUTu31_7mXh7FjeV4,1995
+torchaudio/models/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/models/__pycache__/_hdemucs.cpython-312.pyc,,
+torchaudio/models/__pycache__/conformer.cpython-312.pyc,,
+torchaudio/models/__pycache__/conv_tasnet.cpython-312.pyc,,
+torchaudio/models/__pycache__/deepspeech.cpython-312.pyc,,
+torchaudio/models/__pycache__/emformer.cpython-312.pyc,,
+torchaudio/models/__pycache__/rnnt.cpython-312.pyc,,
+torchaudio/models/__pycache__/rnnt_decoder.cpython-312.pyc,,
+torchaudio/models/__pycache__/tacotron2.cpython-312.pyc,,
+torchaudio/models/__pycache__/wav2letter.cpython-312.pyc,,
+torchaudio/models/__pycache__/wavernn.cpython-312.pyc,,
+torchaudio/models/_hdemucs.py,sha256=VPnQ73lA9lfAxRjZ85NCGJYP36mPNwTjS-TU4qelu_k,38242
+torchaudio/models/conformer.py,sha256=5IceU-jcZKofkHTTqRKoytubQ75MzZPrPlfkLsIlxeA,10068
+torchaudio/models/conv_tasnet.py,sha256=v-DI_Ej9FCBBbSH-Spkh3tzq8rkBhbQNA-Wp52Uf32E,12540
+torchaudio/models/decoder/__init__.py,sha256=4IS_DyQageh2_uY3YE1aBCYEE3HArCFd8ZUfbgww-Tc,1206
+torchaudio/models/decoder/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/models/decoder/__pycache__/_ctc_decoder.cpython-312.pyc,,
+torchaudio/models/decoder/__pycache__/_cuda_ctc_decoder.cpython-312.pyc,,
+torchaudio/models/decoder/_ctc_decoder.py,sha256=K3gSsG9htU08fe7tKSuIJPDIs7ruY50pJ3eNdNhXSVY,20082
+torchaudio/models/decoder/_cuda_ctc_decoder.py,sha256=rtpN1Z_Xni1LlHgHx6jJ1Jap4TnQ0rRRMvwGWa-xnvA,7186
+torchaudio/models/deepspeech.py,sha256=kQW3B6YcjYuq7xRzWjRJFGr7ZNraY9gMYDTxII7Cgtg,2746
+torchaudio/models/emformer.py,sha256=ncDeEcYegUmIKQoDBoufUhVWj4dYpZAXxLX0qmEqt1A,37766
+torchaudio/models/rnnt.py,sha256=jz66nwDd1qGT6KQR1lbA_urPktygewhm0FH66T7P3Ek,35541
+torchaudio/models/rnnt_decoder.py,sha256=IwlDsuw1SA-uCRrXGMBqm05auGFSha2bZ-8BOImnK0c,12839
+torchaudio/models/squim/__init__.py,sha256=b98nAaL28Q4w3lrqd_6wUd0An-xNhhJn4Tj8oZlzQnc,346
+torchaudio/models/squim/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/models/squim/__pycache__/objective.cpython-312.pyc,,
+torchaudio/models/squim/__pycache__/subjective.cpython-312.pyc,,
+torchaudio/models/squim/objective.py,sha256=YPkEWdDMyeP7GcR0OzUPHr2wKhIKFMjy4peYsABmZQk,12289
+torchaudio/models/squim/subjective.py,sha256=N00kILSPm0akWyNsrNYKmHgZmooo8gbyUm5IVLf7bx8,5797
+torchaudio/models/tacotron2.py,sha256=FimYhGSI8FKwWb87CLk4h3yKWatCU2HvFmU1t5WUn4E,45914
+torchaudio/models/wav2letter.py,sha256=KNcq4p0qZG2Bwfdakv7YwLCvi_yGT-qB4fJwGMuFQhg,3278
+torchaudio/models/wav2vec2/__init__.py,sha256=WlafukV6GwuSNh0CZifrYUt4V5l59kjvGX7AZNonjfk,927
+torchaudio/models/wav2vec2/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/models/wav2vec2/__pycache__/components.cpython-312.pyc,,
+torchaudio/models/wav2vec2/__pycache__/model.cpython-312.pyc,,
+torchaudio/models/wav2vec2/__pycache__/wavlm_attention.cpython-312.pyc,,
+torchaudio/models/wav2vec2/components.py,sha256=DRmW-GHYf-JReCg_0l1ovNWJBnAavePO3S2vPY-1ze4,47077
+torchaudio/models/wav2vec2/model.py,sha256=Z2VN6KbDOOdq5JtP7lxPQebwYqsxKms1Eu4IjDJtZaQ,60092
+torchaudio/models/wav2vec2/utils/__init__.py,sha256=qmMbz4HAN5kEEyl4cSGm_JQZI47beyh4witydPC_qns,181
+torchaudio/models/wav2vec2/utils/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/models/wav2vec2/utils/__pycache__/import_fairseq.cpython-312.pyc,,
+torchaudio/models/wav2vec2/utils/__pycache__/import_huggingface.cpython-312.pyc,,
+torchaudio/models/wav2vec2/utils/import_fairseq.py,sha256=oCwG6qpG0bCXue2V56fjDcC8cA2rgy4b3O_nu_FI9ZY,9198
+torchaudio/models/wav2vec2/utils/import_huggingface.py,sha256=1nVCipp-lOUAyl_-P103DWLUeTOZi9X_ffX93bOXxEk,5946
+torchaudio/models/wav2vec2/wavlm_attention.py,sha256=1DU_pkoLCeHQwSF4lJ06cez0PsMVoXNxiYKP0Yv0qFQ,10844
+torchaudio/models/wavernn.py,sha256=5xUyao5g69jRXX4ReNi4mP_aTSIonJPP6XcPrqKybEk,15446
+torchaudio/pipelines/__init__.py,sha256=Xy8NmInKwTcNBHwLTTjHjrfczRLuQq8a67ENt1OTVXM,2745
+torchaudio/pipelines/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/pipelines/__pycache__/_source_separation_pipeline.cpython-312.pyc,,
+torchaudio/pipelines/__pycache__/_squim_pipeline.cpython-312.pyc,,
+torchaudio/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc,,
+torchaudio/pipelines/_source_separation_pipeline.py,sha256=WxngB1d13H5IVqs4RPqSL53ZGYsJ3tnfCpxgc5FNSOM,4224
+torchaudio/pipelines/_squim_pipeline.py,sha256=eCimTeoqNX8LilIQNGmb8UaRtnLIXa4LNShXFjodcZM,6280
+torchaudio/pipelines/_tts/__init__.py,sha256=PP7l8XzVURqelwuMJFgfOCv4fvzZunDiy90ZQlRkv7g,426
+torchaudio/pipelines/_tts/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/pipelines/_tts/__pycache__/impl.cpython-312.pyc,,
+torchaudio/pipelines/_tts/__pycache__/interface.cpython-312.pyc,,
+torchaudio/pipelines/_tts/__pycache__/utils.cpython-312.pyc,,
+torchaudio/pipelines/_tts/impl.py,sha256=Tig4_5sITJADwxN5eZGek7Ath_-e3sV8CTM5t6UpeUU,15374
+torchaudio/pipelines/_tts/interface.py,sha256=yUaS0UK3PTRruYXRWFil7lAhr-1iYiyBaDBLmEnJPUQ,10224
+torchaudio/pipelines/_tts/utils.py,sha256=0rLyoFWS78n5jn9AC99pgIwAjaXSw-MVbj_pjSaOHiM,4616
+torchaudio/pipelines/_wav2vec2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchaudio/pipelines/_wav2vec2/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/pipelines/_wav2vec2/__pycache__/aligner.cpython-312.pyc,,
+torchaudio/pipelines/_wav2vec2/__pycache__/impl.cpython-312.pyc,,
+torchaudio/pipelines/_wav2vec2/__pycache__/utils.cpython-312.pyc,,
+torchaudio/pipelines/_wav2vec2/aligner.py,sha256=pIWRgQ-kdYUxtL8bdc0qk9wBjwRrHY1uSWL3L4e2vxs,2709
+torchaudio/pipelines/_wav2vec2/impl.py,sha256=zdXFjytJO5MvnB-3aygzUUFKxCTkQGU_OX_rhUh9c0k,65561
+torchaudio/pipelines/_wav2vec2/utils.py,sha256=Q8_fWOR2JDnHu0TTRmHzRjI3BOJa0hGIAl0cjtALgsQ,6971
+torchaudio/pipelines/rnnt_pipeline.py,sha256=Qy37z7v6d1jLOHd67zbRu21dgL6Fml1rTd7j4Jl1NsM,13749
+torchaudio/prototype/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchaudio/prototype/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/datasets/__init__.py,sha256=GSMcp2CykcBc-krhlHTrPm5DCvDFwnA7_6GFNCGwsaQ,47
+torchaudio/prototype/datasets/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/datasets/__pycache__/musan.cpython-312.pyc,,
+torchaudio/prototype/datasets/musan.py,sha256=g68PIJCtJM_mXK8vngJ4PRzMqvp-YShPLN9qTMgeiKw,2096
+torchaudio/prototype/functional/__init__.py,sha256=GlbhnDHcNyUWdRd3R-ATzRkG2FXsbqjL56OptyTXpec,562
+torchaudio/prototype/functional/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/functional/__pycache__/_dsp.cpython-312.pyc,,
+torchaudio/prototype/functional/__pycache__/_rir.cpython-312.pyc,,
+torchaudio/prototype/functional/__pycache__/functional.cpython-312.pyc,,
+torchaudio/prototype/functional/_dsp.py,sha256=H4IZgQYjrmV6ITb7iex3F4qwBSFDyPbdrb0e4ZXbkMY,16638
+torchaudio/prototype/functional/_rir.py,sha256=k-svDQK56U1WNpj4dNUxWArBiVM7sZ_BZ98oOop4NNg,17255
+torchaudio/prototype/functional/functional.py,sha256=xd8ZQe69Utl7HzC-VDyhniS0K-dment-Z7FrEwTrfYk,6464
+torchaudio/prototype/models/__init__.py,sha256=zrav5cgVlM51xvocGlL10GoR2r5UuQrenPDlYRUzv40,1254
+torchaudio/prototype/models/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/_conformer_wav2vec2.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/_emformer_hubert.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/conv_emformer.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/hifi_gan.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/rnnt.cpython-312.pyc,,
+torchaudio/prototype/models/__pycache__/rnnt_decoder.cpython-312.pyc,,
+torchaudio/prototype/models/_conformer_wav2vec2.py,sha256=J7ZJ0dPIFLj9RyPsnuSQC9Y5OVJ9xL6F4JS44zua9zA,29522
+torchaudio/prototype/models/_emformer_hubert.py,sha256=D1WlL1S5xNrN5zOWYnzGyRUtPWnFZOkJprhbkYln0fM,13498
+torchaudio/prototype/models/conv_emformer.py,sha256=tdUz8WwhNlmGXpmki4voZg5nrg749xi23rmfrq2XRCk,23076
+torchaudio/prototype/models/hifi_gan.py,sha256=-ZMA722hoYabIbJl3OGqlxyhhqAHEL7UsTEkOyy8w5I,12480
+torchaudio/prototype/models/rnnt.py,sha256=MTsXxGGv8xOIlH_zhOeSUdedI29CHBIsJ0Pcr8D6yK0,30859
+torchaudio/prototype/models/rnnt_decoder.py,sha256=lIacC7qCjMxjAuBHpgrPXlNI3eERo11fYgaEwPDT7ms,15735
+torchaudio/prototype/pipelines/__init__.py,sha256=yo19xKvIW3XDdDo19thGSMkPRuR8xTwSe0qMWEPS9bE,382
+torchaudio/prototype/pipelines/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/pipelines/__pycache__/hifigan_pipeline.cpython-312.pyc,,
+torchaudio/prototype/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc,,
+torchaudio/prototype/pipelines/_vggish/__init__.py,sha256=yi9HO_14_YWFOEvQOhTXb9eqF3JGJ9FtM5-J-a3nEnA,89
+torchaudio/prototype/pipelines/_vggish/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_impl.cpython-312.pyc,,
+torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_pipeline.cpython-312.pyc,,
+torchaudio/prototype/pipelines/_vggish/_vggish_impl.py,sha256=EvpbNx4oHk7Zr7LQo-c191XRDLuklj3dZLi2Lgrpe_0,8497
+torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py,sha256=M47A-KFargc9PxZUeu0zqaZsIapAV4XhJqj_P2IHs9g,2713
+torchaudio/prototype/pipelines/hifigan_pipeline.py,sha256=mEmTdVBn50Zyyy0hcjp-fGv9oIbvJFBYBjShXgjFvhU,9654
+torchaudio/prototype/pipelines/rnnt_pipeline.py,sha256=3wwlCg3refzmsAA69XulFRy4GIixA2EH5ZqFbNFKXhk,2184
+torchaudio/prototype/transforms/__init__.py,sha256=1PcKxc8AWlSRkmchT6IDYDseO7itWY7ueZrlcAUZlkY,225
+torchaudio/prototype/transforms/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/prototype/transforms/__pycache__/_transforms.cpython-312.pyc,,
+torchaudio/prototype/transforms/_transforms.py,sha256=K2MjxUrTxlk0nO0NcJA07bDCoFolsQh-TabE84eKC6k,19144
+torchaudio/sox_effects/__init__.py,sha256=gCxdiwHK3ldlGCeYc9VatJW5HyzjWIgw_Sz_krp_rOw,262
+torchaudio/sox_effects/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/sox_effects/__pycache__/sox_effects.cpython-312.pyc,,
+torchaudio/sox_effects/sox_effects.py,sha256=7cHpPFRJ_pZuohHMnX9JIhiVmIJGYntSmgT6QH5GNMA,10981
+torchaudio/transforms/__init__.py,sha256=Tp1o4haiJAV3MRJenmvGXFbmt-RE4qM_pd6U3Ghohqw,1270
+torchaudio/transforms/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/transforms/__pycache__/_multi_channel.cpython-312.pyc,,
+torchaudio/transforms/__pycache__/_transforms.cpython-312.pyc,,
+torchaudio/transforms/_multi_channel.py,sha256=GZ2rrwFt2KtSG7At7kS9Bqh1KmYYw0HwcUnEjc-AWr8,22221
+torchaudio/transforms/_transforms.py,sha256=QHrEsxxxm1bPd5dltPeTcNOsMBu0Ecxa2oe6GIX-nvk,86872
+torchaudio/utils/__init__.py,sha256=NCtfdIUxDi1u0zaamscSbiWzbxn2TOI-MHHWOKU0RnQ,174
+torchaudio/utils/__pycache__/__init__.cpython-312.pyc,,
+torchaudio/utils/__pycache__/download.cpython-312.pyc,,
+torchaudio/utils/__pycache__/ffmpeg_utils.cpython-312.pyc,,
+torchaudio/utils/__pycache__/sox_utils.cpython-312.pyc,,
+torchaudio/utils/download.py,sha256=2IFKD1rsWBFE31HTiyUgpE5y7AJh8_AUPdc-btNQuKw,2882
+torchaudio/utils/ffmpeg_utils.py,sha256=3I6YM95eNyOAg2K-ebEy9kjBzEDq3_OBqggXztPIDcU,319
+torchaudio/utils/sox_utils.py,sha256=WGSj_RfELpol8U2XPABGsAjO7yrPHS3_MgHkx7oHYQU,2421
+torchaudio/version.py,sha256=SF5cvuzciY5C8y2uV8NtCi5xDUUft5kIDdTyXGMI43k,85
+torio/__init__.py,sha256=aX9s0XAHxHhEXE1akQt74BZ0cMUDgBPhaYHQH1lCbXQ,111
+torio/__pycache__/__init__.cpython-312.pyc,,
+torio/_extension/__init__.py,sha256=q5jjeOhSrzqn0WTEwrx61Fr13aCjb7IQCDGsBqAdGEU,313
+torio/_extension/__pycache__/__init__.cpython-312.pyc,,
+torio/_extension/__pycache__/utils.cpython-312.pyc,,
+torio/_extension/utils.py,sha256=ktE0L_z-RF1qkpLVGgdG4DEGHa2Zn6uokOAmwC7Evvo,4904
+torio/io/__init__.py,sha256=xz7REkkyfRhAASzVCAfoNruFtAGIx1I--usPAa2tMww,226
+torio/io/__pycache__/__init__.cpython-312.pyc,,
+torio/io/__pycache__/_streaming_media_decoder.cpython-312.pyc,,
+torio/io/__pycache__/_streaming_media_encoder.cpython-312.pyc,,
+torio/io/_streaming_media_decoder.py,sha256=vSylEWAB_JXOW-0E1n0zDM3Q3Vf1jc1-CNpdUSs13XU,34376
+torio/io/_streaming_media_encoder.py,sha256=rSTYmHdi7RPJ6YPgAyGJhbQvn4mcxLem3nlnr_ophTs,19722
+torio/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torio/lib/__pycache__/__init__.cpython-312.pyc,,
+torio/lib/_torio_ffmpeg4.so,sha256=rRJ2U173HGLuNGOZJRAoBAelJvBQjDWH4atGrpAiIDY,492272
+torio/lib/_torio_ffmpeg5.so,sha256=emgN3KeIjsahgCsK500WpqXSBq71JsUfi3yTAwRsqg8,492272
+torio/lib/_torio_ffmpeg6.so,sha256=zT_Xss-ax9YV_wVEy2BFjmhkvmCEPVDjXQlowNETa6k,492272
+torio/lib/libtorio_ffmpeg4.so,sha256=vs3qairUrO6Dwj3F2bY44OzrpDji2-GJ0LVyriPxZWs,652080
+torio/lib/libtorio_ffmpeg5.so,sha256=-1y_zSQZ7A7rXotxfzotukRSIqa0uQONWrP9DZGQNeY,652080
+torio/lib/libtorio_ffmpeg6.so,sha256=6gbh7nxjlO7wL9m9aqFN82P7zzx8l83nECfRaG6G-Mc,652080
+torio/utils/__init__.py,sha256=ScHtnontymRDNn9qEIC0neue5mfG82yhB8bwETOb0Z4,56
+torio/utils/__pycache__/__init__.cpython-312.pyc,,
+torio/utils/__pycache__/ffmpeg_utils.cpython-312.pyc,,
+torio/utils/ffmpeg_utils.py,sha256=JsP2ptjQAE4U7Z_CSauQKH_k72wdu6nrBMfNHl9pIXQ,8026
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/REQUESTED b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/REQUESTED
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/WHEEL b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/WHEEL
new file mode 100644
index 0000000000000000000000000000000000000000..5b87859d16e8f3962c98089aebacbb6a0a0c0759
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/WHEEL
@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: false
+Tag: cp312-cp312-linux_x86_64
+
diff --git a/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/top_level.txt b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/top_level.txt
new file mode 100644
index 0000000000000000000000000000000000000000..9f448fc64e7113394edf208556101c579616cc18
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio-2.6.0.dist-info/top_level.txt
@@ -0,0 +1,2 @@
+torchaudio
+torio
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ce68ef19753f2abad949c294fcf847e454a848c
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/__main__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1206916cbabcfe5454bee2b448f94f0744b794b5
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/__main__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_dist_ver.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_dist_ver.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f1ca0fb233ceb52be8f9aa714f8f43ac7da2ed0c
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_dist_ver.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_main.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_main.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..66cffae42c07a362b335d8e26e554b9bfaf03000
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_main.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_monitor.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_monitor.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..585a79d1617d5ac5cbee98d250041d68a4b153c7
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_monitor.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cbee433318d8d9facc2f3193be144b3881dcc112
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_gui.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_gui.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c36f1c8c5aa4a4b7d67f7909c7a2263f7aea488c
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_gui.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_notebook.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_notebook.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..58f5a8296187680018ffef5745b10d48e79a4761
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_notebook.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_pandas.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_pandas.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..30e87c2e23fa18f3c1f4d4b52c11663dbffce161
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_tqdm_pandas.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fe8058e37a8d429c2d125974806de00c4aa54f01
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/asyncio.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/asyncio.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..58e1754d56876200182b0fed986a2eee0be7a2d3
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/asyncio.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/auto.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/auto.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ab0abc2936a39c71046463c95a7a6f6dbec32b7
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/auto.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/autonotebook.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/autonotebook.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1ab8d14ba910e19a9884e61b65f63c28019946d8
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/autonotebook.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/cli.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/cli.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1b1b5df49c732ddb24b541d7774c7f9e28cf400c
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/cli.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/dask.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/dask.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..887b122f0dfb1783b9c4ce02bf9f516a04de3b33
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/dask.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/gui.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/gui.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0b02b4c26d31f7f8784a7d615720b3599b1dd596
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/gui.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/keras.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/keras.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e9b04a625755aa6ffa04b1a9debbbf0b79bb2156
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/keras.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/notebook.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/notebook.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9f2587d999b879c5f0e688f3d6479f8c15d62bcf
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/notebook.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/rich.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/rich.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..43f735617cf71a010145afb636e09a745061ccbd
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/rich.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/std.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/std.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..14b6cfcb6f12a4307cc52cf1f91baff85b228123
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/std.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/tk.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/tk.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9dbeb2a7150d1698ba4c07dd67fc633d6106cd82
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/tk.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b66cc847a2451aac14ef71fa8d6d2cfa7aee2234
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/__pycache__/version.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bbb41613eb989cc9cacbba920066fa89a70430af
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/__pycache__/version.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/_dist_ver.py b/lib/python3.12/site-packages/tqdm/_dist_ver.py
new file mode 100644
index 0000000000000000000000000000000000000000..61af7d5bb0b25d8dc934b45b18ea35bd32dbb465
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/_dist_ver.py
@@ -0,0 +1 @@
+__version__ = '4.67.1'
diff --git a/lib/python3.12/site-packages/tqdm/completion.sh b/lib/python3.12/site-packages/tqdm/completion.sh
new file mode 100644
index 0000000000000000000000000000000000000000..9f61c7f14bb8c1f6099b9eb75dce28ece6a7ae96
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/completion.sh
@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+_tqdm(){
+  local cur prv
+  cur="${COMP_WORDS[COMP_CWORD]}"
+  prv="${COMP_WORDS[COMP_CWORD - 1]}"
+
+  case ${prv} in
+  --bar_format|--buf_size|--colour|--comppath|--delay|--delim|--desc|--initial|--lock_args|--manpath|--maxinterval|--mininterval|--miniters|--ncols|--nrows|--position|--postfix|--smoothing|--total|--unit|--unit_divisor)
+    # await user input
+    ;;
+  "--log")
+    COMPREPLY=($(compgen -W       'CRITICAL FATAL ERROR WARN WARNING INFO DEBUG NOTSET' -- ${cur}))
+    ;;
+  *)
+    COMPREPLY=($(compgen -W '--ascii --bar_format --buf_size --bytes --colour --comppath --delay --delim --desc --disable --dynamic_ncols --help --initial --leave --lock_args --log --manpath --maxinterval --mininterval --miniters --ncols --nrows --null --position --postfix --smoothing --tee --total --unit --unit_divisor --unit_scale --update --update_to --version --write_bytes -h -v' -- ${cur}))
+    ;;
+  esac
+}
+complete -F _tqdm tqdm
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__init__.py b/lib/python3.12/site-packages/tqdm/contrib/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d059461f91fb79115263c16314c3487e16ab98c2
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/__init__.py
@@ -0,0 +1,92 @@
+"""
+Thin wrappers around common functions.
+
+Subpackages contain potentially unstable extensions.
+"""
+from warnings import warn
+
+from ..auto import tqdm as tqdm_auto
+from ..std import TqdmDeprecationWarning, tqdm
+from ..utils import ObjectWrapper
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['tenumerate', 'tzip', 'tmap']
+
+
+class DummyTqdmFile(ObjectWrapper):
+    """Dummy file-like that will write to tqdm"""
+
+    def __init__(self, wrapped):
+        super().__init__(wrapped)
+        self._buf = []
+
+    def write(self, x, nolock=False):
+        nl = b"\n" if isinstance(x, bytes) else "\n"
+        pre, sep, post = x.rpartition(nl)
+        if sep:
+            blank = type(nl)()
+            tqdm.write(blank.join(self._buf + [pre, sep]),
+                       end=blank, file=self._wrapped, nolock=nolock)
+            self._buf = [post]
+        else:
+            self._buf.append(x)
+
+    def __del__(self):
+        if self._buf:
+            blank = type(self._buf[0])()
+            try:
+                tqdm.write(blank.join(self._buf), end=blank, file=self._wrapped)
+            except (OSError, ValueError):
+                pass
+
+
+def builtin_iterable(func):
+    """Returns `func`"""
+    warn("This function has no effect, and will be removed in tqdm==5.0.0",
+         TqdmDeprecationWarning, stacklevel=2)
+    return func
+
+
+def tenumerate(iterable, start=0, total=None, tqdm_class=tqdm_auto, **tqdm_kwargs):
+    """
+    Equivalent of `numpy.ndenumerate` or builtin `enumerate`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    try:
+        import numpy as np
+    except ImportError:
+        pass
+    else:
+        if isinstance(iterable, np.ndarray):
+            return tqdm_class(np.ndenumerate(iterable), total=total or iterable.size,
+                              **tqdm_kwargs)
+    return enumerate(tqdm_class(iterable, total=total, **tqdm_kwargs), start)
+
+
+def tzip(iter1, *iter2plus, **tqdm_kwargs):
+    """
+    Equivalent of builtin `zip`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    kwargs = tqdm_kwargs.copy()
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    for i in zip(tqdm_class(iter1, **kwargs), *iter2plus):
+        yield i
+
+
+def tmap(function, *sequences, **tqdm_kwargs):
+    """
+    Equivalent of builtin `map`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    for i in tzip(*sequences, **tqdm_kwargs):
+        yield function(*i)
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a72341abee924869624b359da3a0a76cb3f21d4f
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/bells.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/bells.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ef0156ddcdf547259d57198dbd063f2d7129aca
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/bells.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/concurrent.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/concurrent.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2a59d7095ae9c109131ed1a092156932163c0005
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/concurrent.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/discord.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/discord.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5f90a40dfee9ec6e82b804e57f8a10c1539b42f9
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/discord.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/itertools.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/itertools.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..508e757932049f036915342b7d0060cc705e6268
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/itertools.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/logging.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..673ca39d8e521f934032b6eafd3cfd93c96b003f
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/logging.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/slack.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/slack.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..322ef03234de0c53987ff093e5877c6aa30eaada
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/slack.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/telegram.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/telegram.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5f32fc4c9cc317424dd84aadc1d8220941aac8d8
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/telegram.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/__pycache__/utils_worker.cpython-312.pyc b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/utils_worker.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b69dd8a771106e8860fb7edacb09700230b70842
Binary files /dev/null and b/lib/python3.12/site-packages/tqdm/contrib/__pycache__/utils_worker.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/tqdm/contrib/bells.py b/lib/python3.12/site-packages/tqdm/contrib/bells.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b8f4b9ecd894f1edfaa08d9fe730b8d7c8b93e0
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/bells.py
@@ -0,0 +1,26 @@
+"""
+Even more features than `tqdm.auto` (all the bells & whistles):
+
+- `tqdm.auto`
+- `tqdm.tqdm.pandas`
+- `tqdm.contrib.telegram`
+    + uses `${TQDM_TELEGRAM_TOKEN}` and `${TQDM_TELEGRAM_CHAT_ID}`
+- `tqdm.contrib.discord`
+    + uses `${TQDM_DISCORD_TOKEN}` and `${TQDM_DISCORD_CHANNEL_ID}`
+"""
+__all__ = ['tqdm', 'trange']
+import warnings
+from os import getenv
+
+if getenv("TQDM_SLACK_TOKEN") and getenv("TQDM_SLACK_CHANNEL"):
+    from .slack import tqdm, trange
+elif getenv("TQDM_TELEGRAM_TOKEN") and getenv("TQDM_TELEGRAM_CHAT_ID"):
+    from .telegram import tqdm, trange
+elif getenv("TQDM_DISCORD_TOKEN") and getenv("TQDM_DISCORD_CHANNEL_ID"):
+    from .discord import tqdm, trange
+else:
+    from ..auto import tqdm, trange
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", category=FutureWarning)
+    tqdm.pandas()
diff --git a/lib/python3.12/site-packages/tqdm/contrib/concurrent.py b/lib/python3.12/site-packages/tqdm/contrib/concurrent.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd81d622a1309df179042159a56cef4f8c309224
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/concurrent.py
@@ -0,0 +1,105 @@
+"""
+Thin wrappers around `concurrent.futures`.
+"""
+from contextlib import contextmanager
+from operator import length_hint
+from os import cpu_count
+
+from ..auto import tqdm as tqdm_auto
+from ..std import TqdmWarning
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['thread_map', 'process_map']
+
+
+@contextmanager
+def ensure_lock(tqdm_class, lock_name=""):
+    """get (create if necessary) and then restore `tqdm_class`'s lock"""
+    old_lock = getattr(tqdm_class, '_lock', None)  # don't create a new lock
+    lock = old_lock or tqdm_class.get_lock()  # maybe create a new lock
+    lock = getattr(lock, lock_name, lock)  # maybe subtype
+    tqdm_class.set_lock(lock)
+    yield lock
+    if old_lock is None:
+        del tqdm_class._lock
+    else:
+        tqdm_class.set_lock(old_lock)
+
+
+def _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs):
+    """
+    Implementation of `thread_map` and `process_map`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    max_workers  : [default: min(32, cpu_count() + 4)].
+    chunksize  : [default: 1].
+    lock_name  : [default: "":str].
+    """
+    kwargs = tqdm_kwargs.copy()
+    if "total" not in kwargs:
+        kwargs["total"] = length_hint(iterables[0])
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4))
+    chunksize = kwargs.pop("chunksize", 1)
+    lock_name = kwargs.pop("lock_name", "")
+    with ensure_lock(tqdm_class, lock_name=lock_name) as lk:
+        # share lock in case workers are already using `tqdm`
+        with PoolExecutor(max_workers=max_workers, initializer=tqdm_class.set_lock,
+                          initargs=(lk,)) as ex:
+            return list(tqdm_class(ex.map(fn, *iterables, chunksize=chunksize), **kwargs))
+
+
+def thread_map(fn, *iterables, **tqdm_kwargs):
+    """
+    Equivalent of `list(map(fn, *iterables))`
+    driven by `concurrent.futures.ThreadPoolExecutor`.
+
+    Parameters
+    ----------
+    tqdm_class  : optional
+        `tqdm` class to use for bars [default: tqdm.auto.tqdm].
+    max_workers  : int, optional
+        Maximum number of workers to spawn; passed to
+        `concurrent.futures.ThreadPoolExecutor.__init__`.
+        [default: max(32, cpu_count() + 4)].
+    """
+    from concurrent.futures import ThreadPoolExecutor
+    return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs)
+
+
+def process_map(fn, *iterables, **tqdm_kwargs):
+    """
+    Equivalent of `list(map(fn, *iterables))`
+    driven by `concurrent.futures.ProcessPoolExecutor`.
+
+    Parameters
+    ----------
+    tqdm_class  : optional
+        `tqdm` class to use for bars [default: tqdm.auto.tqdm].
+    max_workers  : int, optional
+        Maximum number of workers to spawn; passed to
+        `concurrent.futures.ProcessPoolExecutor.__init__`.
+        [default: min(32, cpu_count() + 4)].
+    chunksize  : int, optional
+        Size of chunks sent to worker processes; passed to
+        `concurrent.futures.ProcessPoolExecutor.map`. [default: 1].
+    lock_name  : str, optional
+        Member of `tqdm_class.get_lock()` to use [default: mp_lock].
+    """
+    from concurrent.futures import ProcessPoolExecutor
+    if iterables and "chunksize" not in tqdm_kwargs:
+        # default `chunksize=1` has poor performance for large iterables
+        # (most time spent dispatching items to workers).
+        longest_iterable_len = max(map(length_hint, iterables))
+        if longest_iterable_len > 1000:
+            from warnings import warn
+            warn("Iterable length %d > 1000 but `chunksize` is not set."
+                 " This may seriously degrade multiprocess performance."
+                 " Set `chunksize=1` or more." % longest_iterable_len,
+                 TqdmWarning, stacklevel=2)
+    if "lock_name" not in tqdm_kwargs:
+        tqdm_kwargs = tqdm_kwargs.copy()
+        tqdm_kwargs["lock_name"] = "mp_lock"
+    return _executor_map(ProcessPoolExecutor, fn, *iterables, **tqdm_kwargs)
diff --git a/lib/python3.12/site-packages/tqdm/contrib/discord.py b/lib/python3.12/site-packages/tqdm/contrib/discord.py
new file mode 100644
index 0000000000000000000000000000000000000000..574baa84bbbeb5afce4a49f23edac894d680ca82
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/discord.py
@@ -0,0 +1,156 @@
+"""
+Sends updates to a Discord bot.
+
+Usage:
+>>> from tqdm.contrib.discord import tqdm, trange
+>>> for i in trange(10, token='{token}', channel_id='{channel_id}'):
+...     ...
+
+![screenshot](https://tqdm.github.io/img/screenshot-discord.png)
+"""
+from os import getenv
+from warnings import warn
+
+from requests import Session
+from requests.utils import default_user_agent
+
+from ..auto import tqdm as tqdm_auto
+from ..std import TqdmWarning
+from ..version import __version__
+from .utils_worker import MonoWorker
+
+__author__ = {"github.com/": ["casperdcl", "guigoruiz1"]}
+__all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange']
+
+
+class DiscordIO(MonoWorker):
+    """Non-blocking file-like IO using a Discord Bot."""
+    API = "https://discord.com/api/v10"
+    UA = f"tqdm (https://tqdm.github.io, {__version__}) {default_user_agent()}"
+
+    def __init__(self, token, channel_id):
+        """Creates a new message in the given `channel_id`."""
+        super().__init__()
+        self.token = token
+        self.channel_id = channel_id
+        self.session = Session()
+        self.text = self.__class__.__name__
+        self.message_id
+
+    @property
+    def message_id(self):
+        if hasattr(self, '_message_id'):
+            return self._message_id
+        try:
+            res = self.session.post(
+                f'{self.API}/channels/{self.channel_id}/messages',
+                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},
+                json={'content': f"`{self.text}`"}).json()
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            if res.get('error_code') == 429:
+                warn("Creation rate limit: try increasing `mininterval`.",
+                     TqdmWarning, stacklevel=2)
+            else:
+                self._message_id = res['id']
+                return self._message_id
+
+    def write(self, s):
+        """Replaces internal `message_id`'s text with `s`."""
+        if not s:
+            s = "..."
+        s = s.replace('\r', '').strip()
+        if s == self.text:
+            return  # avoid duplicate message Bot error
+        message_id = self.message_id
+        if message_id is None:
+            return
+        self.text = s
+        try:
+            future = self.submit(
+                self.session.patch,
+                f'{self.API}/channels/{self.channel_id}/messages/{message_id}',
+                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},
+                json={'content': f"`{self.text}`"})
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            return future
+
+    def delete(self):
+        """Deletes internal `message_id`."""
+        try:
+            future = self.submit(
+                self.session.delete,
+                f'{self.API}/channels/{self.channel_id}/messages/{self.message_id}',
+                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA})
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            return future
+
+
+class tqdm_discord(tqdm_auto):
+    """
+    Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot.
+    May take a few seconds to create (`__init__`).
+
+    - create a discord bot (not public, no requirement of OAuth2 code
+      grant, only send message permissions) & invite it to a channel:
+      
+    - copy the bot `{token}` & `{channel_id}` and paste below
+
+    >>> from tqdm.contrib.discord import tqdm, trange
+    >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'):
+    ...     ...
+    """
+    def __init__(self, *args, **kwargs):
+        """
+        Parameters
+        ----------
+        token  : str, required. Discord bot token
+            [default: ${TQDM_DISCORD_TOKEN}].
+        channel_id  : int, required. Discord channel ID
+            [default: ${TQDM_DISCORD_CHANNEL_ID}].
+
+        See `tqdm.auto.tqdm.__init__` for other parameters.
+        """
+        if not kwargs.get('disable'):
+            kwargs = kwargs.copy()
+            self.dio = DiscordIO(
+                kwargs.pop('token', getenv('TQDM_DISCORD_TOKEN')),
+                kwargs.pop('channel_id', getenv('TQDM_DISCORD_CHANNEL_ID')))
+        super().__init__(*args, **kwargs)
+
+    def display(self, **kwargs):
+        super().display(**kwargs)
+        fmt = self.format_dict
+        if fmt.get('bar_format', None):
+            fmt['bar_format'] = fmt['bar_format'].replace(
+                '', '{bar:10u}').replace('{bar}', '{bar:10u}')
+        else:
+            fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'
+        self.dio.write(self.format_meter(**fmt))
+
+    def clear(self, *args, **kwargs):
+        super().clear(*args, **kwargs)
+        if not self.disable:
+            self.dio.write("")
+
+    def close(self):
+        if self.disable:
+            return
+        super().close()
+        if not (self.leave or (self.leave is None and self.pos == 0)):
+            self.dio.delete()
+
+
+def tdrange(*args, **kwargs):
+    """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`."""
+    return tqdm_discord(range(*args), **kwargs)
+
+
+# Aliases
+tqdm = tqdm_discord
+trange = tdrange
diff --git a/lib/python3.12/site-packages/tqdm/contrib/itertools.py b/lib/python3.12/site-packages/tqdm/contrib/itertools.py
new file mode 100644
index 0000000000000000000000000000000000000000..e67651a41a6b8760d9b928ea48239e4611d70315
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/itertools.py
@@ -0,0 +1,35 @@
+"""
+Thin wrappers around `itertools`.
+"""
+import itertools
+
+from ..auto import tqdm as tqdm_auto
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['product']
+
+
+def product(*iterables, **tqdm_kwargs):
+    """
+    Equivalent of `itertools.product`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    kwargs = tqdm_kwargs.copy()
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    try:
+        lens = list(map(len, iterables))
+    except TypeError:
+        total = None
+    else:
+        total = 1
+        for i in lens:
+            total *= i
+        kwargs.setdefault("total", total)
+    with tqdm_class(**kwargs) as t:
+        it = itertools.product(*iterables)
+        for i in it:
+            yield i
+            t.update()
diff --git a/lib/python3.12/site-packages/tqdm/contrib/logging.py b/lib/python3.12/site-packages/tqdm/contrib/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..e06febe37b5d70b5296804c55dca48a397c250e3
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/logging.py
@@ -0,0 +1,126 @@
+"""
+Helper functionality for interoperability with stdlib `logging`.
+"""
+import logging
+import sys
+from contextlib import contextmanager
+
+try:
+    from typing import Iterator, List, Optional, Type  # noqa: F401
+except ImportError:
+    pass
+
+from ..std import tqdm as std_tqdm
+
+
+class _TqdmLoggingHandler(logging.StreamHandler):
+    def __init__(
+        self,
+        tqdm_class=std_tqdm  # type: Type[std_tqdm]
+    ):
+        super().__init__()
+        self.tqdm_class = tqdm_class
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            self.tqdm_class.write(msg, file=self.stream)
+            self.flush()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except:  # noqa pylint: disable=bare-except
+            self.handleError(record)
+
+
+def _is_console_logging_handler(handler):
+    return (isinstance(handler, logging.StreamHandler)
+            and handler.stream in {sys.stdout, sys.stderr})
+
+
+def _get_first_found_console_logging_handler(handlers):
+    for handler in handlers:
+        if _is_console_logging_handler(handler):
+            return handler
+
+
+@contextmanager
+def logging_redirect_tqdm(
+    loggers=None,  # type: Optional[List[logging.Logger]],
+    tqdm_class=std_tqdm  # type: Type[std_tqdm]
+):
+    # type: (...) -> Iterator[None]
+    """
+    Context manager redirecting console logging to `tqdm.write()`, leaving
+    other logging handlers (e.g. log files) unaffected.
+
+    Parameters
+    ----------
+    loggers  : list, optional
+      Which handlers to redirect (default: [logging.root]).
+    tqdm_class  : optional
+
+    Example
+    -------
+    ```python
+    import logging
+    from tqdm import trange
+    from tqdm.contrib.logging import logging_redirect_tqdm
+
+    LOG = logging.getLogger(__name__)
+
+    if __name__ == '__main__':
+        logging.basicConfig(level=logging.INFO)
+        with logging_redirect_tqdm():
+            for i in trange(9):
+                if i == 4:
+                    LOG.info("console logging redirected to `tqdm.write()`")
+        # logging restored
+    ```
+    """
+    if loggers is None:
+        loggers = [logging.root]
+    original_handlers_list = [logger.handlers for logger in loggers]
+    try:
+        for logger in loggers:
+            tqdm_handler = _TqdmLoggingHandler(tqdm_class)
+            orig_handler = _get_first_found_console_logging_handler(logger.handlers)
+            if orig_handler is not None:
+                tqdm_handler.setFormatter(orig_handler.formatter)
+                tqdm_handler.stream = orig_handler.stream
+            logger.handlers = [
+                handler for handler in logger.handlers
+                if not _is_console_logging_handler(handler)] + [tqdm_handler]
+        yield
+    finally:
+        for logger, original_handlers in zip(loggers, original_handlers_list):
+            logger.handlers = original_handlers
+
+
+@contextmanager
+def tqdm_logging_redirect(
+    *args,
+    # loggers=None,  # type: Optional[List[logging.Logger]]
+    # tqdm=None,  # type: Optional[Type[tqdm.tqdm]]
+    **kwargs
+):
+    # type: (...) -> Iterator[None]
+    """
+    Convenience shortcut for:
+    ```python
+    with tqdm_class(*args, **tqdm_kwargs) as pbar:
+        with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class):
+            yield pbar
+    ```
+
+    Parameters
+    ----------
+    tqdm_class  : optional, (default: tqdm.std.tqdm).
+    loggers  : optional, list.
+    **tqdm_kwargs  : passed to `tqdm_class`.
+    """
+    tqdm_kwargs = kwargs.copy()
+    loggers = tqdm_kwargs.pop('loggers', None)
+    tqdm_class = tqdm_kwargs.pop('tqdm_class', std_tqdm)
+    with tqdm_class(*args, **tqdm_kwargs) as pbar:
+        with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class):
+            yield pbar
diff --git a/lib/python3.12/site-packages/tqdm/contrib/slack.py b/lib/python3.12/site-packages/tqdm/contrib/slack.py
new file mode 100644
index 0000000000000000000000000000000000000000..9bca8ee98904ce869a4f8d6417bbdc4f00b38751
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/slack.py
@@ -0,0 +1,120 @@
+"""
+Sends updates to a Slack app.
+
+Usage:
+>>> from tqdm.contrib.slack import tqdm, trange
+>>> for i in trange(10, token='{token}', channel='{channel}'):
+...     ...
+
+![screenshot](https://tqdm.github.io/img/screenshot-slack.png)
+"""
+import logging
+from os import getenv
+
+try:
+    from slack_sdk import WebClient
+except ImportError:
+    raise ImportError("Please `pip install slack-sdk`")
+
+from ..auto import tqdm as tqdm_auto
+from .utils_worker import MonoWorker
+
+__author__ = {"github.com/": ["0x2b3bfa0", "casperdcl"]}
+__all__ = ['SlackIO', 'tqdm_slack', 'tsrange', 'tqdm', 'trange']
+
+
+class SlackIO(MonoWorker):
+    """Non-blocking file-like IO using a Slack app."""
+    def __init__(self, token, channel):
+        """Creates a new message in the given `channel`."""
+        super().__init__()
+        self.client = WebClient(token=token)
+        self.text = self.__class__.__name__
+        try:
+            self.message = self.client.chat_postMessage(channel=channel, text=self.text)
+        except Exception as e:
+            tqdm_auto.write(str(e))
+            self.message = None
+
+    def write(self, s):
+        """Replaces internal `message`'s text with `s`."""
+        if not s:
+            s = "..."
+        s = s.replace('\r', '').strip()
+        if s == self.text:
+            return  # skip duplicate message
+        message = self.message
+        if message is None:
+            return
+        self.text = s
+        try:
+            future = self.submit(self.client.chat_update, channel=message['channel'],
+                                 ts=message['ts'], text='`' + s + '`')
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            return future
+
+
+class tqdm_slack(tqdm_auto):
+    """
+    Standard `tqdm.auto.tqdm` but also sends updates to a Slack app.
+    May take a few seconds to create (`__init__`).
+
+    - create a Slack app with the `chat:write` scope & invite it to a
+      channel: 
+    - copy the bot `{token}` & `{channel}` and paste below
+    >>> from tqdm.contrib.slack import tqdm, trange
+    >>> for i in tqdm(iterable, token='{token}', channel='{channel}'):
+    ...     ...
+    """
+    def __init__(self, *args, **kwargs):
+        """
+        Parameters
+        ----------
+        token  : str, required. Slack token
+            [default: ${TQDM_SLACK_TOKEN}].
+        channel  : int, required. Slack channel
+            [default: ${TQDM_SLACK_CHANNEL}].
+        mininterval  : float, optional.
+          Minimum of [default: 1.5] to avoid rate limit.
+
+        See `tqdm.auto.tqdm.__init__` for other parameters.
+        """
+        if not kwargs.get('disable'):
+            kwargs = kwargs.copy()
+            logging.getLogger("HTTPClient").setLevel(logging.WARNING)
+            self.sio = SlackIO(
+                kwargs.pop('token', getenv("TQDM_SLACK_TOKEN")),
+                kwargs.pop('channel', getenv("TQDM_SLACK_CHANNEL")))
+            kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5))
+        super().__init__(*args, **kwargs)
+
+    def display(self, **kwargs):
+        super().display(**kwargs)
+        fmt = self.format_dict
+        if fmt.get('bar_format', None):
+            fmt['bar_format'] = fmt['bar_format'].replace(
+                '', '`{bar:10}`').replace('{bar}', '`{bar:10u}`')
+        else:
+            fmt['bar_format'] = '{l_bar}`{bar:10}`{r_bar}'
+        if fmt['ascii'] is False:
+            fmt['ascii'] = [":black_square:", ":small_blue_diamond:", ":large_blue_diamond:",
+                            ":large_blue_square:"]
+            fmt['ncols'] = 336
+        self.sio.write(self.format_meter(**fmt))
+
+    def clear(self, *args, **kwargs):
+        super().clear(*args, **kwargs)
+        if not self.disable:
+            self.sio.write("")
+
+
+def tsrange(*args, **kwargs):
+    """Shortcut for `tqdm.contrib.slack.tqdm(range(*args), **kwargs)`."""
+    return tqdm_slack(range(*args), **kwargs)
+
+
+# Aliases
+tqdm = tqdm_slack
+trange = tsrange
diff --git a/lib/python3.12/site-packages/tqdm/contrib/telegram.py b/lib/python3.12/site-packages/tqdm/contrib/telegram.py
new file mode 100644
index 0000000000000000000000000000000000000000..019151800bc0c4c4fc543314b6398aa602b0692a
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/telegram.py
@@ -0,0 +1,153 @@
+"""
+Sends updates to a Telegram bot.
+
+Usage:
+>>> from tqdm.contrib.telegram import tqdm, trange
+>>> for i in trange(10, token='{token}', chat_id='{chat_id}'):
+...     ...
+
+![screenshot](https://tqdm.github.io/img/screenshot-telegram.gif)
+"""
+from os import getenv
+from warnings import warn
+
+from requests import Session
+
+from ..auto import tqdm as tqdm_auto
+from ..std import TqdmWarning
+from .utils_worker import MonoWorker
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['TelegramIO', 'tqdm_telegram', 'ttgrange', 'tqdm', 'trange']
+
+
+class TelegramIO(MonoWorker):
+    """Non-blocking file-like IO using a Telegram Bot."""
+    API = 'https://api.telegram.org/bot'
+
+    def __init__(self, token, chat_id):
+        """Creates a new message in the given `chat_id`."""
+        super().__init__()
+        self.token = token
+        self.chat_id = chat_id
+        self.session = Session()
+        self.text = self.__class__.__name__
+        self.message_id
+
+    @property
+    def message_id(self):
+        if hasattr(self, '_message_id'):
+            return self._message_id
+        try:
+            res = self.session.post(
+                self.API + '%s/sendMessage' % self.token,
+                data={'text': '`' + self.text + '`', 'chat_id': self.chat_id,
+                      'parse_mode': 'MarkdownV2'}).json()
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            if res.get('error_code') == 429:
+                warn("Creation rate limit: try increasing `mininterval`.",
+                     TqdmWarning, stacklevel=2)
+            else:
+                self._message_id = res['result']['message_id']
+                return self._message_id
+
+    def write(self, s):
+        """Replaces internal `message_id`'s text with `s`."""
+        if not s:
+            s = "..."
+        s = s.replace('\r', '').strip()
+        if s == self.text:
+            return  # avoid duplicate message Bot error
+        message_id = self.message_id
+        if message_id is None:
+            return
+        self.text = s
+        try:
+            future = self.submit(
+                self.session.post, self.API + '%s/editMessageText' % self.token,
+                data={'text': '`' + s + '`', 'chat_id': self.chat_id,
+                      'message_id': message_id, 'parse_mode': 'MarkdownV2'})
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            return future
+
+    def delete(self):
+        """Deletes internal `message_id`."""
+        try:
+            future = self.submit(
+                self.session.post, self.API + '%s/deleteMessage' % self.token,
+                data={'chat_id': self.chat_id, 'message_id': self.message_id})
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            return future
+
+
+class tqdm_telegram(tqdm_auto):
+    """
+    Standard `tqdm.auto.tqdm` but also sends updates to a Telegram Bot.
+    May take a few seconds to create (`__init__`).
+
+    - create a bot 
+    - copy its `{token}`
+    - add the bot to a chat and send it a message such as `/start`
+    - go to  to find out
+      the `{chat_id}`
+    - paste the `{token}` & `{chat_id}` below
+
+    >>> from tqdm.contrib.telegram import tqdm, trange
+    >>> for i in tqdm(iterable, token='{token}', chat_id='{chat_id}'):
+    ...     ...
+    """
+    def __init__(self, *args, **kwargs):
+        """
+        Parameters
+        ----------
+        token  : str, required. Telegram token
+            [default: ${TQDM_TELEGRAM_TOKEN}].
+        chat_id  : str, required. Telegram chat ID
+            [default: ${TQDM_TELEGRAM_CHAT_ID}].
+
+        See `tqdm.auto.tqdm.__init__` for other parameters.
+        """
+        if not kwargs.get('disable'):
+            kwargs = kwargs.copy()
+            self.tgio = TelegramIO(
+                kwargs.pop('token', getenv('TQDM_TELEGRAM_TOKEN')),
+                kwargs.pop('chat_id', getenv('TQDM_TELEGRAM_CHAT_ID')))
+        super().__init__(*args, **kwargs)
+
+    def display(self, **kwargs):
+        super().display(**kwargs)
+        fmt = self.format_dict
+        if fmt.get('bar_format', None):
+            fmt['bar_format'] = fmt['bar_format'].replace(
+                '', '{bar:10u}').replace('{bar}', '{bar:10u}')
+        else:
+            fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'
+        self.tgio.write(self.format_meter(**fmt))
+
+    def clear(self, *args, **kwargs):
+        super().clear(*args, **kwargs)
+        if not self.disable:
+            self.tgio.write("")
+
+    def close(self):
+        if self.disable:
+            return
+        super().close()
+        if not (self.leave or (self.leave is None and self.pos == 0)):
+            self.tgio.delete()
+
+
+def ttgrange(*args, **kwargs):
+    """Shortcut for `tqdm.contrib.telegram.tqdm(range(*args), **kwargs)`."""
+    return tqdm_telegram(range(*args), **kwargs)
+
+
+# Aliases
+tqdm = tqdm_telegram
+trange = ttgrange
diff --git a/lib/python3.12/site-packages/tqdm/contrib/utils_worker.py b/lib/python3.12/site-packages/tqdm/contrib/utils_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a03a2a8930001e37938836196e0d15b649b07a8
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/contrib/utils_worker.py
@@ -0,0 +1,38 @@
+"""
+IO/concurrency helpers for `tqdm.contrib`.
+"""
+from collections import deque
+from concurrent.futures import ThreadPoolExecutor
+
+from ..auto import tqdm as tqdm_auto
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['MonoWorker']
+
+
+class MonoWorker(object):
+    """
+    Supports one running task and one waiting task.
+    The waiting task is the most recent submitted (others are discarded).
+    """
+    def __init__(self):
+        self.pool = ThreadPoolExecutor(max_workers=1)
+        self.futures = deque([], 2)
+
+    def submit(self, func, *args, **kwargs):
+        """`func(*args, **kwargs)` may replace currently waiting task."""
+        futures = self.futures
+        if len(futures) == futures.maxlen:
+            running = futures.popleft()
+            if not running.done():
+                if len(futures):  # clear waiting
+                    waiting = futures.pop()
+                    waiting.cancel()
+                futures.appendleft(running)  # re-insert running
+        try:
+            waiting = self.pool.submit(func, *args, **kwargs)
+        except Exception as e:
+            tqdm_auto.write(str(e))
+        else:
+            futures.append(waiting)
+            return waiting
diff --git a/lib/python3.12/site-packages/tqdm/gui.py b/lib/python3.12/site-packages/tqdm/gui.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb52fb91a8661f4c73edd352bbc6f21b877dcfee
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/gui.py
@@ -0,0 +1,179 @@
+"""
+Matplotlib GUI progressbar decorator for iterators.
+
+Usage:
+>>> from tqdm.gui import trange, tqdm
+>>> for i in trange(10):
+...     ...
+"""
+# future division is important to divide integers and get as
+# a result precise floating numbers (instead of truncated int)
+import re
+from warnings import warn
+
+# to inherit from the tqdm class
+from .std import TqdmExperimentalWarning
+from .std import tqdm as std_tqdm
+
+# import compatibility functions and utilities
+
+__author__ = {"github.com/": ["casperdcl", "lrq3000"]}
+__all__ = ['tqdm_gui', 'tgrange', 'tqdm', 'trange']
+
+
+class tqdm_gui(std_tqdm):  # pragma: no cover
+    """Experimental Matplotlib GUI version of tqdm!"""
+    # TODO: @classmethod: write() on GUI?
+    def __init__(self, *args, **kwargs):
+        from collections import deque
+
+        import matplotlib as mpl
+        import matplotlib.pyplot as plt
+        kwargs = kwargs.copy()
+        kwargs['gui'] = True
+        colour = kwargs.pop('colour', 'g')
+        super().__init__(*args, **kwargs)
+
+        if self.disable:
+            return
+
+        warn("GUI is experimental/alpha", TqdmExperimentalWarning, stacklevel=2)
+        self.mpl = mpl
+        self.plt = plt
+
+        # Remember if external environment uses toolbars
+        self.toolbar = self.mpl.rcParams['toolbar']
+        self.mpl.rcParams['toolbar'] = 'None'
+
+        self.mininterval = max(self.mininterval, 0.5)
+        self.fig, ax = plt.subplots(figsize=(9, 2.2))
+        # self.fig.subplots_adjust(bottom=0.2)
+        total = self.__len__()  # avoids TypeError on None #971
+        if total is not None:
+            self.xdata = []
+            self.ydata = []
+            self.zdata = []
+        else:
+            self.xdata = deque([])
+            self.ydata = deque([])
+            self.zdata = deque([])
+        self.line1, = ax.plot(self.xdata, self.ydata, color='b')
+        self.line2, = ax.plot(self.xdata, self.zdata, color='k')
+        ax.set_ylim(0, 0.001)
+        if total is not None:
+            ax.set_xlim(0, 100)
+            ax.set_xlabel("percent")
+            self.fig.legend((self.line1, self.line2), ("cur", "est"),
+                            loc='center right')
+            # progressbar
+            self.hspan = plt.axhspan(0, 0.001, xmin=0, xmax=0, color=colour)
+        else:
+            # ax.set_xlim(-60, 0)
+            ax.set_xlim(0, 60)
+            ax.invert_xaxis()
+            ax.set_xlabel("seconds")
+            ax.legend(("cur", "est"), loc='lower left')
+        ax.grid()
+        # ax.set_xlabel('seconds')
+        ax.set_ylabel((self.unit if self.unit else "it") + "/s")
+        if self.unit_scale:
+            plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 0))
+            ax.yaxis.get_offset_text().set_x(-0.15)
+
+        # Remember if external environment is interactive
+        self.wasion = plt.isinteractive()
+        plt.ion()
+        self.ax = ax
+
+    def close(self):
+        if self.disable:
+            return
+
+        self.disable = True
+
+        with self.get_lock():
+            self._instances.remove(self)
+
+        # Restore toolbars
+        self.mpl.rcParams['toolbar'] = self.toolbar
+        # Return to non-interactive mode
+        if not self.wasion:
+            self.plt.ioff()
+        if self.leave:
+            self.display()
+        else:
+            self.plt.close(self.fig)
+
+    def clear(self, *_, **__):
+        pass
+
+    def display(self, *_, **__):
+        n = self.n
+        cur_t = self._time()
+        elapsed = cur_t - self.start_t
+        delta_it = n - self.last_print_n
+        delta_t = cur_t - self.last_print_t
+
+        # Inline due to multiple calls
+        total = self.total
+        xdata = self.xdata
+        ydata = self.ydata
+        zdata = self.zdata
+        ax = self.ax
+        line1 = self.line1
+        line2 = self.line2
+        hspan = getattr(self, 'hspan', None)
+        # instantaneous rate
+        y = delta_it / delta_t
+        # overall rate
+        z = n / elapsed
+        # update line data
+        xdata.append(n * 100.0 / total if total else cur_t)
+        ydata.append(y)
+        zdata.append(z)
+
+        # Discard old values
+        # xmin, xmax = ax.get_xlim()
+        # if (not total) and elapsed > xmin * 1.1:
+        if (not total) and elapsed > 66:
+            xdata.popleft()
+            ydata.popleft()
+            zdata.popleft()
+
+        ymin, ymax = ax.get_ylim()
+        if y > ymax or z > ymax:
+            ymax = 1.1 * y
+            ax.set_ylim(ymin, ymax)
+            ax.figure.canvas.draw()
+
+        if total:
+            line1.set_data(xdata, ydata)
+            line2.set_data(xdata, zdata)
+            if hspan:
+                hspan.set_xy((0, ymin))
+                hspan.set_height(ymax - ymin)
+                hspan.set_width(n / total)
+        else:
+            t_ago = [cur_t - i for i in xdata]
+            line1.set_data(t_ago, ydata)
+            line2.set_data(t_ago, zdata)
+
+        d = self.format_dict
+        # remove {bar}
+        d['bar_format'] = (d['bar_format'] or "{l_bar}{r_bar}").replace(
+            "{bar}", "")
+        msg = self.format_meter(**d)
+        if '' in msg:
+            msg = "".join(re.split(r'\|?\|?', msg, maxsplit=1))
+        ax.set_title(msg, fontname="DejaVu Sans Mono", fontsize=11)
+        self.plt.pause(1e-9)
+
+
+def tgrange(*args, **kwargs):
+    """Shortcut for `tqdm.gui.tqdm(range(*args), **kwargs)`."""
+    return tqdm_gui(range(*args), **kwargs)
+
+
+# Aliases
+tqdm = tqdm_gui
+trange = tgrange
diff --git a/lib/python3.12/site-packages/tqdm/keras.py b/lib/python3.12/site-packages/tqdm/keras.py
new file mode 100644
index 0000000000000000000000000000000000000000..cce9467c51a95388aaa502d1da9a42f3ebf0af24
--- /dev/null
+++ b/lib/python3.12/site-packages/tqdm/keras.py
@@ -0,0 +1,122 @@
+from copy import copy
+from functools import partial
+
+from .auto import tqdm as tqdm_auto
+
+try:
+    import keras
+except (ImportError, AttributeError) as e:
+    try:
+        from tensorflow import keras
+    except ImportError:
+        raise e
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['TqdmCallback']
+
+
+class TqdmCallback(keras.callbacks.Callback):
+    """Keras callback for epoch and batch progress."""
+    @staticmethod
+    def bar2callback(bar, pop=None, delta=(lambda logs: 1)):
+        def callback(_, logs=None):
+            n = delta(logs)
+            if logs:
+                if pop:
+                    logs = copy(logs)
+                    [logs.pop(i, 0) for i in pop]
+                bar.set_postfix(logs, refresh=False)
+            bar.update(n)
+
+        return callback
+
+    def __init__(self, epochs=None, data_size=None, batch_size=None, verbose=1,
+                 tqdm_class=tqdm_auto, **tqdm_kwargs):
+        """
+        Parameters
+        ----------
+        epochs  : int, optional
+        data_size  : int, optional
+            Number of training pairs.
+        batch_size  : int, optional
+            Number of training pairs per batch.
+        verbose  : int
+            0: epoch, 1: batch (transient), 2: batch. [default: 1].
+            Will be set to `0` unless both `data_size` and `batch_size`
+            are given.
+        tqdm_class  : optional
+            `tqdm` class to use for bars [default: `tqdm.auto.tqdm`].
+        tqdm_kwargs  : optional
+            Any other arguments used for all bars.
+        """
+        if tqdm_kwargs:
+            tqdm_class = partial(tqdm_class, **tqdm_kwargs)
+        self.tqdm_class = tqdm_class
+        self.epoch_bar = tqdm_class(total=epochs, unit='epoch')
+        self.on_epoch_end = self.bar2callback(self.epoch_bar)
+        if data_size and batch_size:
+            self.batches = batches = (data_size + batch_size - 1) // batch_size
+        else:
+            self.batches = batches = None
+        self.verbose = verbose
+        if verbose == 1:
+            self.batch_bar = tqdm_class(total=batches, unit='batch', leave=False)
+            self.on_batch_end = self.bar2callback(
+                self.batch_bar, pop=['batch', 'size'],
+                delta=lambda logs: logs.get('size', 1))
+
+    def on_train_begin(self, *_, **__):
+        params = self.params.get
+        auto_total = params('epochs', params('nb_epoch', None))
+        if auto_total is not None and auto_total != self.epoch_bar.total:
+            self.epoch_bar.reset(total=auto_total)
+
+    def on_epoch_begin(self, epoch, *_, **__):
+        if self.epoch_bar.n < epoch:
+            ebar = self.epoch_bar
+            ebar.n = ebar.last_print_n = ebar.initial = epoch
+        if self.verbose:
+            params = self.params.get
+            total = params('samples', params(
+                'nb_sample', params('steps', None))) or self.batches
+            if self.verbose == 2:
+                if hasattr(self, 'batch_bar'):
+                    self.batch_bar.close()
+                self.batch_bar = self.tqdm_class(
+                    total=total, unit='batch', leave=True,
+                    unit_scale=1 / (params('batch_size', 1) or 1))
+                self.on_batch_end = self.bar2callback(
+                    self.batch_bar, pop=['batch', 'size'],
+                    delta=lambda logs: logs.get('size', 1))
+            elif self.verbose == 1:
+                self.batch_bar.unit_scale = 1 / (params('batch_size', 1) or 1)
+                self.batch_bar.reset(total=total)
+            else:
+                raise KeyError('Unknown verbosity')
+
+    def on_train_end(self, *_, **__):
+        if hasattr(self, 'batch_bar'):
+            self.batch_bar.close()
+        self.epoch_bar.close()
+
+    def display(self):
+        """Displays in the current cell in Notebooks."""
+        container = getattr(self.epoch_bar, 'container', None)
+        if container is None:
+            return
+        from .notebook import display
+        display(container)
+        batch_bar = getattr(self, 'batch_bar', None)
+        if batch_bar is not None:
+            display(batch_bar.container)
+
+    @staticmethod
+    def _implements_train_batch_hooks():
+        return True
+
+    @staticmethod
+    def _implements_test_batch_hooks():
+        return True
+
+    @staticmethod
+    def _implements_predict_batch_hooks():
+        return True