[aws-xray-sdk] Annotate more

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

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

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

@@ -1,15 +1,23 @@
-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):
     path: str
     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: ...