Addition of aws-xray-sdk package (#5961)

pyrightconfig.stricter.json

@@ -25,6 +25,7 @@
         "stdlib/tkinter",
         "stdlib/xml/dom",
         "stdlib/xml/sax",
+        "stubs/aws-xray-sdk",
         "stubs/babel",
         "stubs/backports",
         "stubs/backports_abc",

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

@@ -0,0 +1,4 @@
+aws_xray_sdk.core.async_recorder.subsegment_decorator
+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

stubs/aws-xray-sdk/METADATA.toml

@@ -0,0 +1 @@
+version = "2.8"

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

@@ -0,0 +1,3 @@
+from .sdk_config import SDKConfig
+
+global_sdk_config: SDKConfig

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

@@ -0,0 +1,4 @@
+from .patcher import patch as patch, patch_all as patch_all
+from .recorder import AWSXRayRecorder as AWSXRayRecorder
+
+xray_recorder: AWSXRayRecorder

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

@@ -0,0 +1,15 @@
+from typing import Any
+
+from .context import Context as _Context
+
+class AsyncContext(_Context):
+    def __init__(self, *args, loop: Any | None = ..., use_task_factory: bool = ..., **kwargs) -> None: ...
+    def clear_trace_entities(self) -> None: ...
+
+class TaskLocalStorage:
+    def __init__(self, loop: Any | None = ...) -> None: ...
+    def __setattr__(self, name, value) -> None: ...
+    def __getattribute__(self, item): ...
+    def clear(self) -> None: ...
+
+def task_factory(loop, coro): ...

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

@@ -0,0 +1,25 @@
+from typing import Any
+
+from .models.segment import SegmentContextManager as SegmentContextManager
+from .models.subsegment import (
+    SubsegmentContextManager as SubsegmentContextManager,
+    is_already_recording as is_already_recording,
+    subsegment_decorator as subsegment_decorator,
+)
+from .recorder import AWSXRayRecorder as AWSXRayRecorder
+from .utils import stacktrace as stacktrace
+
+class AsyncSegmentContextManager(SegmentContextManager):
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+
+class AsyncSubsegmentContextManager(SubsegmentContextManager):
+    async def __call__(self, wrapped, instance, args, kwargs): ...
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+
+class AsyncAWSXRayRecorder(AWSXRayRecorder):
+    def capture_async(self, name: Any | None = ...): ...
+    def in_segment_async(self, name: Any | None = ..., **segment_kwargs): ...
+    def in_subsegment_async(self, name: Any | None = ..., **subsegment_kwargs): ...
+    async def record_subsegment_async(self, wrapped, instance, args, kwargs, name, namespace, meta_processor): ...

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

@@ -0,0 +1,30 @@
+import time
+from logging import Logger
+from typing import Any
+
+from .. import global_sdk_config as global_sdk_config
+from .exceptions.exceptions import SegmentNotFoundException as SegmentNotFoundException
+from .models.dummy_entities import DummySegment as DummySegment
+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
+
+class Context:
+    def __init__(self, context_missing: str = ...) -> None: ...
+    def put_segment(self, segment: Segment) -> None: ...
+    def end_segment(self, end_time: time.struct_time | None = ...) -> None: ...
+    def put_subsegment(self, subsegment: Subsegment) -> None: ...
+    def end_subsegment(self, end_time: time.struct_time | None = ...): ...
+    def get_trace_entity(self): ...
+    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): ...
+    @context_missing.setter
+    def context_missing(self, value: str) -> None: ...

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

@@ -0,0 +1,15 @@
+from .exceptions.exceptions import InvalidDaemonAddressException as InvalidDaemonAddressException
+
+DAEMON_ADDRESS_KEY: str
+DEFAULT_ADDRESS: str
+
+class DaemonConfig:
+    def __init__(self, daemon_address=...) -> None: ...
+    @property
+    def udp_ip(self): ...
+    @property
+    def udp_port(self): ...
+    @property
+    def tcp_ip(self): ...
+    @property
+    def tcp_port(self): ...

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

@@ -0,0 +1,19 @@
+from logging import Logger
+
+from aws_xray_sdk.core.daemon_config import DaemonConfig as DaemonConfig
+
+from ..exceptions.exceptions import InvalidDaemonAddressException as InvalidDaemonAddressException
+
+log: Logger
+PROTOCOL_HEADER: str
+PROTOCOL_DELIMITER: str
+DEFAULT_DAEMON_ADDRESS: str
+
+class UDPEmitter:
+    def __init__(self, daemon_address=...) -> None: ...
+    def send_entity(self, entity) -> None: ...
+    def set_daemon_address(self, address) -> None: ...
+    @property
+    def ip(self): ...
+    @property
+    def port(self): ...

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

@@ -0,0 +1,8 @@
+class InvalidSamplingManifestError(Exception): ...
+class SegmentNotFoundException(Exception): ...
+class InvalidDaemonAddressException(Exception): ...
+class SegmentNameMissingException(Exception): ...
+class SubsegmentNameMissingException(Exception): ...
+class FacadeSegmentMutationException(Exception): ...
+class MissingPluginNames(Exception): ...
+class AlreadyEndedException(Exception): ...

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

@@ -0,0 +1,28 @@
+from logging import Logger
+from typing import Any
+
+from aws_xray_sdk import global_sdk_config as global_sdk_config
+
+from .context import Context as Context
+from .models.facade_segment import FacadeSegment as FacadeSegment
+from .models.trace_header import TraceHeader as TraceHeader
+
+log: Logger
+LAMBDA_TRACE_HEADER_KEY: str
+LAMBDA_TASK_ROOT_KEY: str
+TOUCH_FILE_DIR: str
+TOUCH_FILE_PATH: str
+
+def check_in_lambda(): ...
+
+class LambdaContext(Context):
+    def __init__(self) -> None: ...
+    def put_segment(self, segment) -> None: ...
+    def end_segment(self, end_time: Any | None = ...) -> None: ...
+    def put_subsegment(self, subsegment) -> None: ...
+    def get_trace_entity(self): ...
+    @property
+    def context_missing(self) -> None: ...
+    @context_missing.setter
+    def context_missing(self, value) -> None: ...
+    def handle_context_missing(self) -> None: ...

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

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

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

@@ -0,0 +1,29 @@
+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
+
+class DummySegment(Segment):
+    sampled: bool
+    def __init__(self, name: str = ...) -> 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 = ...) -> 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 = ...) -> None: ...
+    def serialize(self) -> None: ...
+
+class DummySubsegment(Subsegment):
+    sampled: bool
+    def __init__(self, segment, name: str = ...) -> 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 = ...) -> None: ...
+    def set_sql(self, sql) -> None: ...
+    def apply_status_code(self, status_code) -> None: ...
+    def add_exception(self, exception, stack, remote: bool = ...) -> None: ...
+    def serialize(self) -> None: ...

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

@@ -0,0 +1,47 @@
+from logging import Logger
+from traceback import StackSummary
+from typing import Any
+
+from ..exceptions.exceptions import AlreadyEndedException as AlreadyEndedException
+from ..utils.compat import annotation_value_types as annotation_value_types, string_types as string_types
+from ..utils.conversion import metadata_to_dict as metadata_to_dict
+from . import http as http
+from .throwable import Throwable as Throwable
+
+log: Logger
+ORIGIN_TRACE_HEADER_ATTR_KEY: str
+
+class Entity:
+    id: Any
+    name: Any
+    start_time: Any
+    parent_id: Any
+    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: Any | None = ...) -> None: ...
+    def close(self, end_time: Any | 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 = ...) -> 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 add_exception(self, exception: Exception, stack: StackSummary, remote: bool = ...) -> None: ...
+    def save_origin_trace_header(self, trace_header) -> None: ...
+    def get_origin_trace_header(self): ...
+    def serialize(self): ...
+    def to_dict(self): ...

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

