[aws-xray-sdk] Add missing stubs

stubs/aws-xray-sdk/@tests/stubtest_allowlist.txt

@@ -1,3 +1,34 @@
 aws_xray_sdk.core.models.subsegment.subsegment_decorator
 aws_xray_sdk.core.sampling.connector.ServiceConnector.fetch_sampling_rules
 aws_xray_sdk.core.sampling.sampler.ServiceConnector.fetch_sampling_rules
+
+# We can not import 3rd-party libraries in teststubs runtime,
+# but we can use Protocol to replace this types:
+aws_xray_sdk.ext.aiobotocore
+aws_xray_sdk.ext.aiobotocore.patch
+aws_xray_sdk.ext.aiohttp.client
+aws_xray_sdk.ext.aiohttp.middleware
+aws_xray_sdk.ext.bottle.middleware
+aws_xray_sdk.ext.django.apps
+aws_xray_sdk.ext.django.conf
+aws_xray_sdk.ext.django.db
+aws_xray_sdk.ext.django.middleware
+aws_xray_sdk.ext.django.templates
+aws_xray_sdk.ext.flask.middleware
+aws_xray_sdk.ext.flask_sqlalchemy.query
+aws_xray_sdk.ext.httpx
+aws_xray_sdk.ext.httpx.patch
+aws_xray_sdk.ext.mysql
+aws_xray_sdk.ext.mysql.patch
+aws_xray_sdk.ext.pg8000
+aws_xray_sdk.ext.pg8000.patch
+aws_xray_sdk.ext.pymongo
+aws_xray_sdk.ext.pymongo.patch
+aws_xray_sdk.ext.pymysql
+aws_xray_sdk.ext.pymysql.patch
+aws_xray_sdk.ext.pynamodb
+aws_xray_sdk.ext.pynamodb.patch
+aws_xray_sdk.ext.sqlalchemy.query
+aws_xray_sdk.ext.sqlalchemy.util.decorators
+aws_xray_sdk.ext.sqlalchemy_core
+aws_xray_sdk.ext.sqlalchemy_core.patch

stubs/aws-xray-sdk/METADATA.toml

@@ -1,6 +1,2 @@
 version = "2.14.*"
 upstream_repository = "https://github.com/aws/aws-xray-sdk-python"
-partial_stub = true
-
-[tool.stubtest]
-ignore_missing_stub = true

stubs/aws-xray-sdk/aws_xray_sdk/core/__init__.pyi

@@ -2,3 +2,5 @@ from .patcher import patch as patch, patch_all as patch_all
 from .recorder import AWSXRayRecorder as AWSXRayRecorder
 
 xray_recorder: AWSXRayRecorder
+
+__all__ = ["patch", "patch_all", "xray_recorder", "AWSXRayRecorder"]

stubs/aws-xray-sdk/aws_xray_sdk/core/recorder.pyi

@@ -47,6 +47,7 @@ class AWSXRayRecorder:
     def end_segment(self, end_time: time.struct_time | None = None) -> None: ...
     def current_segment(self) -> Segment: ...
     def begin_subsegment(self, name: str, namespace: str = "local") -> DummySubsegment | Subsegment | None: ...
+    def begin_subsegment_without_sampling(self, name: str) -> DummySubsegment | Subsegment | None: ...
     def current_subsegment(self) -> Subsegment | DummySubsegment | None: ...
     def end_subsegment(self, end_time: time.struct_time | None = None) -> None: ...
     def put_annotation(self, key: str, value: Any) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/utils/sqs_message_helper.pyi

@@ -0,0 +1,9 @@
+from _typeshed import Incomplete
+from collections.abc import Mapping
+from typing import Final
+
+SQS_XRAY_HEADER: Final = "AWSTraceHeader"
+
+class SqsMessageHelper:
+    @staticmethod
+    def isSampled(sqs_message: Mapping[str, Incomplete]) -> bool: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/aiobotocore/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/aiobotocore/patch.pyi

@@ -0,0 +1 @@
+def patch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/aiohttp/client.pyi

@@ -0,0 +1,10 @@
+from typing import Final
+
+REMOTE_NAMESPACE: Final = "remote"
+LOCAL_NAMESPACE: Final = "local"
+LOCAL_EXCEPTIONS: tuple[type[Exception], ...]
+
+async def begin_subsegment(session, trace_config_ctx, params): ...
+async def end_subsegment(session, trace_config_ctx, params): ...
+async def end_subsegment_with_exception(session, trace_config_ctx, params): ...
+def aws_xray_trace_config(name=None): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/aiohttp/middleware.pyi

@@ -0,0 +1 @@
+async def middleware(request, handler): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/boto_utils.pyi

@@ -0,0 +1,6 @@
+from _typeshed import Incomplete
+
+whitelist: Incomplete
+
+def inject_header(wrapped, instance, args, kwargs): ...
+def aws_meta_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/botocore/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/botocore/patch.pyi

@@ -0,0 +1 @@
+def patch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/bottle/middleware.pyi

@@ -0,0 +1,9 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+from typing import ClassVar
+
+class XRayMiddleware:
+    name: ClassVar[str]
+    api: ClassVar[int]
+    def __init__(self, recorder) -> None: ...
+    def apply(self, callback: Callable[..., Incomplete], route) -> Callable[..., Incomplete]: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/dbapi2.pyi

@@ -0,0 +1,14 @@
+from typing_extensions import Self
+
+class XRayTracedConn:
+    def __init__(self, conn, meta={}) -> None: ...
+    def cursor(self, *args, **kwargs) -> XRayTracedCursor: ...
+
+class XRayTracedCursor:
+    def __init__(self, cursor, meta={}) -> None: ...
+    def __enter__(self) -> Self: ...
+    def execute(self, query, *args, **kwargs): ...
+    def executemany(self, query, *args, **kwargs): ...
+    def callproc(self, proc, args): ...
+
+def add_sql_meta(meta) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/__init__.pyi

@@ -0,0 +1,3 @@
+from typing import Final
+
+default_app_config: Final = "aws_xray_sdk.ext.django.apps.XRayConfig"

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/apps.pyi

@@ -0,0 +1,8 @@
+from logging import Logger
+from typing import ClassVar
+
+log: Logger
+
+class XRayConfig:
+    name: ClassVar[str]
+    def ready(self) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/conf.pyi

@@ -0,0 +1,17 @@
+from _typeshed import Incomplete
+from typing import Final
+
+DEFAULTS: dict[str, str | bool | tuple[Incomplete] | list[Incomplete] | None]
+XRAY_NAMESPACE: Final = "XRAY_RECORDER"
+SUPPORTED_ENV_VARS: tuple[str, ...]
+
+class XRaySettings:
+    defaults: dict[str, str | bool | tuple[Incomplete] | list[Incomplete] | None]
+    def __init__(self, user_settings=None) -> None: ...
+    @property
+    def user_settings(self): ...
+    def __getattr__(self, attr): ...
+
+settings: XRaySettings
+
+def reload_settings(*, settings: str | None = None, value=None) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/db.pyi

@@ -0,0 +1,12 @@
+from logging import Logger
+
+from aws_xray_sdk.ext.dbapi2 import XRayTracedCursor
+
+log: Logger
+
+def patch_db() -> None: ...
+
+class DjangoXRayTracedCursor(XRayTracedCursor):
+    def execute(self, query, *args, **kwargs): ...
+    def executemany(self, query, *args, **kwargs): ...
+    def callproc(self, proc, args): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/middleware.pyi

@@ -0,0 +1,19 @@
+from _typeshed import Incomplete
+from logging import Logger
+from typing import Final
+
+log: Logger
+
+USER_AGENT_KEY: Final = "HTTP_USER_AGENT"
+X_FORWARDED_KEY: Final = "HTTP_X_FORWARDED_FOR"
+REMOTE_ADDR_KEY: Final = "REMOTE_ADDR"
+HOST_KEY: Final = "HTTP_HOST"
+CONTENT_LENGTH_KEY: Final = "content-length"
+
+class XRayMiddleware:
+    get_response: Incomplete
+    in_lambda_ctx: bool
+
+    def __init__(self, get_response) -> None: ...
+    def __call__(self, request): ...
+    def process_exception(self, request, exception: Exception) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/django/templates.pyi

@@ -0,0 +1,5 @@
+from logging import Logger
+
+log: Logger
+
+def patch_template() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/flask/middleware.pyi

@@ -0,0 +1,6 @@
+from _typeshed import Incomplete
+
+class XRayMiddleware:
+    app: Incomplete
+    in_lambda_ctx: bool
+    def __init__(self, app, recorder) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/flask_sqlalchemy/query.pyi

