diff --git a/aws_lambda_powertools/tracing/base.py b/aws_lambda_powertools/tracing/base.py index 6ea58da6b5a..74b146ad6e8 100644 --- a/aws_lambda_powertools/tracing/base.py +++ b/aws_lambda_powertools/tracing/base.py @@ -1,15 +1,19 @@ +from __future__ import annotations + import abc -import numbers -import traceback from contextlib import contextmanager -from typing import Any, Generator, List, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Generator, Sequence + +if TYPE_CHECKING: + import numbers + import traceback class BaseSegment(abc.ABC): """Holds common properties and methods on segment and subsegment.""" @abc.abstractmethod - def close(self, end_time: Optional[int] = None): + def close(self, end_time: int | None = None): """Close the trace entity by setting `end_time` and flip the in progress flag to False. @@ -28,7 +32,7 @@ def remove_subsegment(self, subsegment: Any): """Remove input subsegment from child subsegments.""" @abc.abstractmethod - def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> None: + def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None: """Annotate segment or subsegment with a key-value pair. Note: Annotations will be indexed for later search query. @@ -37,7 +41,7 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> N ---------- key: str Metadata key - value: Union[str, numbers.Number, bool] + value: str | numbers.Number | bool Annotation value """ @@ -52,19 +56,19 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None Metadata key value: Any Any object that can be serialized into a JSON string - namespace: Set[str] + namespace: set[str] Metadata namespace, by default 'default' """ @abc.abstractmethod - def add_exception(self, exception: BaseException, stack: List[traceback.StackSummary], remote: bool = False): + def add_exception(self, exception: BaseException, stack: list[traceback.StackSummary], remote: bool = False): """Add an exception to trace entities. Parameters ---------- exception: Exception Caught exception - stack: List[traceback.StackSummary] + stack: list[traceback.StackSummary] List of traceback summaries Output from `traceback.extract_stack()`. @@ -83,7 +87,7 @@ def in_subsegment(self, name=None, **kwargs) -> Generator[BaseSegment, None, Non ---------- name: str Subsegment name - kwargs: Optional[dict] + kwargs: dict | None Optional parameters to be propagated to segment """ @@ -96,12 +100,12 @@ def in_subsegment_async(self, name=None, **kwargs) -> Generator[BaseSegment, Non ---------- name: str Subsegment name - kwargs: Optional[dict] + kwargs: dict | None Optional parameters to be propagated to segment """ @abc.abstractmethod - def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> None: + def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None: """Annotate current active trace entity with a key-value pair. Note: Annotations will be indexed for later search query. @@ -110,7 +114,7 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> N ---------- key: str Metadata key - value: Union[str, numbers.Number, bool] + value: str | numbers.Number | bool Annotation value """ @@ -126,7 +130,7 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None Metadata key value: Any Any object that can be serialized into a JSON string - namespace: Set[str] + namespace: set[str] Metadata namespace, by default 'default' """ @@ -136,7 +140,7 @@ def patch(self, modules: Sequence[str]) -> None: Parameters ---------- - modules: Set[str] + modules: set[str] Set of modules to be patched """ diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index a79ac4ec738..0e2bbeb4a3e 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/aws_lambda_powertools/tracing/tracer.py @@ -1,11 +1,12 @@ +from __future__ import annotations + import contextlib import copy import functools import inspect import logging -import numbers import os -from typing import Any, Callable, Dict, List, Optional, Sequence, Union, cast, overload +from typing import TYPE_CHECKING, Any, Callable, Sequence, cast, overload from aws_lambda_powertools.shared import constants from aws_lambda_powertools.shared.functions import ( @@ -15,7 +16,11 @@ ) from aws_lambda_powertools.shared.lazy_import import LazyLoader from aws_lambda_powertools.shared.types import AnyCallableT -from aws_lambda_powertools.tracing.base import BaseProvider, BaseSegment + +if TYPE_CHECKING: + import numbers + + from aws_lambda_powertools.tracing.base import BaseProvider, BaseSegment is_cold_start = True logger = logging.getLogger(__name__) @@ -57,7 +62,7 @@ class Tracer: disabled: bool Flag to explicitly disable tracing, useful when running/testing locally `Env POWERTOOLS_TRACE_DISABLED="true"` - patch_modules: Optional[Sequence[str]] + patch_modules: Sequence[str] | None Tuple of modules supported by tracing provider to patch, by default all modules are patched provider: BaseProvider Tracing provider, by default it is aws_xray_sdk.core.xray_recorder @@ -75,13 +80,13 @@ class Tracer: tracer = Tracer(service="greeting") @tracer.capture_method - def greeting(name: str) -> Dict: + def greeting(name: str) -> dict: return { "name": name } @tracer.capture_lambda_handler - def handler(event: dict, context: Any) -> Dict: + def handler(event: dict, context: Any) -> dict: print("Received event from Lambda...") response = greeting(name="Heitor") return response @@ -92,7 +97,7 @@ def handler(event: dict, context: Any) -> Dict: tracer = Tracer(service="booking") @tracer.capture_method - def confirm_booking(booking_id: str) -> Dict: + def confirm_booking(booking_id: str) -> dict: resp = add_confirmation(booking_id) tracer.put_annotation("BookingConfirmation", resp["requestId"]) @@ -101,7 +106,7 @@ def confirm_booking(booking_id: str) -> Dict: return resp @tracer.capture_lambda_handler - def handler(event: dict, context: Any) -> Dict: + def handler(event: dict, context: Any) -> dict: print("Received event from Lambda...") booking_id = event.get("booking_id") response = confirm_booking(booking_id=booking_id) @@ -114,7 +119,7 @@ def handler(event: dict, context: Any) -> Dict: tracer = Tracer() @tracer.capture_lambda_handler - def handler(event: dict, context: Any) -> Dict: + def handler(event: dict, context: Any) -> dict: print("Received event from Lambda...") response = greeting(name="Lessa") return response @@ -126,7 +131,7 @@ def handler(event: dict, context: Any) -> Dict: tracer = Tracer() @tracer.capture_lambda_handler - def handler(event: dict, context: Any) -> Dict: + def handler(event: dict, context: Any) -> dict: ... # utils.py @@ -139,7 +144,7 @@ def handler(event: dict, context: Any) -> Dict: * Async handler not supported """ - _default_config: Dict[str, Any] = { + _default_config: dict[str, Any] = { "service": "", "disabled": False, "auto_patch": True, @@ -150,11 +155,11 @@ def handler(event: dict, context: Any) -> Dict: def __init__( self, - service: Optional[str] = None, - disabled: Optional[bool] = None, - auto_patch: Optional[bool] = None, - patch_modules: Optional[Sequence[str]] = None, - provider: Optional[BaseProvider] = None, + service: str | None = None, + disabled: bool | None = None, + auto_patch: bool | None = None, + patch_modules: Sequence[str] | None = None, + provider: BaseProvider | None = None, ): self.__build_config( service=service, @@ -177,14 +182,14 @@ def __init__( if self._is_xray_provider(): self._disable_xray_trace_batching() - def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]): + def put_annotation(self, key: str, value: str | numbers.Number | bool): """Adds annotation to existing segment or subsegment Parameters ---------- key : str Annotation key - value : Union[str, numbers.Number, bool] + value : str | numbers.Number | bool Value for annotation Example @@ -201,7 +206,7 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]): logger.debug(f"Annotating on key '{key}' with '{value}'") self.provider.put_annotation(key=key, value=value) - def put_metadata(self, key: str, value: Any, namespace: Optional[str] = None): + def put_metadata(self, key: str, value: Any, namespace: str | None = None): """Adds metadata to existing segment or subsegment Parameters @@ -229,14 +234,14 @@ def put_metadata(self, key: str, value: Any, namespace: Optional[str] = None): logger.debug(f"Adding metadata on key '{key}' with '{value}' at namespace '{namespace}'") self.provider.put_metadata(key=key, value=value, namespace=namespace) - def patch(self, modules: Optional[Sequence[str]] = None): + def patch(self, modules: Sequence[str] | None = None): """Patch modules for instrumentation. Patches all supported modules by default if none are given. Parameters ---------- - modules : Optional[Sequence[str]] + modules : Sequence[str] | None List of modules to be patched, optional by default """ if self.disabled: @@ -250,9 +255,9 @@ def patch(self, modules: Optional[Sequence[str]] = None): def capture_lambda_handler( self, - lambda_handler: Union[Callable[[Dict, Any], Any], Optional[Callable[[Dict, Any, Optional[Dict]], Any]]] = None, - capture_response: Optional[bool] = None, - capture_error: Optional[bool] = None, + lambda_handler: Callable[[dict, Any], Any] | Callable[[dict, Any, dict | None], Any] | None = None, + capture_response: bool | None = None, + capture_error: bool | None = None, ): """Decorator to create subsegment for lambda handlers @@ -349,21 +354,21 @@ def decorate(event, context, **kwargs): # see #465 @overload - def capture_method(self, method: "AnyCallableT") -> "AnyCallableT": ... # pragma: no cover + def capture_method(self, method: AnyCallableT) -> AnyCallableT: ... # pragma: no cover @overload def capture_method( self, method: None = None, - capture_response: Optional[bool] = None, - capture_error: Optional[bool] = None, - ) -> Callable[["AnyCallableT"], "AnyCallableT"]: ... # pragma: no cover + capture_response: bool | None = None, + capture_error: bool | None = None, + ) -> Callable[[AnyCallableT], AnyCallableT]: ... # pragma: no cover def capture_method( self, - method: Optional[AnyCallableT] = None, - capture_response: Optional[bool] = None, - capture_error: Optional[bool] = None, + method: AnyCallableT | None = None, + capture_response: bool | None = None, + capture_error: bool | None = None, ) -> AnyCallableT: """Decorator to create subsegment for arbitrary functions @@ -402,7 +407,7 @@ def some_function() tracer = Tracer(service="booking") @tracer.capture_method - async def confirm_booking(booking_id: str) -> Dict: + async def confirm_booking(booking_id: str) -> dict: resp = call_to_booking_service() tracer.put_annotation("BookingConfirmation", resp["requestId"]) @@ -410,7 +415,7 @@ async def confirm_booking(booking_id: str) -> Dict: return resp - def lambda_handler(event: dict, context: Any) -> Dict: + def lambda_handler(event: dict, context: Any) -> dict: booking_id = event.get("booking_id") asyncio.run(confirm_booking(booking_id=booking_id)) @@ -425,7 +430,7 @@ def bookings_generator(booking_id): yield resp[0] yield resp[1] - def lambda_handler(event: dict, context: Any) -> Dict: + def lambda_handler(event: dict, context: Any) -> dict: gen = bookings_generator(booking_id=booking_id) result = list(gen) @@ -441,7 +446,7 @@ def booking_actions(booking_id): yield "example result" cleanup_stuff() - def lambda_handler(event: dict, context: Any) -> Dict: + def lambda_handler(event: dict, context: Any) -> dict: booking_id = event.get("booking_id") with booking_actions(booking_id=booking_id) as booking: @@ -569,9 +574,9 @@ async def async_tasks(): def _decorate_async_function( self, method: Callable, - capture_response: Optional[Union[bool, str]] = None, - capture_error: Optional[Union[bool, str]] = None, - method_name: Optional[str] = None, + capture_response: bool | str | None = None, + capture_error: bool | str | None = None, + method_name: str | None = None, ): @functools.wraps(method) async def decorate(*args, **kwargs): @@ -602,9 +607,9 @@ async def decorate(*args, **kwargs): def _decorate_generator_function( self, method: Callable, - capture_response: Optional[Union[bool, str]] = None, - capture_error: Optional[Union[bool, str]] = None, - method_name: Optional[str] = None, + capture_response: bool | str | None = None, + capture_error: bool | str | None = None, + method_name: str | None = None, ): @functools.wraps(method) def decorate(*args, **kwargs): @@ -635,9 +640,9 @@ def decorate(*args, **kwargs): def _decorate_generator_function_with_context_manager( self, method: Callable, - capture_response: Optional[Union[bool, str]] = None, - capture_error: Optional[Union[bool, str]] = None, - method_name: Optional[str] = None, + capture_response: bool | str | None = None, + capture_error: bool | str | None = None, + method_name: str | None = None, ): @functools.wraps(method) @contextlib.contextmanager @@ -669,9 +674,9 @@ def decorate(*args, **kwargs): def _decorate_sync_function( self, method: AnyCallableT, - capture_response: Optional[Union[bool, str]] = None, - capture_error: Optional[Union[bool, str]] = None, - method_name: Optional[str] = None, + capture_response: bool | str | None = None, + capture_error: bool | str | None = None, + method_name: str | None = None, ) -> AnyCallableT: @functools.wraps(method) def decorate(*args, **kwargs): @@ -701,10 +706,10 @@ def decorate(*args, **kwargs): def _add_response_as_metadata( self, - method_name: Optional[str] = None, - data: Optional[Any] = None, - subsegment: Optional[BaseSegment] = None, - capture_response: Optional[Union[bool, str]] = None, + method_name: str | None = None, + data: Any | None = None, + subsegment: BaseSegment | None = None, + capture_response: bool | str | None = None, ): """Add response as metadata for given subsegment @@ -729,7 +734,7 @@ def _add_full_exception_as_metadata( method_name: str, error: Exception, subsegment: BaseSegment, - capture_error: Optional[bool] = None, + capture_error: bool | None = None, ): """Add full exception object as metadata for given subsegment @@ -756,7 +761,7 @@ def _disable_tracer_provider(): aws_xray_sdk.global_sdk_config.set_sdk_enabled(False) @staticmethod - def _is_tracer_disabled() -> Union[bool, str]: + def _is_tracer_disabled() -> bool | str: """Detects whether trace has been disabled Tracing is automatically disabled in the following conditions: @@ -767,7 +772,7 @@ def _is_tracer_disabled() -> Union[bool, str]: Returns ------- - Union[bool, str] + bool | str """ logger.debug("Verifying whether Tracing has been disabled") is_lambda_env = os.getenv(constants.LAMBDA_TASK_ROOT_ENV) @@ -787,11 +792,11 @@ def _is_tracer_disabled() -> Union[bool, str]: def __build_config( self, - service: Optional[str] = None, - disabled: Optional[bool] = None, - auto_patch: Optional[bool] = None, - patch_modules: Optional[Sequence[str]] = None, - provider: Optional[BaseProvider] = None, + service: str | None = None, + disabled: bool | None = None, + auto_patch: bool | None = None, + patch_modules: Sequence[str] | None = None, + provider: BaseProvider | None = None, ): """Populates Tracer config for new and existing initializations""" is_disabled = disabled if disabled is not None else self._is_tracer_disabled() @@ -833,7 +838,7 @@ def _disable_xray_trace_batching(self): def _is_xray_provider(self): return "aws_xray_sdk" in self.provider.__module__ - def ignore_endpoint(self, hostname: Optional[str] = None, urls: Optional[List[str]] = None): + def ignore_endpoint(self, hostname: str | None = None, urls: list[str] | None = None): """If you want to ignore certain httplib requests you can do so based on the hostname or URL that is being requested. @@ -847,7 +852,7 @@ def ignore_endpoint(self, hostname: Optional[str] = None, urls: Optional[List[st ---------- hostname : Optional, str The hostname is matched using the Python fnmatch library which does Unix glob style matching. - urls: Optional, List[str] + urls: Optional, list[str] List of urls to ignore. Example `tracer.ignore_endpoint(urls=["/ignored-url"])` """ if not self._is_xray_provider(): diff --git a/tests/e2e/parser/handlers/handler_with_union_tag.py b/tests/e2e/parser/handlers/handler_with_union_tag.py index 9dbc81662b4..5240a902db0 100644 --- a/tests/e2e/parser/handlers/handler_with_union_tag.py +++ b/tests/e2e/parser/handlers/handler_with_union_tag.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Literal +from typing import Literal, Union from pydantic import BaseModel, Field from typing_extensions import Annotated @@ -25,7 +25,7 @@ class PartialFailureCallback(BaseModel): error_msg: str -OrderCallback = Annotated[SuccessCallback | ErrorCallback | PartialFailureCallback, Field(discriminator="status")] +OrderCallback = Annotated[Union[SuccessCallback, ErrorCallback, PartialFailureCallback], Field(discriminator="status")] @event_parser