@@ -0,0 +1,25 @@
+from typing import Any
+
+from ..exceptions.exceptions import FacadeSegmentMutationException as FacadeSegmentMutationException
+from .segment import Segment as Segment
+
+MUTATION_UNSUPPORTED_MESSAGE: str
+
+class FacadeSegment(Segment):
+    initializing: Any
+    def __init__(self, name, entityid, traceid, sampled) -> None: ...
+    def close(self, end_time: Any | 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 = ...) -> 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 = ...) -> 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: ...

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

@@ -0,0 +1,13 @@
+from typing import Any
+
+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

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

@@ -0,0 +1,6 @@
+class NoOpTraceId:
+    VERSION: str
+    DELIMITER: str
+    start_time: str
+    def __init__(self) -> None: ...
+    def to_id(self): ...

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

@@ -0,0 +1,43 @@
+from typing import Any
+
+from ..exceptions.exceptions import SegmentNameMissingException as SegmentNameMissingException
+from ..recorder import AWSXRayRecorder
+from ..utils.atomic_counter import AtomicCounter as AtomicCounter
+from .entity import Entity as Entity
+from .subsegment import Subsegment
+from .traceid import TraceId as TraceId
+
+ORIGIN_TRACE_HEADER_ATTR_KEY: str
+
+class SegmentContextManager:
+    name: str
+    segment_kwargs: dict[str, Any]
+    recorder: AWSXRayRecorder
+    segment: Segment
+    def __init__(self, recorder: AWSXRayRecorder, name: str | None = ..., **segment_kwargs) -> None: ...
+    def __enter__(self): ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...
+
+class Segment(Entity):
+    trace_id: str | None
+    id: str | None
+    in_progress: bool
+    sampled: bool
+    user: str | None
+    ref_counter: AtomicCounter
+    parent_id: str | None
+    service: dict[str, str]
+    def __init__(
+        self, name, entityid: str | None = ..., traceid: str | None = ..., parent_id: str | None = ..., sampled: bool = ...
+    ) -> 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 set_user(self, user) -> None: ...
+    def set_service(self, service_info) -> None: ...
+    def set_rule_name(self, rule_name) -> None: ...
+    def to_dict(self): ...

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

@@ -0,0 +1,36 @@
+import time
+from typing import Any
+
+from ...core import AWSXRayRecorder
+from ..exceptions.exceptions import SegmentNotFoundException as SegmentNotFoundException
+from .entity import Entity as Entity
+from .segment import Segment
+
+SUBSEGMENT_RECORDING_ATTRIBUTE: 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
+    recorder: AWSXRayRecorder
+    subsegment: Subsegment
+    def __init__(self, recorder: AWSXRayRecorder, name: Any | None = ..., **subsegment_kwargs) -> None: ...
+    def __call__(self, wrapped, instance, args: list[Any], kwargs: dict[str, Any]): ...
+    def __enter__(self) -> Subsegment: ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...
+
+class Subsegment(Entity):
+    parent_segment: Segment
+    trace_id: str
+    type: str
+    namespace: str
+    sql: dict[str, Any]
+    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: ...
+    def set_sql(self, sql: dict[str, Any]) -> None: ...
+    def to_dict(self) -> dict[str, Any]: ...

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

@@ -0,0 +1,14 @@
+from typing import Any
+
+from ..utils.compat import string_types as string_types
+
+log: Any
+
+class Throwable:
+    id: Any
+    message: Any
+    type: Any
+    remote: Any
+    stack: Any
+    def __init__(self, exception, stack, remote: bool = ...) -> None: ...
+    def to_dict(self): ...

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

