[aws-xray-sdk] Annotate more (#15004)

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

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

@@ -1,13 +1,23 @@
+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 __setattr__(self, name: str, value) -> None: ...
-    def __getattribute__(self, item: str): ...
+    def __init__(self, loop: AbstractEventLoop | None = None) -> None: ...
+    # 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, coro): ...
+def task_factory(loop: AbstractEventLoop | None, coro: _TaskCompatibleCoro[_T_co]) -> Task[_T_co]: ...

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: ...

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: ...

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: ...

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: ...

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: ...

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

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

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, ...]

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

stubs/aws-xray-sdk/aws_xray_sdk/core/models/subsegment.pyi

@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from collections.abc import Callable
 from types import TracebackType
 from typing import Final
 
@@ -7,10 +8,10 @@ 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[..., object]) -> bool: ...
 def subsegment_decorator(wrapped, instance, args, kwargs): ...
 
 class SubsegmentContextManager:

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: ...

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(

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: ...

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: ...

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: ...

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

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

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

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: ...

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: ...

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): ...