Skip to content

Commit 017aed4

Browse files
donBarbosdonBarbos
authored
[aws-xray-sdk] Annotate more (#15004)
1 parent 4239e94 commit 017aed4

23 files changed

+161
-101
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
aws_xray_sdk.core.models.subsegment.subsegment_decorator
22
aws_xray_sdk.core.sampling.connector.ServiceConnector.fetch_sampling_rules
33

4+
# Inconsistency because `context_missing` param can be passed in *args or **kwargs:
5+
aws_xray_sdk.core.async_context.AsyncContext.__init__
6+
47
# We can not import 3rd-party libraries in teststubs runtime,
58
# but we can use Protocol to replace this types:
69
aws_xray_sdk.ext.aiobotocore

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

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,23 @@
1+
from asyncio.events import AbstractEventLoop
2+
from asyncio.tasks import Task, _TaskCompatibleCoro
3+
from typing import Any, TypeVar
4+
15
from .context import Context as _Context
26

7+
_T_co = TypeVar("_T_co", covariant=True)
8+
39
class AsyncContext(_Context):
4-
def __init__(self, *args, loop=None, use_task_factory: bool = True, **kwargs) -> None: ...
10+
def __init__(
11+
self, context_missing: str = "LOG_ERROR", loop: AbstractEventLoop | None = None, use_task_factory: bool = True
12+
) -> None: ...
513
def clear_trace_entities(self) -> None: ...
614

715
class TaskLocalStorage:
8-
def __init__(self, loop=None) -> None: ...
9-
def __setattr__(self, name: str, value) -> None: ...
10-
def __getattribute__(self, item: str): ...
16+
def __init__(self, loop: AbstractEventLoop | None = None) -> None: ...
17+
# Sets unknown items on the current task's context attribute
18+
def __setattr__(self, name: str, value: Any) -> None: ...
19+
# Returns unknown items from the current tasks context attribute
20+
def __getattribute__(self, item: str) -> Any | None: ...
1121
def clear(self) -> None: ...
1222

13-
def task_factory(loop, coro): ...
23+
def task_factory(loop: AbstractEventLoop | None, coro: _TaskCompatibleCoro[_T_co]) -> Task[_T_co]: ...

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

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,43 @@
1+
from _typeshed import Incomplete
2+
from collections.abc import Awaitable, Callable, Iterable, Mapping
13
from types import TracebackType
4+
from typing import TypeVar
25

3-
from .models.segment import SegmentContextManager
4-
from .models.subsegment import SubsegmentContextManager
6+
from .models.dummy_entities import DummySegment, DummySubsegment
7+
from .models.segment import Segment, SegmentContextManager
8+
from .models.subsegment import Subsegment, SubsegmentContextManager
59
from .recorder import AWSXRayRecorder
610

11+
_T = TypeVar("_T")
12+
713
class AsyncSegmentContextManager(SegmentContextManager):
8-
async def __aenter__(self): ...
14+
async def __aenter__(self) -> DummySegment | Segment: ...
915
async def __aexit__(
1016
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
1117
) -> None: ...
1218

1319
class AsyncSubsegmentContextManager(SubsegmentContextManager):
14-
async def __call__(self, wrapped, instance, args, kwargs): ...
15-
async def __aenter__(self): ...
20+
async def __call__(
21+
self, wrapped: Callable[..., Awaitable[_T]], instance, args: Iterable[Incomplete], kwargs: Mapping[str, Incomplete]
22+
) -> _T: ...
23+
async def __aenter__(self) -> DummySubsegment | Subsegment | None: ...
1624
async def __aexit__(
1725
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
1826
) -> None: ...
1927

2028
class AsyncAWSXRayRecorder(AWSXRayRecorder):
21-
def capture_async(self, name=None): ...
22-
def in_segment_async(self, name=None, **segment_kwargs): ...
23-
def in_subsegment_async(self, name=None, **subsegment_kwargs): ...
24-
async def record_subsegment_async(self, wrapped, instance, args, kwargs, name, namespace, meta_processor): ...
29+
def capture_async(self, name: str | None = None) -> AsyncSubsegmentContextManager: ...
30+
def in_segment_async(
31+
self, name: str | None = None, *, traceid: str | None = None, parent_id: str | None = None, sampling: bool | None = None
32+
) -> AsyncSegmentContextManager: ...
33+
def in_subsegment_async(self, name: str | None = None, *, namespace: str = "local") -> AsyncSubsegmentContextManager: ...
34+
async def record_subsegment_async(
35+
self,
36+
wrapped: Callable[..., Awaitable[_T]],
37+
instance,
38+
args: Iterable[Incomplete],
39+
kwargs: Mapping[str, Incomplete],
40+
name: str,
41+
namespace: str,
42+
meta_processor: Callable[..., object] | None,
43+
) -> _T: ...

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,27 @@
11
import time
22
from logging import Logger
3-
from typing import Any
3+
from typing import Final
44

55
from .models.entity import Entity
66
from .models.segment import Segment
77
from .models.subsegment import Subsegment
88

99
log: Logger
10-
SUPPORTED_CONTEXT_MISSING: Any
11-
MISSING_SEGMENT_MSG: str
12-
CXT_MISSING_STRATEGY_KEY: str
10+
MISSING_SEGMENT_MSG: Final[str]
11+
SUPPORTED_CONTEXT_MISSING: Final = ("RUNTIME_ERROR", "LOG_ERROR", "IGNORE_ERROR")
12+
CXT_MISSING_STRATEGY_KEY: Final = "AWS_XRAY_CONTEXT_MISSING"
1313

1414
class Context:
1515
def __init__(self, context_missing: str = "LOG_ERROR") -> None: ...
1616
def put_segment(self, segment: Segment) -> None: ...
1717
def end_segment(self, end_time: time.struct_time | None = None) -> None: ...
1818
def put_subsegment(self, subsegment: Subsegment) -> None: ...
19-
def end_subsegment(self, end_time: time.struct_time | None = None): ...
20-
def get_trace_entity(self): ...
19+
def end_subsegment(self, end_time: time.struct_time | None = None) -> bool: ...
20+
def get_trace_entity(self) -> Entity: ...
2121
def set_trace_entity(self, trace_entity: Entity) -> None: ...
2222
def clear_trace_entities(self) -> None: ...
2323
def handle_context_missing(self) -> None: ...
2424
@property
25-
def context_missing(self): ...
25+
def context_missing(self) -> str: ...
2626
@context_missing.setter
2727
def context_missing(self, value: str) -> None: ...

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
1-
DAEMON_ADDRESS_KEY: str
2-
DEFAULT_ADDRESS: str
1+
from typing import Final
2+
3+
DAEMON_ADDRESS_KEY: Final = "AWS_XRAY_DAEMON_ADDRESS"
4+
DEFAULT_ADDRESS: Final = "127.0.0.1:2000"
35

46
class DaemonConfig:
5-
def __init__(self, daemon_address="127.0.0.1:2000") -> None: ...
7+
def __init__(self, daemon_address: str | None = "127.0.0.1:2000") -> None: ...
68
@property
7-
def udp_ip(self): ...
9+
def udp_ip(self) -> str: ...
810
@property
9-
def udp_port(self): ...
11+
def udp_port(self) -> int: ...
1012
@property
11-
def tcp_ip(self): ...
13+
def tcp_ip(self) -> str: ...
1214
@property
13-
def tcp_port(self): ...
15+
def tcp_port(self) -> int: ...

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,6 @@ class UDPEmitter:
1313
def send_entity(self, entity: Entity) -> None: ...
1414
def set_daemon_address(self, address: str | None) -> None: ...
1515
@property
16-
def ip(self): ...
16+
def ip(self) -> str: ...
1717
@property
18-
def port(self): ...
18+
def port(self) -> int: ...

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

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,19 @@
11
from logging import Logger
2+
from typing import Final
23

34
from .context import Context
45

56
log: Logger
6-
LAMBDA_TRACE_HEADER_KEY: str
7-
LAMBDA_TASK_ROOT_KEY: str
8-
TOUCH_FILE_DIR: str
9-
TOUCH_FILE_PATH: str
7+
LAMBDA_TRACE_HEADER_KEY: Final = "_X_AMZN_TRACE_ID"
8+
LAMBDA_TASK_ROOT_KEY: Final = "LAMBDA_TASK_ROOT"
9+
TOUCH_FILE_DIR: Final = "/tmp/.aws-xray/"
10+
TOUCH_FILE_PATH: Final = "/tmp/.aws-xray/initialized"
1011

11-
def check_in_lambda(): ...
12+
def check_in_lambda() -> LambdaContext | None: ...
1213

1314
class LambdaContext(Context):
1415
def __init__(self) -> None: ...
15-
def put_segment(self, segment) -> None: ...
16-
def end_segment(self, end_time=None) -> None: ...
17-
def put_subsegment(self, subsegment) -> None: ...
18-
def get_trace_entity(self): ...
19-
@property
16+
@property # type: ignore[override]
2017
def context_missing(self) -> None: ...
2118
@context_missing.setter
22-
def context_missing(self, value) -> None: ...
23-
def handle_context_missing(self) -> None: ...
19+
def context_missing(self, value: str) -> None: ...

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ from .subsegment import Subsegment
77
from .throwable import Throwable
88

99
log: Logger
10-
ORIGIN_TRACE_HEADER_ATTR_KEY: Final[str]
10+
ORIGIN_TRACE_HEADER_ATTR_KEY: Final = "_origin_trace_header"
1111

1212
class Entity:
1313
id: str

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ from typing import Final
22

33
from .segment import Segment
44

5-
MUTATION_UNSUPPORTED_MESSAGE: Final[str]
5+
MUTATION_UNSUPPORTED_MESSAGE: Final = "FacadeSegments cannot be mutated."
66

77
class FacadeSegment(Segment):
88
initializing: bool

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
from typing import Final
22

3-
URL: Final[str]
4-
METHOD: Final[str]
5-
USER_AGENT: Final[str]
6-
CLIENT_IP: Final[str]
7-
X_FORWARDED_FOR: Final[str]
8-
STATUS: Final[str]
9-
CONTENT_LENGTH: Final[str]
10-
XRAY_HEADER: Final[str]
11-
ALT_XRAY_HEADER: Final[str]
3+
URL: Final = "url"
4+
METHOD: Final = "method"
5+
USER_AGENT: Final = "user_agent"
6+
CLIENT_IP: Final = "client_ip"
7+
X_FORWARDED_FOR: Final = "x_forwarded_for"
8+
STATUS: Final = "status"
9+
CONTENT_LENGTH: Final = "content_length"
10+
XRAY_HEADER: Final = "X-Amzn-Trace-Id"
11+
ALT_XRAY_HEADER: Final = "HTTP_X_AMZN_TRACE_ID"
1212
request_keys: tuple[str, ...]
1313
response_keys: tuple[str, ...]

0 commit comments

Comments
 (0)