@@ -0,0 +1,24 @@
+from typing import Any
+
+log: Any
+ROOT: str
+PARENT: str
+SAMPLE: str
+SELF: str
+HEADER_DELIMITER: str
+
+class TraceHeader:
+    def __init__(
+        self, root: str | None = ..., parent: str | None = ..., sampled: bool | None = ..., data: dict[str, Any] | None = ...
+    ) -> None: ...
+    @classmethod
+    def from_header_str(cls, header): ...
+    def to_header_str(self): ...
+    @property
+    def root(self): ...
+    @property
+    def parent(self): ...
+    @property
+    def sampled(self): ...
+    @property
+    def data(self): ...

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

@@ -0,0 +1,8 @@
+from typing import Any
+
+class TraceId:
+    VERSION: str
+    DELIMITER: str
+    start_time: Any
+    def __init__(self) -> None: ...
+    def to_id(self): ...

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

@@ -0,0 +1,13 @@
+from logging import Logger
+from typing import Any, Iterable
+
+from aws_xray_sdk import global_sdk_config as global_sdk_config
+
+from .utils.compat import PY2 as PY2, is_classmethod as is_classmethod, is_instance_method as is_instance_method
+
+log: Logger
+SUPPORTED_MODULES: Any
+NO_DOUBLE_PATCH: Any
+
+def patch_all(double_patch: bool = ...) -> None: ...
+def patch(modules_to_patch: Iterable[str], raise_errors: bool = ..., ignore_module_patterns: str | None = ...) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ec2_plugin.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+
+log: Any
+SERVICE_NAME: str
+ORIGIN: str
+IMDS_URL: str
+
+def initialize() -> None: ...
+def get_token(): ...
+def get_metadata(token: Any | None = ...): ...
+def parse_metadata_json(json_str): ...
+def do_request(url, headers: Any | None = ..., method: str = ...): ...

stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/ecs_plugin.pyi

@@ -0,0 +1,7 @@
+from typing import Any
+
+log: Any
+SERVICE_NAME: str
+ORIGIN: str
+
+def initialize() -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/plugins/elasticbeanstalk_plugin.pyi

@@ -0,0 +1,8 @@
+from typing import Any
+
+log: Any
+CONF_PATH: str
+SERVICE_NAME: str
+ORIGIN: str
+
+def initialize() -> None: ...

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

@@ -0,0 +1,8 @@
+from typing import Any
+
+from ..exceptions.exceptions import MissingPluginNames as MissingPluginNames
+
+module_prefix: str
+PLUGIN_MAPPING: Any
+
+def get_plugin_modules(plugins): ...

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

@@ -0,0 +1,124 @@
+import time
+from logging import Logger
+from typing import Any, Callable, Iterable
+
+from aws_xray_sdk import global_sdk_config as global_sdk_config
+from aws_xray_sdk.version import VERSION as VERSION
+
+from .context import Context as Context
+from .daemon_config import DaemonConfig as DaemonConfig
+from .emitters.udp_emitter import UDPEmitter as UDPEmitter
+from .exceptions.exceptions import (
+    SegmentNameMissingException as SegmentNameMissingException,
+    SegmentNotFoundException as SegmentNotFoundException,
+)
+from .lambda_launcher import check_in_lambda as check_in_lambda
+from .models.default_dynamic_naming import DefaultDynamicNaming as DefaultDynamicNaming
+from .models.dummy_entities import DummySegment as DummySegment, DummySubsegment as DummySubsegment
+from .models.segment import Segment as Segment, SegmentContextManager as SegmentContextManager
+from .models.subsegment import Subsegment as Subsegment, SubsegmentContextManager as SubsegmentContextManager
+from .plugins.utils import get_plugin_modules as get_plugin_modules
+from .sampling.local.sampler import LocalSampler as LocalSampler
+from .sampling.sampler import DefaultSampler as DefaultSampler
+from .streaming.default_streaming import DefaultStreaming as DefaultStreaming
+from .utils import stacktrace as stacktrace
+from .utils.compat import string_types as string_types
+
+log: Logger
+TRACING_NAME_KEY: str
+DAEMON_ADDR_KEY: str
+CONTEXT_MISSING_KEY: str
+XRAY_META: Any
+SERVICE_INFO: Any
+
+class AWSXRayRecorder:
+    def __init__(self) -> None: ...
+    def configure(
+        self,
+        sampling: bool | None = ...,
+        plugins: Iterable[str] | None = ...,
+        context_missing: str | None = ...,
+        sampling_rules: dict[str, Any] | str | None = ...,
+        daemon_address: str | None = ...,
+        service: str | None = ...,
+        context: Context | None = ...,
+        emitter: UDPEmitter | None = ...,
+        streaming: DefaultStreaming | None = ...,
+        dynamic_naming: DefaultDynamicNaming | None = ...,
+        streaming_threshold: int | None = ...,
+        max_trace_back: int | None = ...,
+        sampler: LocalSampler | DefaultSampler | None = ...,
+        stream_sql: bool | None = ...,
+    ) -> None: ...
+    def in_segment(self, name: str | None = ..., **segment_kwargs): ...
+    def in_subsegment(self, name: str | None = ..., **subsegment_kwargs): ...
+    def begin_segment(
+        self, name: str | None = ..., traceid: str | None = ..., parent_id: str | None = ..., sampling: bool | None = ...
+    ): ...
+    def end_segment(self, end_time: time.struct_time | None = ...) -> None: ...
+    def current_segment(self): ...
+    def begin_subsegment(self, name: str, namespace: str = ...): ...
+    def current_subsegment(self): ...
+    def end_subsegment(self, end_time: time.struct_time | None = ...) -> None: ...
+    def put_annotation(self, key: str, value: Any) -> None: ...
+    def put_metadata(self, key: str, value: Any, namespace: str = ...) -> None: ...
+    def is_sampled(self): ...
+    def get_trace_entity(self): ...
+    def set_trace_entity(self, trace_entity: Context) -> None: ...
+    def clear_trace_entities(self) -> None: ...
+    def stream_subsegments(self) -> None: ...
+    def capture(self, name: str | None = ...): ...
+    def record_subsegment(
+        self,
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: list[Any],
+        kwargs: dict[str, Any],
+        name: str,
+        namespace: str,
+        meta_processor: Callable[..., Any],
+    ): ...
+    @property
+    def enabled(self): ...
+    @enabled.setter
+    def enabled(self, value: bool) -> None: ...
+    @property
+    def sampling(self): ...
+    @sampling.setter
+    def sampling(self, value: bool) -> None: ...
+    @property
+    def sampler(self): ...
+    @sampler.setter
+    def sampler(self, value: LocalSampler | DefaultSampler) -> None: ...
+    @property
+    def service(self): ...
+    @service.setter
+    def service(self, value: str) -> None: ...
+    @property
+    def dynamic_naming(self): ...
+    @dynamic_naming.setter
+    def dynamic_naming(self, value) -> None: ...
+    @property
+    def context(self): ...
+    @context.setter
+    def context(self, cxt: Context) -> None: ...
+    @property
+    def emitter(self): ...
+    @emitter.setter
+    def emitter(self, value: UDPEmitter) -> None: ...
+    @property
+    def streaming(self): ...
+    @streaming.setter
+    def streaming(self, value: DefaultStreaming) -> None: ...
+    @property
+    def streaming_threshold(self): ...
+    @streaming_threshold.setter
+    def streaming_threshold(self, value: int) -> None: ...
+    @property
+    def max_trace_back(self): ...
+    @max_trace_back.setter
+    def max_trace_back(self, value: int) -> None: ...
+    @property
+    def stream_sql(self): ...
+    @stream_sql.setter
+    def stream_sql(self, value: bool) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/connector.pyi

@@ -0,0 +1,17 @@
+from datetime import datetime as datetime
+
+from aws_xray_sdk.core.context import Context as Context
+from aws_xray_sdk.core.models.dummy_entities import DummySegment as DummySegment
+from aws_xray_sdk.core.utils.compat import PY2 as PY2
+
+from .sampling_rule import SamplingRule as SamplingRule
+
+class ServiceConnector:
+    def __init__(self) -> None: ...
+    def fetch_sampling_rules(self): ...
+    def fetch_sampling_target(self, rules): ...
+    def setup_xray_client(self, ip, port, client) -> None: ...
+    @property
+    def context(self): ...
+    @context.setter
+    def context(self, v) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/local/reservoir.pyi

@@ -0,0 +1,8 @@
+from typing import Any
+
+class Reservoir:
+    traces_per_sec: Any
+    used_this_sec: int
+    this_sec: Any
+    def __init__(self, traces_per_sec: int = ...) -> None: ...
+    def take(self): ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/local/sampler.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+
+from ...exceptions.exceptions import InvalidSamplingManifestError as InvalidSamplingManifestError
+from .sampling_rule import SamplingRule as SamplingRule
+
+local_sampling_rule: Any
+SUPPORTED_RULE_VERSION: Any
+
+class LocalSampler:
+    def __init__(self, rules=...) -> None: ...
+    def should_trace(self, sampling_req: Any | None = ...): ...
+    def load_local_rules(self, rules) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/local/sampling_rule.pyi

@@ -0,0 +1,28 @@
+from aws_xray_sdk.core.utils.search_pattern import wildcard_match as wildcard_match
+
+from ...exceptions.exceptions import InvalidSamplingManifestError as InvalidSamplingManifestError
+from .reservoir import Reservoir as Reservoir
+
+class SamplingRule:
+    FIXED_TARGET: str
+    RATE: str
+    HOST: str
+    METHOD: str
+    PATH: str
+    SERVICE_NAME: str
+    def __init__(self, rule_dict, version: int = ..., default: bool = ...) -> None: ...
+    def applies(self, host, method, path): ...
+    @property
+    def fixed_target(self): ...
+    @property
+    def rate(self): ...
+    @property
+    def host(self): ...
+    @property
+    def method(self): ...
+    @property
+    def path(self): ...
+    @property
+    def reservoir(self): ...
+    @property
+    def version(self): ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/reservoir.pyi

@@ -0,0 +1,15 @@
+from enum import Enum
+
+class Reservoir:
+    def __init__(self) -> None: ...
+    def borrow_or_take(self, now, can_borrow): ...
+    def load_quota(self, quota, TTL, interval) -> None: ...
+    @property
+    def quota(self): ...
+    @property
+    def TTL(self): ...
+
+class ReservoirDecision(Enum):
+    TAKE: str
+    BORROW: str
+    NO: str

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_cache.pyi

@@ -0,0 +1,17 @@
+from typing import Any
+
+TTL: Any
+
+class RuleCache:
+    def __init__(self) -> None: ...
+    def get_matched_rule(self, sampling_req, now): ...
+    def load_rules(self, rules) -> None: ...
+    def load_targets(self, targets_dict) -> None: ...
+    @property
+    def rules(self): ...
+    @rules.setter
+    def rules(self, v) -> None: ...
+    @property
+    def last_updated(self): ...
+    @last_updated.setter
+    def last_updated(self, v) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/rule_poller.pyi

@@ -0,0 +1,9 @@
+from typing import Any
+
+log: Any
+DEFAULT_INTERVAL: Any
+
+class RulePoller:
+    def __init__(self, cache, connector) -> None: ...
+    def start(self) -> None: ...
+    def wake_up(self) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampler.pyi

@@ -0,0 +1,21 @@
+from typing import Any
+
+from .connector import ServiceConnector as ServiceConnector
+from .local.sampler import LocalSampler as LocalSampler
+from .reservoir import ReservoirDecision as ReservoirDecision
+from .rule_cache import RuleCache as RuleCache
+from .rule_poller import RulePoller as RulePoller
+from .target_poller import TargetPoller as TargetPoller
+
+log: Any
+
+class DefaultSampler:
+    def __init__(self) -> None: ...
+    def start(self) -> None: ...
+    def should_trace(self, sampling_req: Any | None = ...): ...
+    def load_local_rules(self, rules) -> None: ...
+    def load_settings(self, daemon_config, context, origin: Any | None = ...) -> None: ...
+    @property
+    def xray_client(self): ...
+    @xray_client.setter
+    def xray_client(self, v) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/sampling_rule.pyi