@@ -0,0 +1,22 @@
+from _typeshed import Incomplete
+
+from aws_xray_sdk.ext.sqlalchemy.query import XRaySession
+
+class XRayBaseQuery: ...
+
+class XRaySignallingSession(XRaySession):
+    app: Incomplete
+    def __init__(self, db, autocommit: bool = False, autoflush: bool = True, **options) -> None: ...
+    def get_bind(self, mapper=None, clause=None): ...
+
+class XRayFlaskSqlAlchemy:
+    def __init__(
+        self,
+        app=None,
+        use_native_unicode: bool = True,
+        session_options=None,
+        metadata=None,
+        query_class: type = ...,
+        model_class: type = ...,
+    ) -> None: ...
+    def create_session(self, options: dict[str, Incomplete]): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/httplib/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import add_ignored as add_ignored, patch as patch, reset_ignored as reset_ignored, unpatch as unpatch
+
+__all__ = ["patch", "unpatch", "add_ignored", "reset_ignored"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/httplib/patch.pyi

@@ -0,0 +1,12 @@
+from typing import Final
+
+httplib_client_module: Final = "http.client"
+PATCH_FLAG: Final = "__xray_patched"
+
+def add_ignored(subclass=None, hostname=None, urls=None) -> None: ...
+def reset_ignored() -> None: ...
+def http_response_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...
+def http_send_request_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...
+def http_read_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...
+def patch() -> None: ...
+def unpatch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/httpx/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/httpx/patch.pyi

@@ -0,0 +1,9 @@
+def patch() -> None: ...
+
+class SyncInstrumentedTransport:
+    def __init__(self, transport): ...
+    def handle_request(self, request): ...
+
+class AsyncInstrumentedTransport:
+    def __init__(self, transport): ...
+    async def handle_async_request(self, request): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/mysql/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/mysql/patch.pyi

@@ -0,0 +1,6 @@
+from typing import Final
+
+MYSQL_ATTR: Final[dict[str, str]]
+
+def patch() -> None: ...
+def sanitize_db_ver(raw: tuple[str]) -> str: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/pg8000/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch, unpatch as unpatch
+
+__all__ = ["patch", "unpatch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/pg8000/patch.pyi

@@ -0,0 +1,2 @@
+def patch() -> None: ...
+def unpatch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/psycopg2/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/psycopg2/patch.pyi

@@ -0,0 +1 @@
+def patch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/pymongo/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/pymongo/patch.pyi

@@ -0,0 +1,8 @@
+class XrayCommandListener:
+    record_full_documents: bool
+    def __init__(self, record_full_documents: bool) -> None: ...
+    def started(self, event) -> None: ...
+    def succeeded(self, event) -> None: ...
+    def failed(self, event) -> None: ...
+
+def patch(record_full_documents: bool = False) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/pymysql/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch, unpatch as unpatch
+
+__all__ = ["patch", "unpatch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/pymysql/patch.pyi

@@ -0,0 +1,3 @@
+def patch() -> None: ...
+def sanitize_db_ver(raw: tuple[str]) -> str: ...
+def unpatch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/pynamodb/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/pynamodb/patch.pyi

@@ -0,0 +1,6 @@
+from typing import Final
+
+PYNAMODB4: Final[bool]
+
+def patch() -> None: ...
+def pynamodb_meta_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/requests/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/requests/patch.pyi

@@ -0,0 +1,2 @@
+def patch() -> None: ...
+def requests_processor(wrapped, instance, args, kwargs, return_value, exception, subsegment, stack) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlalchemy/query.pyi

@@ -0,0 +1,16 @@
+from typing import Any
+
+class XRaySession: ...
+class XRayQuery: ...
+
+class XRaySessionMaker:
+    def __init__(
+        self,
+        bind=None,
+        class_: type = ...,
+        autoflush: bool = True,
+        autocommit: bool = False,
+        expire_on_commit: bool = True,
+        info: dict[Any, Any] | None = None,  # it was taken from sqlalchemy stubs
+        **kw,
+    ) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlalchemy/util/decorators.pyi

@@ -0,0 +1,6 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+
+def decorate_all_functions(function_decorator: Callable[..., Incomplete]) -> Callable[..., Incomplete]: ...
+def xray_on_call(cls, func: Callable[..., Incomplete]) -> Callable[..., Incomplete]: ...
+def parse_bind(bind) -> dict[str, str]: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlalchemy_core/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch, unpatch as unpatch
+
+__all__ = ["patch", "unpatch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlalchemy_core/patch.pyi

@@ -0,0 +1,2 @@
+def patch() -> None: ...
+def unpatch() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlite3/__init__.pyi

@@ -0,0 +1,3 @@
+from .patch import patch as patch
+
+__all__ = ["patch"]

stubs/aws-xray-sdk/aws_xray_sdk/ext/sqlite3/patch.pyi

@@ -0,0 +1,7 @@
+from aws_xray_sdk.ext.dbapi2 import XRayTracedConn
+
+def patch() -> None: ...
+
+class XRayTracedSQLite(XRayTracedConn):
+    def execute(self, *args, **kwargs): ...
+    def executemany(self, *args, **kwargs): ...

stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi

@@ -0,0 +1,21 @@
+import re
+from typing import Final, overload
+
+from aws_xray_sdk.core.models.trace_header import TraceHeader
+
+first_cap_re: Final[re.Pattern[str]]
+all_cap_re: Final[re.Pattern[str]]
+UNKNOWN_HOSTNAME: str = "UNKNOWN HOST"
+
+def inject_trace_header(headers, entity) -> None: ...
+def calculate_sampling_decision(trace_header, recorder, sampling_req): ...
+def construct_xray_header(headers) -> TraceHeader: ...
+def calculate_segment_name(host_name, recorder): ...
+def prepare_response_header(origin_header, segment) -> str: ...
+def to_snake_case(name: str) -> str: ...
+def strip_url(url): ...
+@overload
+def get_hostname(url: str | None) -> str: ...
+@overload
+def get_hostname(url: bytes | bytearray | None) -> str | bytes: ...
+def unwrap(obj: object, attr: str) -> None: ...