From 50cb26b67884b4fa1b244e326366c3e35e61f90b Mon Sep 17 00:00:00 2001 From: donBarbos Date: Sun, 9 Nov 2025 23:18:12 +0400 Subject: [PATCH 1/2] [aws-xray-sdk] Annotate more --- .../@tests/stubtest_allowlist.txt | 3 ++ .../aws_xray_sdk/core/async_context.pyi | 16 ++++++-- .../aws_xray_sdk/core/async_recorder.pyi | 37 ++++++++++++++----- .../aws_xray_sdk/core/context.pyi | 14 +++---- .../aws_xray_sdk/core/daemon_config.pyi | 16 ++++---- .../core/emitters/udp_emitter.pyi | 4 +- .../aws_xray_sdk/core/lambda_launcher.pyi | 20 ++++------ .../aws_xray_sdk/core/models/entity.pyi | 2 +- .../core/models/facade_segment.pyi | 2 +- .../aws_xray_sdk/core/models/http.pyi | 18 ++++----- .../aws_xray_sdk/core/models/segment.pyi | 2 +- .../aws_xray_sdk/core/models/subsegment.pyi | 7 ++-- .../aws_xray_sdk/core/models/throwable.pyi | 14 +++++-- .../aws_xray_sdk/core/patcher.pyi | 5 ++- .../aws_xray_sdk/core/plugins/ec2_plugin.pyi | 6 +-- .../aws_xray_sdk/core/plugins/ecs_plugin.pyi | 4 +- .../core/plugins/elasticbeanstalk_plugin.pyi | 6 +-- .../aws_xray_sdk/core/recorder.pyi | 34 +++++++++-------- .../aws_xray_sdk/core/sampling/rule_cache.pyi | 2 +- .../aws_xray_sdk/core/sampling/sampler.pyi | 4 +- .../core/sampling/sampling_rule.pyi | 36 +++++++++++------- .../core/sampling/target_poller.pyi | 6 ++- stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi | 2 +- 23 files changed, 159 insertions(+), 101 deletions(-) diff --git a/stubs/aws-xray-sdk/@tests/stubtest_allowlist.txt b/stubs/aws-xray-sdk/@tests/stubtest_allowlist.txt index 75131498945f..0ccd2bd8b0ab 100644 --- a/stubs/aws-xray-sdk/@tests/stubtest_allowlist.txt +++ b/stubs/aws-xray-sdk/@tests/stubtest_allowlist.txt @@ -1,6 +1,9 @@ aws_xray_sdk.core.models.subsegment.subsegment_decorator aws_xray_sdk.core.sampling.connector.ServiceConnector.fetch_sampling_rules +# Inconsistency because `context_missing` param can be passed in *args or **kwargs: +aws_xray_sdk.core.async_context.AsyncContext.__init__ + # We can not import 3rd-party libraries in teststubs runtime, # but we can use Protocol to replace this types: aws_xray_sdk.ext.aiobotocore diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi index 7579d0896224..93e021453f70 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi @@ -1,13 +1,21 @@ +from asyncio.events import AbstractEventLoop +from asyncio.tasks import Task, _TaskCompatibleCoro +from typing import Any, TypeVar + from .context import Context as _Context +_T_co = TypeVar("_T_co", covariant=True) + class AsyncContext(_Context): - def __init__(self, *args, loop=None, use_task_factory: bool = True, **kwargs) -> None: ... + def __init__( + self, context_missing: str = "LOG_ERROR", loop: AbstractEventLoop | None = None, use_task_factory: bool = True + ) -> None: ... def clear_trace_entities(self) -> None: ... class TaskLocalStorage: - def __init__(self, loop=None) -> None: ... + def __init__(self, loop: AbstractEventLoop | None = None) -> None: ... def __setattr__(self, name: str, value) -> None: ... - def __getattribute__(self, item: str): ... + def __getattribute__(self, item: str) -> Any: ... def clear(self) -> None: ... -def task_factory(loop, coro): ... +def task_factory(loop: AbstractEventLoop | None, coro: _TaskCompatibleCoro[_T_co]) -> Task[_T_co]: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_recorder.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_recorder.pyi index 4d4bdb8be332..f7e4d588fa29 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_recorder.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_recorder.pyi @@ -1,24 +1,43 @@ +from _typeshed import Incomplete +from collections.abc import Awaitable, Callable, Iterable, Mapping from types import TracebackType +from typing import TypeVar -from .models.segment import SegmentContextManager -from .models.subsegment import SubsegmentContextManager +from .models.dummy_entities import DummySegment, DummySubsegment +from .models.segment import Segment, SegmentContextManager +from .models.subsegment import Subsegment, SubsegmentContextManager from .recorder import AWSXRayRecorder +_T = TypeVar("_T") + class AsyncSegmentContextManager(SegmentContextManager): - async def __aenter__(self): ... + async def __aenter__(self) -> DummySegment | Segment: ... async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None ) -> None: ... class AsyncSubsegmentContextManager(SubsegmentContextManager): - async def __call__(self, wrapped, instance, args, kwargs): ... - async def __aenter__(self): ... + async def __call__( + self, wrapped: Callable[..., Awaitable[_T]], instance, args: Iterable[Incomplete], kwargs: Mapping[str, Incomplete] + ) -> _T: ... + async def __aenter__(self) -> DummySubsegment | Subsegment | None: ... async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None ) -> None: ... class AsyncAWSXRayRecorder(AWSXRayRecorder): - def capture_async(self, name=None): ... - def in_segment_async(self, name=None, **segment_kwargs): ... - def in_subsegment_async(self, name=None, **subsegment_kwargs): ... - async def record_subsegment_async(self, wrapped, instance, args, kwargs, name, namespace, meta_processor): ... + def capture_async(self, name: str | None = None) -> AsyncSubsegmentContextManager: ... + def in_segment_async( + self, name: str | None = None, *, traceid: str | None = None, parent_id: str | None = None, sampling: bool | None = None + ) -> AsyncSegmentContextManager: ... + def in_subsegment_async(self, name: str | None = None, *, namespace: str = "local") -> AsyncSubsegmentContextManager: ... + async def record_subsegment_async( + self, + wrapped: Callable[..., Awaitable[_T]], + instance, + args: Iterable[Incomplete], + kwargs: Mapping[str, Incomplete], + name: str, + namespace: str, + meta_processor: Callable[..., object] | None, + ) -> _T: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/context.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/context.pyi index 3b389b303cc0..c84574342ab4 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/context.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/context.pyi @@ -1,27 +1,27 @@ import time from logging import Logger -from typing import Any +from typing import Final from .models.entity import Entity from .models.segment import Segment from .models.subsegment import Subsegment log: Logger -SUPPORTED_CONTEXT_MISSING: Any -MISSING_SEGMENT_MSG: str -CXT_MISSING_STRATEGY_KEY: str +MISSING_SEGMENT_MSG: Final[str] +SUPPORTED_CONTEXT_MISSING: Final = ("RUNTIME_ERROR", "LOG_ERROR", "IGNORE_ERROR") +CXT_MISSING_STRATEGY_KEY: Final = "AWS_XRAY_CONTEXT_MISSING" class Context: def __init__(self, context_missing: str = "LOG_ERROR") -> None: ... def put_segment(self, segment: Segment) -> None: ... def end_segment(self, end_time: time.struct_time | None = None) -> None: ... def put_subsegment(self, subsegment: Subsegment) -> None: ... - def end_subsegment(self, end_time: time.struct_time | None = None): ... - def get_trace_entity(self): ... + def end_subsegment(self, end_time: time.struct_time | None = None) -> bool: ... + def get_trace_entity(self) -> Entity: ... def set_trace_entity(self, trace_entity: Entity) -> None: ... def clear_trace_entities(self) -> None: ... def handle_context_missing(self) -> None: ... @property - def context_missing(self): ... + def context_missing(self) -> str: ... @context_missing.setter def context_missing(self, value: str) -> None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/daemon_config.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/daemon_config.pyi index e5c4b9609def..35a393a7712b 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/daemon_config.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/daemon_config.pyi @@ -1,13 +1,15 @@ -DAEMON_ADDRESS_KEY: str -DEFAULT_ADDRESS: str +from typing import Final + +DAEMON_ADDRESS_KEY: Final = "AWS_XRAY_DAEMON_ADDRESS" +DEFAULT_ADDRESS: Final = "127.0.0.1:2000" class DaemonConfig: - def __init__(self, daemon_address="127.0.0.1:2000") -> None: ... + def __init__(self, daemon_address: str | None = "127.0.0.1:2000") -> None: ... @property - def udp_ip(self): ... + def udp_ip(self) -> str: ... @property - def udp_port(self): ... + def udp_port(self) -> int: ... @property - def tcp_ip(self): ... + def tcp_ip(self) -> str: ... @property - def tcp_port(self): ... + def tcp_port(self) -> int: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/emitters/udp_emitter.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/emitters/udp_emitter.pyi index fa5c2b24ca4e..affa22d2a484 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/emitters/udp_emitter.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/emitters/udp_emitter.pyi @@ -13,6 +13,6 @@ class UDPEmitter: def send_entity(self, entity: Entity) -> None: ... def set_daemon_address(self, address: str | None) -> None: ... @property - def ip(self): ... + def ip(self) -> str: ... @property - def port(self): ... + def port(self) -> int: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/lambda_launcher.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/lambda_launcher.pyi index 9d0cbf5e3f17..50e94b257de5 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/lambda_launcher.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/lambda_launcher.pyi @@ -1,23 +1,19 @@ from logging import Logger +from typing import Final from .context import Context log: Logger -LAMBDA_TRACE_HEADER_KEY: str -LAMBDA_TASK_ROOT_KEY: str -TOUCH_FILE_DIR: str -TOUCH_FILE_PATH: str +LAMBDA_TRACE_HEADER_KEY: Final = "_X_AMZN_TRACE_ID" +LAMBDA_TASK_ROOT_KEY: Final = "LAMBDA_TASK_ROOT" +TOUCH_FILE_DIR: Final = "/tmp/.aws-xray/" +TOUCH_FILE_PATH: Final = "/tmp/.aws-xray/initialized" -def check_in_lambda(): ... +def check_in_lambda() -> LambdaContext | None: ... class LambdaContext(Context): def __init__(self) -> None: ... - def put_segment(self, segment) -> None: ... - def end_segment(self, end_time=None) -> None: ... - def put_subsegment(self, subsegment) -> None: ... - def get_trace_entity(self): ... - @property + @property # type: ignore[override] def context_missing(self) -> None: ... @context_missing.setter - def context_missing(self, value) -> None: ... - def handle_context_missing(self) -> None: ... + def context_missing(self, value: str) -> None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/entity.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/entity.pyi index 0e19e4d65fa1..5ea373e689fd 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/entity.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/entity.pyi @@ -7,7 +7,7 @@ from .subsegment import Subsegment from .throwable import Throwable log: Logger -ORIGIN_TRACE_HEADER_ATTR_KEY: Final[str] +ORIGIN_TRACE_HEADER_ATTR_KEY: Final = "_origin_trace_header" class Entity: id: str diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/facade_segment.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/facade_segment.pyi index 69c156c46cfe..2ae599e4c968 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/facade_segment.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/facade_segment.pyi @@ -2,7 +2,7 @@ from typing import Final from .segment import Segment -MUTATION_UNSUPPORTED_MESSAGE: Final[str] +MUTATION_UNSUPPORTED_MESSAGE: Final = "FacadeSegments cannot be mutated." class FacadeSegment(Segment): initializing: bool diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/http.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/http.pyi index 4aa14206214e..360493f2f829 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/http.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/http.pyi @@ -1,13 +1,13 @@ from typing import Final -URL: Final[str] -METHOD: Final[str] -USER_AGENT: Final[str] -CLIENT_IP: Final[str] -X_FORWARDED_FOR: Final[str] -STATUS: Final[str] -CONTENT_LENGTH: Final[str] -XRAY_HEADER: Final[str] -ALT_XRAY_HEADER: Final[str] +URL: Final = "url" +METHOD: Final = "method" +USER_AGENT: Final = "user_agent" +CLIENT_IP: Final = "client_ip" +X_FORWARDED_FOR: Final = "x_forwarded_for" +STATUS: Final = "status" +CONTENT_LENGTH: Final = "content_length" +XRAY_HEADER: Final = "X-Amzn-Trace-Id" +ALT_XRAY_HEADER: Final = "HTTP_X_AMZN_TRACE_ID" request_keys: tuple[str, ...] response_keys: tuple[str, ...] diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/segment.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/segment.pyi index 903378f6e6b6..1851d8ea3282 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/segment.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/segment.pyi @@ -8,7 +8,7 @@ from .dummy_entities import DummySegment from .entity import Entity from .subsegment import Subsegment -ORIGIN_TRACE_HEADER_ATTR_KEY: Final[str] +ORIGIN_TRACE_HEADER_ATTR_KEY: Final = "_origin_trace_header" class SegmentContextManager: name: str | None diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi index 57ceefa1cb12..3312e7703a61 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi @@ -1,16 +1,17 @@ from _typeshed import Incomplete +from collections.abc import Callable from types import TracebackType -from typing import Final +from typing import Any, Final from ..recorder import AWSXRayRecorder from .dummy_entities import DummySubsegment from .entity import Entity from .segment import Segment -SUBSEGMENT_RECORDING_ATTRIBUTE: Final[str] +SUBSEGMENT_RECORDING_ATTRIBUTE: Final = "_self___SUBSEGMENT_RECORDING_ATTRIBUTE__" def set_as_recording(decorated_func, wrapped) -> None: ... -def is_already_recording(func): ... +def is_already_recording(func: Callable[..., Any]) -> bool: ... # return type does not matter def subsegment_decorator(wrapped, instance, args, kwargs): ... class SubsegmentContextManager: diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/throwable.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/throwable.pyi index 54b918634306..e4e0f748fd59 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/throwable.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/throwable.pyi @@ -1,7 +1,9 @@ -from _typeshed import Incomplete from logging import Logger from traceback import StackSummary from typing import TypedDict, type_check_only +from typing_extensions import NotRequired + +log: Logger @type_check_only class _StackInfo(TypedDict): @@ -9,7 +11,13 @@ class _StackInfo(TypedDict): line: int label: str -log: Logger +@type_check_only +class _ThrowableAttrs(TypedDict): + id: str + message: NotRequired[str] + type: str + remote: bool + stack: NotRequired[list[_StackInfo]] class Throwable: id: str @@ -18,4 +26,4 @@ class Throwable: remote: bool stack: list[_StackInfo] | None def __init__(self, exception: Exception, stack: StackSummary, remote: bool = False) -> None: ... - def to_dict(self) -> dict[str, Incomplete]: ... + def to_dict(self) -> _ThrowableAttrs: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/patcher.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/patcher.pyi index 20cee42f8ef9..6b9f5d73226a 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/patcher.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/patcher.pyi @@ -1,9 +1,10 @@ from collections.abc import Iterable from logging import Logger +from typing import Final log: Logger -SUPPORTED_MODULES: tuple[str, ...] -NO_DOUBLE_PATCH: tuple[str, ...] +SUPPORTED_MODULES: Final[tuple[str, ...]] +NO_DOUBLE_PATCH: Final[tuple[str, ...]] def patch_all(double_patch: bool = False) -> None: ... def patch( diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ec2_plugin.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ec2_plugin.pyi index aae2f13951b7..84c29f5631cf 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ec2_plugin.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ec2_plugin.pyi @@ -3,9 +3,9 @@ from logging import Logger from typing import Any, Final, overload log: Logger -SERVICE_NAME: Final[str] -ORIGIN: Final[str] -IMDS_URL: Final[str] +SERVICE_NAME: Final = "ec2" +ORIGIN: Final = "AWS::EC2::Instance" +IMDS_URL: Final = "http://169.254.169.254/latest/" def initialize() -> None: ... def get_token() -> str | None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ecs_plugin.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ecs_plugin.pyi index b2cd96d3dcba..019494fb0c57 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ecs_plugin.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ecs_plugin.pyi @@ -2,7 +2,7 @@ from logging import Logger from typing import Final log: Logger -SERVICE_NAME: Final[str] -ORIGIN: Final[str] +SERVICE_NAME: Final = "ecs" +ORIGIN: Final = "AWS::ECS::Container" def initialize() -> None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/elasticbeanstalk_plugin.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/elasticbeanstalk_plugin.pyi index 3ae45f58cded..9ac631478457 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/elasticbeanstalk_plugin.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/elasticbeanstalk_plugin.pyi @@ -2,8 +2,8 @@ from logging import Logger from typing import Final log: Logger -CONF_PATH: Final[str] -SERVICE_NAME: Final[str] -ORIGIN: Final[str] +CONF_PATH: Final = "/var/elasticbeanstalk/xray/environment.conf" +SERVICE_NAME: Final = "elastic_beanstalk" +ORIGIN: Final = "AWS::ElasticBeanstalk::Environment" def initialize() -> None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/recorder.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/recorder.pyi index ee8f964f3e30..8b05f527a5a1 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/recorder.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/recorder.pyi @@ -1,8 +1,8 @@ import time -from _typeshed import FileDescriptorOrPath -from collections.abc import Callable, Iterable +from _typeshed import FileDescriptorOrPath, Incomplete +from collections.abc import Callable, Iterable, Mapping from logging import Logger -from typing import Any +from typing import Any, Final, TypeVar from .context import Context from .emitters.udp_emitter import UDPEmitter @@ -15,11 +15,13 @@ from .sampling.sampler import DefaultSampler from .streaming.default_streaming import DefaultStreaming log: Logger -TRACING_NAME_KEY: str -DAEMON_ADDR_KEY: str -CONTEXT_MISSING_KEY: str -XRAY_META: dict[str, dict[str, str]] -SERVICE_INFO: dict[str, str] +TRACING_NAME_KEY: Final = "AWS_XRAY_TRACING_NAME" +DAEMON_ADDR_KEY: Final = "AWS_XRAY_DAEMON_ADDRESS" +CONTEXT_MISSING_KEY: Final = "AWS_XRAY_CONTEXT_MISSING" +XRAY_META: Final[dict[str, dict[str, str]]] +SERVICE_INFO: Final[dict[str, str]] + +_T = TypeVar("_T") class AWSXRayRecorder: def __init__(self) -> None: ... @@ -40,8 +42,10 @@ class AWSXRayRecorder: sampler: LocalSampler | DefaultSampler | None = None, stream_sql: bool | None = True, ) -> None: ... - def in_segment(self, name: str | None = None, **segment_kwargs) -> SegmentContextManager: ... - def in_subsegment(self, name: str | None = None, **subsegment_kwargs) -> SubsegmentContextManager: ... + def in_segment( + self, name: str | None = None, *, traceid: str | None = None, parent_id: str | None = None, sampling: bool | None = None + ) -> SegmentContextManager: ... + def in_subsegment(self, name: str | None = None, *, namespace: str = "local") -> SubsegmentContextManager: ... def begin_segment( self, name: str | None = None, traceid: str | None = None, parent_id: str | None = None, sampling: bool | None = None ) -> Segment | DummySegment: ... @@ -61,14 +65,14 @@ class AWSXRayRecorder: def capture(self, name: str | None = None) -> SubsegmentContextManager: ... def record_subsegment( self, - wrapped: Callable[..., Any], + wrapped: Callable[..., _T], instance: Any, - args: list[Any], - kwargs: dict[str, Any], + args: Iterable[Incomplete], + kwargs: Mapping[str, Incomplete], name: str, namespace: str, - meta_processor: Callable[..., object], - ) -> Any: ... + meta_processor: Callable[..., object] | None, + ) -> _T: ... @property def enabled(self) -> bool: ... @enabled.setter diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_cache.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_cache.pyi index 6fb06c428f2e..c6e06134ace7 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_cache.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_cache.pyi @@ -4,7 +4,7 @@ TTL: Final = 3600 class RuleCache: def __init__(self) -> None: ... - def get_matched_rule(self, sampling_req, now): ... + def get_matched_rule(self, sampling_req, now: float): ... def load_rules(self, rules) -> None: ... def load_targets(self, targets_dict) -> None: ... @property diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampler.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampler.pyi index 53235a4db593..0bc379e1cb0b 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampler.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampler.pyi @@ -1,5 +1,7 @@ from logging import Logger +from aws_xray_sdk.core.daemon_config import DaemonConfig + log: Logger class DefaultSampler: @@ -7,7 +9,7 @@ class DefaultSampler: def start(self) -> None: ... def should_trace(self, sampling_req=None): ... def load_local_rules(self, rules) -> None: ... - def load_settings(self, daemon_config, context, origin=None) -> None: ... + def load_settings(self, daemon_config: DaemonConfig, context, origin=None) -> None: ... @property def xray_client(self): ... @xray_client.setter diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampling_rule.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampling_rule.pyi index 1ea487cfb0c7..f67af0b60997 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampling_rule.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampling_rule.pyi @@ -1,13 +1,23 @@ +from typing import Literal, TypedDict, type_check_only + +from .reservoir import Reservoir + +@type_check_only +class _Stats(TypedDict): + request_count: int + borrow_count: int + sampled_count: int + class SamplingRule: def __init__( - self, name, priority, rate, reservoir_size, host=None, method=None, path=None, service=None, service_type=None + self, name: str, priority, rate, reservoir_size, host=None, method=None, path=None, service=None, service_type=None ) -> None: ... - def match(self, sampling_req): ... - def is_default(self): ... - def snapshot_statistics(self): ... + def match(self, sampling_req) -> bool: ... + def is_default(self) -> bool: ... + def snapshot_statistics(self) -> _Stats: ... def merge(self, rule) -> None: ... - def ever_matched(self): ... - def time_to_report(self): ... + def ever_matched(self) -> bool: ... + def time_to_report(self) -> Literal[True] | None: ... def increment_request_count(self) -> None: ... def increment_borrow_count(self) -> None: ... def increment_sampled_count(self) -> None: ... @@ -16,18 +26,18 @@ class SamplingRule: @rate.setter def rate(self, v) -> None: ... @property - def name(self): ... + def name(self) -> str: ... @property def priority(self): ... @property - def reservoir(self): ... + def reservoir(self) -> Reservoir: ... @reservoir.setter - def reservoir(self, v) -> None: ... + def reservoir(self, v: Reservoir) -> None: ... @property - def can_borrow(self): ... + def can_borrow(self) -> bool: ... @property - def request_count(self): ... + def request_count(self) -> int: ... @property - def borrow_count(self): ... + def borrow_count(self) -> int: ... @property - def sampled_count(self): ... + def sampled_count(self) -> int: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/target_poller.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/target_poller.pyi index 7fdc45bf8f0b..7a1a5ff753b1 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/target_poller.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/target_poller.pyi @@ -1,7 +1,11 @@ from logging import Logger +from .connector import ServiceConnector +from .rule_cache import RuleCache +from .rule_poller import RulePoller + log: Logger class TargetPoller: - def __init__(self, cache, rule_poller, connector) -> None: ... + def __init__(self, cache: RuleCache, rule_poller: RulePoller, connector: ServiceConnector) -> None: ... def start(self) -> None: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi index daf9f4501bcf..d61f173ceb75 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/ext/util.pyi @@ -5,7 +5,7 @@ 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" +UNKNOWN_HOSTNAME: Final = "UNKNOWN HOST" def inject_trace_header(headers, entity) -> None: ... def calculate_sampling_decision(trace_header, recorder, sampling_req): ... From ee8ec01053e6bc72a9bb2d5965c72a997323a3e4 Mon Sep 17 00:00:00 2001 From: donBarbos Date: Mon, 10 Nov 2025 21:09:09 +0400 Subject: [PATCH 2/2] Accept suggestions --- stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi | 6 ++++-- stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi index 93e021453f70..85c366e821ee 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/async_context.pyi @@ -14,8 +14,10 @@ class AsyncContext(_Context): class TaskLocalStorage: def __init__(self, loop: AbstractEventLoop | None = None) -> None: ... - def __setattr__(self, name: str, value) -> None: ... - def __getattribute__(self, item: str) -> Any: ... + # Sets unknown items on the current task's context attribute + def __setattr__(self, name: str, value: Any) -> None: ... + # Returns unknown items from the current tasks context attribute + def __getattribute__(self, item: str) -> Any | None: ... def clear(self) -> None: ... def task_factory(loop: AbstractEventLoop | None, coro: _TaskCompatibleCoro[_T_co]) -> Task[_T_co]: ... diff --git a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi index 3312e7703a61..27481f85870b 100644 --- a/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi +++ b/stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi @@ -1,7 +1,7 @@ from _typeshed import Incomplete from collections.abc import Callable from types import TracebackType -from typing import Any, Final +from typing import Final from ..recorder import AWSXRayRecorder from .dummy_entities import DummySubsegment @@ -11,7 +11,7 @@ from .segment import Segment SUBSEGMENT_RECORDING_ATTRIBUTE: Final = "_self___SUBSEGMENT_RECORDING_ATTRIBUTE__" def set_as_recording(decorated_func, wrapped) -> None: ... -def is_already_recording(func: Callable[..., Any]) -> bool: ... # return type does not matter +def is_already_recording(func: Callable[..., object]) -> bool: ... def subsegment_decorator(wrapped, instance, args, kwargs): ... class SubsegmentContextManager: