[aws-xray-sdk] Improve stubs (#14277)

Author: donbarbos

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

@@ -1,6 +1,5 @@
 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:

stubs/aws-xray-sdk/aws_xray_sdk/core/emitters/udp_emitter.pyi

@@ -1,18 +1,17 @@
 from logging import Logger
+from typing import Final
 
-from aws_xray_sdk.core.daemon_config import DaemonConfig as DaemonConfig
-
-from ..exceptions.exceptions import InvalidDaemonAddressException as InvalidDaemonAddressException
+from aws_xray_sdk.core.models.entity import Entity
 
 log: Logger
-PROTOCOL_HEADER: str
-PROTOCOL_DELIMITER: str
-DEFAULT_DAEMON_ADDRESS: str
+PROTOCOL_HEADER: Final[str]
+PROTOCOL_DELIMITER: Final[str]
+DEFAULT_DAEMON_ADDRESS: Final[str]
 
 class UDPEmitter:
-    def __init__(self, daemon_address="127.0.0.1:2000") -> None: ...
-    def send_entity(self, entity) -> None: ...
-    def set_daemon_address(self, address) -> None: ...
+    def __init__(self, daemon_address: str = "127.0.0.1:2000") -> None: ...
+    def send_entity(self, entity: Entity) -> None: ...
+    def set_daemon_address(self, address: str | None) -> None: ...
     @property
     def ip(self): ...
     @property

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

@@ -1,5 +1,3 @@
-from ..utils.search_pattern import wildcard_match as wildcard_match
-
 class DefaultDynamicNaming:
-    def __init__(self, pattern, fallback) -> None: ...
-    def get_name(self, host_name): ...
+    def __init__(self, pattern: str, fallback: str) -> None: ...
+    def get_name(self, host_name: str | None) -> str: ...

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

@@ -1,29 +1,8 @@
-from .noop_traceid import NoOpTraceId as NoOpTraceId
-from .segment import Segment as Segment
-from .subsegment import Subsegment as Subsegment
-from .traceid import TraceId as TraceId
+from .segment import Segment
+from .subsegment import Subsegment
 
 class DummySegment(Segment):
-    sampled: bool
     def __init__(self, name: str = "dummy") -> None: ...
-    def set_aws(self, aws_meta) -> None: ...
-    def put_http_meta(self, key, value) -> None: ...
-    def put_annotation(self, key, value) -> None: ...
-    def put_metadata(self, key, value, namespace: str = "default") -> None: ...
-    def set_user(self, user) -> None: ...
-    def set_service(self, service_info) -> None: ...
-    def apply_status_code(self, status_code) -> None: ...
-    def add_exception(self, exception, stack, remote: bool = False) -> None: ...
-    def serialize(self) -> None: ...
 
 class DummySubsegment(Subsegment):
-    sampled: bool
     def __init__(self, segment, name: str = "dummy") -> None: ...
-    def set_aws(self, aws_meta) -> None: ...
-    def put_http_meta(self, key, value) -> None: ...
-    def put_annotation(self, key, value) -> None: ...
-    def put_metadata(self, key, value, namespace: str = "default") -> None: ...
-    def set_sql(self, sql) -> None: ...
-    def apply_status_code(self, status_code) -> None: ...
-    def add_exception(self, exception, stack, remote: bool = False) -> None: ...
-    def serialize(self) -> None: ...

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

@@ -1,41 +1,50 @@
+from _typeshed import Incomplete
 from logging import Logger
 from traceback import StackSummary
-from typing import Any
+from typing import Any, Final, Literal, overload
+
+from .subsegment import Subsegment
+from .throwable import Throwable
 
 log: Logger
-ORIGIN_TRACE_HEADER_ATTR_KEY: str
+ORIGIN_TRACE_HEADER_ATTR_KEY: Final[str]
 
 class Entity:
-    id: Any
-    name: Any
-    start_time: Any
-    parent_id: Any
+    id: str
+    name: str
+    start_time: float
+    parent_id: str | None
     sampled: bool
     in_progress: bool
-    http: Any
-    annotations: Any
-    metadata: Any
-    aws: Any
-    cause: Any
-    subsegments: Any
-    end_time: Any
-    def __init__(self, name, entity_id=None) -> None: ...
-    def close(self, end_time=None) -> None: ...
-    def add_subsegment(self, subsegment) -> None: ...
-    def remove_subsegment(self, subsegment) -> None: ...
-    def put_http_meta(self, key, value) -> None: ...
-    def put_annotation(self, key, value) -> None: ...
-    def put_metadata(self, key, value, namespace: str = "default") -> None: ...
+    http: dict[str, dict[str, str | int]]
+    annotations: dict[str, float | str | bool]
+    metadata: dict[str, dict[str, Any]]  # value is any object that can be serialized into JSON string
+    aws: dict[str, Incomplete]
+    cause: dict[str, str | list[Throwable]]
+    subsegments: list[Subsegment]
+    end_time: float
+    def __init__(self, name: str, entity_id: str | None = None) -> None: ...
+    def close(self, end_time: float | None = None) -> None: ...
+    def add_subsegment(self, subsegment: Subsegment) -> None: ...
+    def remove_subsegment(self, subsegment: Subsegment) -> None: ...
+    @overload
+    def put_http_meta(self, key: Literal["status", "content_length"], value: int) -> None: ...
+    @overload
+    def put_http_meta(self, key: Literal["url", "method", "user_agent", "client_ip", "x_forwarded_for"], value: str) -> None: ...
+    def put_annotation(self, key: str, value: float | str | bool) -> None: ...
+    def put_metadata(
+        self, key: str, value: Any, namespace: str = "default"  # value is any object that can be serialized into JSON string
+    ) -> None: ...
     def set_aws(self, aws_meta) -> None: ...
     throttle: bool
     def add_throttle_flag(self) -> None: ...
     fault: bool
     def add_fault_flag(self) -> None: ...
     error: bool
     def add_error_flag(self) -> None: ...
-    def apply_status_code(self, status_code) -> None: ...
+    def apply_status_code(self, status_code: int | None) -> None: ...
     def add_exception(self, exception: Exception, stack: StackSummary, remote: bool = False) -> None: ...
     def save_origin_trace_header(self, trace_header) -> None: ...
     def get_origin_trace_header(self): ...
-    def serialize(self): ...
-    def to_dict(self): ...
+    def serialize(self) -> str: ...
+    def to_dict(self) -> dict[str, Incomplete]: ...

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

@@ -1,24 +1,9 @@
-from typing import Any
+from typing import Final
 
 from .segment import Segment
 
-MUTATION_UNSUPPORTED_MESSAGE: str
+MUTATION_UNSUPPORTED_MESSAGE: Final[str]
 
 class FacadeSegment(Segment):
-    initializing: Any
-    def __init__(self, name, entityid, traceid, sampled) -> None: ...
-    def close(self, end_time=None) -> None: ...
-    def put_http_meta(self, key, value) -> None: ...
-    def put_annotation(self, key, value) -> None: ...
-    def put_metadata(self, key, value, namespace: str = "default") -> None: ...
-    def set_aws(self, aws_meta) -> None: ...
-    def set_user(self, user) -> None: ...
-    def add_throttle_flag(self) -> None: ...
-    def add_fault_flag(self) -> None: ...
-    def add_error_flag(self) -> None: ...
-    def add_exception(self, exception, stack, remote: bool = False) -> None: ...
-    def apply_status_code(self, status_code) -> None: ...
-    def serialize(self) -> None: ...
-    def ready_to_send(self): ...
-    def increment(self) -> None: ...
-    def decrement_ref_counter(self) -> None: ...
+    initializing: bool
+    def __init__(self, name: str, entityid: str | None, traceid: str | None, sampled: bool | None) -> None: ...

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

@@ -1,13 +1,13 @@
-from typing import Any
+from typing import Final
 
-URL: str
-METHOD: str
-USER_AGENT: str
-CLIENT_IP: str
-X_FORWARDED_FOR: str
-STATUS: str
-CONTENT_LENGTH: str
-XRAY_HEADER: str
-ALT_XRAY_HEADER: str
-request_keys: Any
-response_keys: Any
+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]
+request_keys: tuple[str, ...]
+response_keys: tuple[str, ...]

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

@@ -1,6 +1,8 @@
+from typing import ClassVar
+
 class NoOpTraceId:
-    VERSION: str
-    DELIMITER: str
+    VERSION: ClassVar[str]
+    DELIMITER: ClassVar[str]
     start_time: str
     def __init__(self) -> None: ...
-    def to_id(self): ...
+    def to_id(self) -> str: ...

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

@@ -1,44 +1,59 @@
+from _typeshed import Incomplete
 from types import TracebackType
-from typing import Any
+from typing import Final
 
 from ..recorder import AWSXRayRecorder
 from ..utils.atomic_counter import AtomicCounter
+from .dummy_entities import DummySegment
 from .entity import Entity
 from .subsegment import Subsegment
 
-ORIGIN_TRACE_HEADER_ATTR_KEY: str
+ORIGIN_TRACE_HEADER_ATTR_KEY: Final[str]
 
 class SegmentContextManager:
-    name: str
-    segment_kwargs: dict[str, Any]
+    name: str | None
+    segment_kwargs: dict[str, str | bool | None]
     recorder: AWSXRayRecorder
-    segment: Segment
-    def __init__(self, recorder: AWSXRayRecorder, name: str | None = None, **segment_kwargs) -> None: ...
-    def __enter__(self): ...
+    segment: Segment | None
+    def __init__(
+        self,
+        recorder: AWSXRayRecorder,
+        name: str | None = None,
+        *,
+        traceid: str | None = None,
+        parent_id: str | None = None,
+        sampling: bool | None = None,
+    ) -> None: ...
+    def __enter__(self) -> DummySegment | Segment: ...
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None: ...
 
 class Segment(Entity):
-    trace_id: str | None
-    id: str | None
+    trace_id: str
+    id: str
     in_progress: bool
     sampled: bool
     user: str | None
     ref_counter: AtomicCounter
-    parent_id: str | None
+    parent_id: str
     service: dict[str, str]
     def __init__(
-        self, name, entityid: str | None = None, traceid: str | None = None, parent_id: str | None = None, sampled: bool = True
+        self,
+        name: str,
+        entityid: str | None = None,
+        traceid: str | None = None,
+        parent_id: str | None = None,
+        sampled: bool = True,
     ) -> None: ...
     def add_subsegment(self, subsegment: Subsegment) -> None: ...
     def increment(self) -> None: ...
     def decrement_ref_counter(self) -> None: ...
-    def ready_to_send(self): ...
-    def get_total_subsegments_size(self): ...
-    def decrement_subsegments_size(self): ...
-    def remove_subsegment(self, subsegment) -> None: ...
+    def ready_to_send(self) -> bool: ...
+    def get_total_subsegments_size(self) -> int: ...
+    def decrement_subsegments_size(self) -> int: ...
+    def remove_subsegment(self, subsegment: Subsegment) -> None: ...
     def set_user(self, user) -> None: ...
-    def set_service(self, service_info) -> None: ...
-    def set_rule_name(self, rule_name) -> None: ...
-    def to_dict(self): ...
+    def set_service(self, service_info: dict[str, str]) -> None: ...
+    def set_rule_name(self, rule_name: str) -> None: ...
+    def to_dict(self) -> dict[str, Incomplete]: ...

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

@@ -1,25 +1,26 @@
-import time
+from _typeshed import Incomplete
 from types import TracebackType
-from typing import Any
+from typing import Final
 
-from ...core import AWSXRayRecorder
+from ..recorder import AWSXRayRecorder
+from .dummy_entities import DummySubsegment
 from .entity import Entity
 from .segment import Segment
 
-SUBSEGMENT_RECORDING_ATTRIBUTE: str
+SUBSEGMENT_RECORDING_ATTRIBUTE: Final[str]
 
 def set_as_recording(decorated_func, wrapped) -> None: ...
 def is_already_recording(func): ...
 def subsegment_decorator(wrapped, instance, args, kwargs): ...
 
 class SubsegmentContextManager:
     name: str | None
-    subsegment_kwargs: dict[str, Any] | None
+    subsegment_kwargs: dict[str, str]
     recorder: AWSXRayRecorder
-    subsegment: Subsegment
-    def __init__(self, recorder: AWSXRayRecorder, name=None, **subsegment_kwargs) -> None: ...
-    def __call__(self, wrapped, instance, args: list[Any], kwargs: dict[str, Any]): ...
-    def __enter__(self) -> Subsegment | None: ...
+    subsegment: Subsegment | None
+    def __init__(self, recorder: AWSXRayRecorder, name: str | None = None, *, namespace: str = "local") -> None: ...
+    def __call__(self, wrapped, instance, args: list[Incomplete], kwargs: dict[str, Incomplete]): ...
+    def __enter__(self) -> DummySubsegment | Subsegment | None: ...
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None: ...
@@ -29,10 +30,10 @@ class Subsegment(Entity):
     trace_id: str
     type: str
     namespace: str
-    sql: dict[str, Any]
+    sql: dict[str, Incomplete]
     def __init__(self, name: str, namespace: str, segment: Segment) -> None: ...
     def add_subsegment(self, subsegment: Subsegment) -> None: ...
     def remove_subsegment(self, subsegment: Subsegment) -> None: ...
-    def close(self, end_time: time.struct_time | None = None) -> None: ...
-    def set_sql(self, sql: dict[str, Any]) -> None: ...
-    def to_dict(self) -> dict[str, Any]: ...
+    def close(self, end_time: float | None = None) -> None: ...
+    def set_sql(self, sql: dict[str, Incomplete]) -> None: ...
+    def to_dict(self) -> dict[str, Incomplete]: ...