@@ -0,0 +1,47 @@
+from typing import Any
+
+from ..utils.search_pattern import wildcard_match as wildcard_match
+from .reservoir import Reservoir as Reservoir
+
+class SamplingRule:
+    def __init__(
+        self,
+        name,
+        priority,
+        rate,
+        reservoir_size,
+        host: Any | None = ...,
+        method: Any | None = ...,
+        path: Any | None = ...,
+        service: Any | None = ...,
+        service_type: Any | None = ...,
+    ) -> None: ...
+    def match(self, sampling_req): ...
+    def is_default(self): ...
+    def snapshot_statistics(self): ...
+    def merge(self, rule) -> None: ...
+    def ever_matched(self): ...
+    def time_to_report(self): ...
+    def increment_request_count(self) -> None: ...
+    def increment_borrow_count(self) -> None: ...
+    def increment_sampled_count(self) -> None: ...
+    @property
+    def rate(self): ...
+    @rate.setter
+    def rate(self, v) -> None: ...
+    @property
+    def name(self): ...
+    @property
+    def priority(self): ...
+    @property
+    def reservoir(self): ...
+    @reservoir.setter
+    def reservoir(self, v) -> None: ...
+    @property
+    def can_borrow(self): ...
+    @property
+    def request_count(self): ...
+    @property
+    def borrow_count(self): ...
+    @property
+    def sampled_count(self): ...

stubs/aws-xray-sdk/aws_xray_sdk/core/sampling/target_poller.pyi

@@ -0,0 +1,7 @@
+from typing import Any
+
+log: Any
+
+class TargetPoller:
+    def __init__(self, cache, rule_poller, connector) -> None: ...
+    def start(self) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/core/streaming/default_streaming.pyi

@@ -0,0 +1,8 @@
+class DefaultStreaming:
+    def __init__(self, streaming_threshold: int = ...) -> None: ...
+    def is_eligible(self, segment): ...
+    def stream(self, entity, callback) -> None: ...
+    @property
+    def streaming_threshold(self): ...
+    @streaming_threshold.setter
+    def streaming_threshold(self, value) -> None: ...

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

@@ -0,0 +1,9 @@
+from typing import Any
+
+class AtomicCounter:
+    value: Any
+    def __init__(self, initial: int = ...) -> None: ...
+    def increment(self, num: int = ...): ...
+    def decrement(self, num: int = ...): ...
+    def get_current(self): ...
+    def reset(self): ...

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

@@ -0,0 +1,9 @@
+from typing import Any
+
+PY2: Any
+PY35: Any
+annotation_value_types: Any
+string_types = str
+
+def is_classmethod(func): ...
+def is_instance_method(parent_class, func_name, func): ...

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

@@ -0,0 +1,5 @@
+from typing import Any
+
+log: Any
+
+def metadata_to_dict(obj): ...

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

@@ -0,0 +1 @@
+def wildcard_match(pattern, text, case_insensitive: bool = ...): ...

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

@@ -0,0 +1,3 @@
+from typing import Any
+
+def get_stacktrace(limit: Any | None = ...): ...

stubs/aws-xray-sdk/aws_xray_sdk/sdk_config.pyi

@@ -0,0 +1,11 @@
+from logging import Logger
+
+log: Logger
+
+class SDKConfig:
+    XRAY_ENABLED_KEY: str
+    DISABLED_ENTITY_NAME: str
+    @classmethod
+    def sdk_enabled(cls): ...
+    @classmethod
+    def set_sdk_enabled(cls, value) -> None: ...

stubs/aws-xray-sdk/aws_xray_sdk/version.pyi

@@ -0,0 +1 @@
+VERSION: str