From 3ed41f8f88338a90e2d2325bd6048831ee9d0b50 Mon Sep 17 00:00:00 2001
From: Eric Nielsen <4120606+ericbn@users.noreply.github.com>
Date: Tue, 6 Aug 2024 12:26:33 -0500
Subject: [PATCH 1/2] refactor(typing): reduce
 aws_lambda_powertools.shared.types usage

As discussed in #4607. This simplifies linting and refactoring so we can
introduce

    from __future__ import annotations

to all files, which is the plan as the next step.
---
 .../event_handler/api_gateway.py              |  2 +-
 .../event_handler/middlewares/base.py         |  3 +--
 .../event_handler/openapi/models.py           |  4 ++--
 .../event_handler/openapi/params.py           |  4 ++--
 .../event_handler/openapi/types.py            |  4 ++--
 aws_lambda_powertools/logging/types.py        |  4 ++--
 aws_lambda_powertools/metrics/functions.py    |  2 +-
 .../metrics/provider/cloudwatch_emf/types.py  |  4 +++-
 aws_lambda_powertools/metrics/types.py        |  4 +++-
 aws_lambda_powertools/shared/cookies.py       |  4 +---
 aws_lambda_powertools/shared/types.py         | 22 +------------------
 .../utilities/batch/types.py                  |  4 +---
 .../data_classes/s3_batch_operation_event.py  |  3 +--
 .../data_classes/secrets_manager_event.py     |  3 ++-
 .../utilities/idempotency/hook.py             |  3 +--
 .../idempotency/persistence/redis.py          |  3 +--
 .../utilities/parameters/ssm.py               |  3 +--
 .../utilities/parameters/types.py             |  2 +-
 .../utilities/parser/models/cloudwatch.py     |  4 +---
 .../utilities/parser/types.py                 |  4 +---
 .../utilities/streaming/s3_object.py          |  2 +-
 .../assert_bedrock_agent_response_module.py   |  3 ++-
 .../assert_async_graphql_response_module.py   |  3 +--
 .../src/assert_graphql_response_module.py     |  3 +--
 .../src/async_resolvers.py                    |  3 +--
 .../src/custom_models.py                      |  3 +--
 .../getting_started_graphql_api_resolver.py   |  3 +--
 .../src/graphql_transformer_merchant_info.py  |  3 +--
 .../graphql_transformer_search_merchant.py    |  3 +--
 .../src/nested_mappings.py                    |  3 +--
 .../split_operation_append_context_module.py  |  3 +--
 .../src/split_operation_module.py             |  3 +--
 .../src/validating_headers.py                 |  2 +-
 .../event_handler_rest/src/validating_path.py |  2 +-
 .../src/validating_payload_subset.py          |  2 +-
 .../src/validating_query_strings.py           |  2 +-
 .../src/working_with_headers_multi_value.py   |  3 ++-
 .../src/working_with_multi_query_values.py    |  3 ++-
 examples/parser/src/multiple_model_parsing.py |  2 +-
 .../event_handler/test_bedrock_agent.py       |  3 ++-
 .../event_handler/test_openapi_params.py      |  2 +-
 .../test_openapi_schema_pydantic_v1.py        |  4 ++--
 .../test_openapi_schema_pydantic_v2.py        |  4 ++--
 .../test_openapi_validation_middleware.py     |  2 +-
 tests/functional/parser/test_parser.py        |  2 +-
 .../parser/test_parser_performance.py         |  4 ++--
 46 files changed, 62 insertions(+), 96 deletions(-)

diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py
index a20154b4bbf..ad6f124c385 100644
--- a/aws_lambda_powertools/event_handler/api_gateway.py
+++ b/aws_lambda_powertools/event_handler/api_gateway.py
@@ -17,6 +17,7 @@
     Dict,
     Generic,
     List,
+    Literal,
     Mapping,
     Match,
     Optional,
@@ -47,7 +48,6 @@
 from aws_lambda_powertools.shared.cookies import Cookie
 from aws_lambda_powertools.shared.functions import powertools_dev_is_set
 from aws_lambda_powertools.shared.json_encoder import Encoder
-from aws_lambda_powertools.shared.types import Literal
 from aws_lambda_powertools.utilities.data_classes import (
     ALBEvent,
     APIGatewayProxyEvent,
diff --git a/aws_lambda_powertools/event_handler/middlewares/base.py b/aws_lambda_powertools/event_handler/middlewares/base.py
index 342b033ec1f..700f02e8a30 100644
--- a/aws_lambda_powertools/event_handler/middlewares/base.py
+++ b/aws_lambda_powertools/event_handler/middlewares/base.py
@@ -1,9 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import Generic
+from typing import Generic, Protocol
 
 from aws_lambda_powertools.event_handler.api_gateway import Response
 from aws_lambda_powertools.event_handler.types import EventHandlerInstance
-from aws_lambda_powertools.shared.types import Protocol
 
 
 class NextMiddleware(Protocol):
diff --git a/aws_lambda_powertools/event_handler/openapi/models.py b/aws_lambda_powertools/event_handler/openapi/models.py
index 72d6c00bb60..57310fd3267 100644
--- a/aws_lambda_powertools/event_handler/openapi/models.py
+++ b/aws_lambda_powertools/event_handler/openapi/models.py
@@ -1,14 +1,14 @@
 from enum import Enum
-from typing import Any, Dict, List, Optional, Set, Union
+from typing import Any, Dict, List, Literal, Optional, Set, Union
 
 from pydantic import AnyUrl, BaseModel, Field
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler.openapi.compat import model_rebuild
 from aws_lambda_powertools.event_handler.openapi.constants import (
     MODEL_CONFIG_ALLOW,
     MODEL_CONFIG_IGNORE,
 )
-from aws_lambda_powertools.shared.types import Annotated, Literal
 
 """
 The code defines Pydantic models for the various OpenAPI objects like OpenAPI, PathItem, Operation, Parameter etc.
diff --git a/aws_lambda_powertools/event_handler/openapi/params.py b/aws_lambda_powertools/event_handler/openapi/params.py
index 62a23ba0dd6..4aa22882c3d 100644
--- a/aws_lambda_powertools/event_handler/openapi/params.py
+++ b/aws_lambda_powertools/event_handler/openapi/params.py
@@ -1,9 +1,10 @@
 import inspect
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type, Union
 
 from pydantic import BaseConfig
 from pydantic.fields import FieldInfo
+from typing_extensions import Annotated, get_args, get_origin
 
 from aws_lambda_powertools.event_handler import Response
 from aws_lambda_powertools.event_handler.openapi.compat import (
@@ -16,7 +17,6 @@
     get_annotation_from_field_info,
 )
 from aws_lambda_powertools.event_handler.openapi.types import CacheKey
-from aws_lambda_powertools.shared.types import Annotated, Literal, get_args, get_origin
 
 """
 This turns the low-level function signature into typed, validated Pydantic models for consumption.
diff --git a/aws_lambda_powertools/event_handler/openapi/types.py b/aws_lambda_powertools/event_handler/openapi/types.py
index 105d4cca2a0..33a6ad4df4e 100644
--- a/aws_lambda_powertools/event_handler/openapi/types.py
+++ b/aws_lambda_powertools/event_handler/openapi/types.py
@@ -1,8 +1,8 @@
 import types
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type, TypedDict, Union
 
-from aws_lambda_powertools.shared.types import NotRequired, TypedDict
+from typing_extensions import NotRequired
 
 if TYPE_CHECKING:
     from pydantic import BaseModel  # noqa: F401
diff --git a/aws_lambda_powertools/logging/types.py b/aws_lambda_powertools/logging/types.py
index eb2b39afe69..fcfec998ac2 100644
--- a/aws_lambda_powertools/logging/types.py
+++ b/aws_lambda_powertools/logging/types.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, TypedDict, Union
 
-from aws_lambda_powertools.shared.types import NotRequired, TypeAlias, TypedDict
+from typing_extensions import NotRequired, TypeAlias
 
 LogRecord: TypeAlias = Union[Dict[str, Any], "PowertoolsLogRecord"]
 LogStackTrace: TypeAlias = Union[Dict[str, Any], "PowertoolsStackTrace"]
diff --git a/aws_lambda_powertools/metrics/functions.py b/aws_lambda_powertools/metrics/functions.py
index ea8dc3603d1..6b00c608c36 100644
--- a/aws_lambda_powertools/metrics/functions.py
+++ b/aws_lambda_powertools/metrics/functions.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime
+from typing import List
 
 from aws_lambda_powertools.metrics.provider.cloudwatch_emf.exceptions import (
     MetricResolutionError,
@@ -8,7 +9,6 @@
 )
 from aws_lambda_powertools.metrics.provider.cloudwatch_emf.metric_properties import MetricResolution, MetricUnit
 from aws_lambda_powertools.shared import constants
-from aws_lambda_powertools.shared.types import List
 
 
 def extract_cloudwatch_metric_resolution_value(metric_resolutions: List, resolution: int | MetricResolution) -> int:
diff --git a/aws_lambda_powertools/metrics/provider/cloudwatch_emf/types.py b/aws_lambda_powertools/metrics/provider/cloudwatch_emf/types.py
index 359fdc4ee6c..669e931ff3e 100644
--- a/aws_lambda_powertools/metrics/provider/cloudwatch_emf/types.py
+++ b/aws_lambda_powertools/metrics/provider/cloudwatch_emf/types.py
@@ -1,4 +1,6 @@
-from aws_lambda_powertools.shared.types import List, NotRequired, TypedDict
+from typing import List, TypedDict
+
+from typing_extensions import NotRequired
 
 
 class CloudWatchEMFMetric(TypedDict):
diff --git a/aws_lambda_powertools/metrics/types.py b/aws_lambda_powertools/metrics/types.py
index d9eea6fe51e..3d29c0a2407 100644
--- a/aws_lambda_powertools/metrics/types.py
+++ b/aws_lambda_powertools/metrics/types.py
@@ -1,4 +1,6 @@
-from aws_lambda_powertools.shared.types import NotRequired, TypedDict
+from typing import TypedDict
+
+from typing_extensions import NotRequired
 
 
 class MetricNameUnitResolution(TypedDict):
diff --git a/aws_lambda_powertools/shared/cookies.py b/aws_lambda_powertools/shared/cookies.py
index 1b57d860201..944bcb5dc9f 100644
--- a/aws_lambda_powertools/shared/cookies.py
+++ b/aws_lambda_powertools/shared/cookies.py
@@ -1,9 +1,7 @@
 from datetime import datetime
 from enum import Enum
 from io import StringIO
-from typing import Optional
-
-from aws_lambda_powertools.shared.types import List
+from typing import List, Optional
 
 
 class SameSite(Enum):
diff --git a/aws_lambda_powertools/shared/types.py b/aws_lambda_powertools/shared/types.py
index d5014c4c467..e4e10192e55 100644
--- a/aws_lambda_powertools/shared/types.py
+++ b/aws_lambda_powertools/shared/types.py
@@ -1,25 +1,5 @@
-import sys
-from typing import Any, Callable, Dict, List, Literal, Protocol, TypedDict, TypeVar, Union
-
-if sys.version_info >= (3, 9):
-    from typing import Annotated
-else:
-    from typing_extensions import Annotated
-
-if sys.version_info >= (3, 11):
-    from typing import NotRequired
-else:
-    from typing_extensions import NotRequired
-
-# Even though `get_args` and `get_origin` were added in Python 3.8, they only handle Annotated correctly on 3.10.
-# So for python < 3.10 we use the backport from typing_extensions.
-if sys.version_info >= (3, 10):
-    from typing import TypeAlias, get_args, get_origin
-else:
-    from typing_extensions import TypeAlias, get_args, get_origin
+from typing import Any, Callable, Dict, List, TypeVar, Union
 
 AnyCallableT = TypeVar("AnyCallableT", bound=Callable[..., Any])  # noqa: VNE001
 # JSON primitives only, mypy doesn't support recursive tho
 JSONType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
-
-__all__ = ["get_args", "get_origin", "Annotated", "Protocol", "TypedDict", "Literal", "NotRequired", "TypeAlias"]
diff --git a/aws_lambda_powertools/utilities/batch/types.py b/aws_lambda_powertools/utilities/batch/types.py
index d48f768a6b8..d737480bf8f 100644
--- a/aws_lambda_powertools/utilities/batch/types.py
+++ b/aws_lambda_powertools/utilities/batch/types.py
@@ -1,7 +1,5 @@
 import sys
-from typing import Optional, Type, Union
-
-from aws_lambda_powertools.shared.types import List, TypedDict
+from typing import List, Optional, Type, TypedDict, Union
 
 has_pydantic = "pydantic" in sys.modules
 
diff --git a/aws_lambda_powertools/utilities/data_classes/s3_batch_operation_event.py b/aws_lambda_powertools/utilities/data_classes/s3_batch_operation_event.py
index 5419f6f8088..bd9a07c4484 100644
--- a/aws_lambda_powertools/utilities/data_classes/s3_batch_operation_event.py
+++ b/aws_lambda_powertools/utilities/data_classes/s3_batch_operation_event.py
@@ -1,9 +1,8 @@
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple
 from urllib.parse import unquote_plus
 
-from aws_lambda_powertools.shared.types import Literal
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 # list of valid result code. Used both in S3BatchOperationResponse and S3BatchOperationResponseRecord
diff --git a/aws_lambda_powertools/utilities/data_classes/secrets_manager_event.py b/aws_lambda_powertools/utilities/data_classes/secrets_manager_event.py
index 1a3a1c5b7f4..35b32685c15 100644
--- a/aws_lambda_powertools/utilities/data_classes/secrets_manager_event.py
+++ b/aws_lambda_powertools/utilities/data_classes/secrets_manager_event.py
@@ -1,4 +1,5 @@
-from aws_lambda_powertools.shared.types import Literal
+from typing import Literal
+
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 
diff --git a/aws_lambda_powertools/utilities/idempotency/hook.py b/aws_lambda_powertools/utilities/idempotency/hook.py
index 0027399b937..1167e6264bf 100644
--- a/aws_lambda_powertools/utilities/idempotency/hook.py
+++ b/aws_lambda_powertools/utilities/idempotency/hook.py
@@ -1,6 +1,5 @@
-from typing import Any
+from typing import Any, Protocol
 
-from aws_lambda_powertools.shared.types import Protocol
 from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
 
 
diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/redis.py b/aws_lambda_powertools/utilities/idempotency/persistence/redis.py
index 6dda3b7fbcd..44e3767312a 100644
--- a/aws_lambda_powertools/utilities/idempotency/persistence/redis.py
+++ b/aws_lambda_powertools/utilities/idempotency/persistence/redis.py
@@ -5,11 +5,10 @@
 import logging
 from contextlib import contextmanager
 from datetime import timedelta
-from typing import Any, Dict
+from typing import Any, Dict, Literal, Protocol
 
 import redis
 
-from aws_lambda_powertools.shared.types import Literal, Protocol
 from aws_lambda_powertools.utilities.idempotency import BasePersistenceLayer
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyItemAlreadyExistsError,
diff --git a/aws_lambda_powertools/utilities/parameters/ssm.py b/aws_lambda_powertools/utilities/parameters/ssm.py
index 60db342d8b9..c646d27f3dc 100644
--- a/aws_lambda_powertools/utilities/parameters/ssm.py
+++ b/aws_lambda_powertools/utilities/parameters/ssm.py
@@ -6,7 +6,7 @@
 
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, overload
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, overload
 
 import boto3
 from botocore.config import Config
@@ -17,7 +17,6 @@
     resolve_truthy_env_var_choice,
     slice_dictionary,
 )
-from aws_lambda_powertools.shared.types import Literal
 from aws_lambda_powertools.utilities.parameters.base import (
     DEFAULT_MAX_AGE_SECS,
     DEFAULT_PROVIDERS,
diff --git a/aws_lambda_powertools/utilities/parameters/types.py b/aws_lambda_powertools/utilities/parameters/types.py
index faa06cee89e..84ea46fe3db 100644
--- a/aws_lambda_powertools/utilities/parameters/types.py
+++ b/aws_lambda_powertools/utilities/parameters/types.py
@@ -1,3 +1,3 @@
-from aws_lambda_powertools.shared.types import Literal
+from typing import Literal
 
 TransformOptions = Literal["json", "binary", "auto", None]
diff --git a/aws_lambda_powertools/utilities/parser/models/cloudwatch.py b/aws_lambda_powertools/utilities/parser/models/cloudwatch.py
index c6b297c7118..df464edd65e 100644
--- a/aws_lambda_powertools/utilities/parser/models/cloudwatch.py
+++ b/aws_lambda_powertools/utilities/parser/models/cloudwatch.py
@@ -3,12 +3,10 @@
 import logging
 import zlib
 from datetime import datetime
-from typing import Optional, Type, Union
+from typing import List, Optional, Type, Union
 
 from pydantic import BaseModel, Field, field_validator
 
-from aws_lambda_powertools.shared.types import List
-
 logger = logging.getLogger(__name__)
 
 
diff --git a/aws_lambda_powertools/utilities/parser/types.py b/aws_lambda_powertools/utilities/parser/types.py
index e7654e3acc2..91bf9a9119e 100644
--- a/aws_lambda_powertools/utilities/parser/types.py
+++ b/aws_lambda_powertools/utilities/parser/types.py
@@ -1,11 +1,9 @@
 """Generics and other shared types used across parser"""
 
-from typing import Any, Dict, Type, TypeVar, Union
+from typing import Any, Dict, Literal, Type, TypeVar, Union
 
 from pydantic import BaseModel, Json
 
-from aws_lambda_powertools.shared.types import Literal
-
 Model = TypeVar("Model", bound=BaseModel)
 EnvelopeModel = TypeVar("EnvelopeModel")
 EventParserReturnType = TypeVar("EventParserReturnType")
diff --git a/aws_lambda_powertools/utilities/streaming/s3_object.py b/aws_lambda_powertools/utilities/streaming/s3_object.py
index 2c8e7f5ea6a..62bc4385c3b 100644
--- a/aws_lambda_powertools/utilities/streaming/s3_object.py
+++ b/aws_lambda_powertools/utilities/streaming/s3_object.py
@@ -7,6 +7,7 @@
     Any,
     Iterable,
     List,
+    Literal,
     Optional,
     Sequence,
     TypeVar,
@@ -15,7 +16,6 @@
     overload,
 )
 
-from aws_lambda_powertools.shared.types import Literal
 from aws_lambda_powertools.utilities.streaming._s3_seekable_io import _S3SeekableIO
 from aws_lambda_powertools.utilities.streaming.transformations import (
     CsvTransform,
diff --git a/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response_module.py b/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response_module.py
index d197e470595..e3901c3e243 100644
--- a/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response_module.py
+++ b/examples/event_handler_bedrock_agents/src/assert_bedrock_agent_response_module.py
@@ -1,9 +1,10 @@
 import time
 
+from typing_extensions import Annotated
+
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import BedrockAgentResolver
 from aws_lambda_powertools.event_handler.openapi.params import Body
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_graphql/src/assert_async_graphql_response_module.py b/examples/event_handler_graphql/src/assert_async_graphql_response_module.py
index 371eeaa23f8..e647f7f1b74 100644
--- a/examples/event_handler_graphql/src/assert_async_graphql_response_module.py
+++ b/examples/event_handler_graphql/src/assert_async_graphql_response_module.py
@@ -1,12 +1,11 @@
 import asyncio
-from typing import List
+from typing import List, TypedDict
 
 import aiohttp
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.tracing import aiohttp_trace_config
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
diff --git a/examples/event_handler_graphql/src/assert_graphql_response_module.py b/examples/event_handler_graphql/src/assert_graphql_response_module.py
index a7cb58c1d98..60c005c95f5 100644
--- a/examples/event_handler_graphql/src/assert_graphql_response_module.py
+++ b/examples/event_handler_graphql/src/assert_graphql_response_module.py
@@ -1,9 +1,8 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_graphql/src/async_resolvers.py b/examples/event_handler_graphql/src/async_resolvers.py
index 08ecbcba85b..610c8afbdc8 100644
--- a/examples/event_handler_graphql/src/async_resolvers.py
+++ b/examples/event_handler_graphql/src/async_resolvers.py
@@ -1,12 +1,11 @@
 import asyncio
-from typing import List
+from typing import List, TypedDict
 
 import aiohttp
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.tracing import aiohttp_trace_config
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
diff --git a/examples/event_handler_graphql/src/custom_models.py b/examples/event_handler_graphql/src/custom_models.py
index 21f5f07af00..6f1b80fe8d0 100644
--- a/examples/event_handler_graphql/src/custom_models.py
+++ b/examples/event_handler_graphql/src/custom_models.py
@@ -1,9 +1,8 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.data_classes.appsync import scalar_types_utils
 from aws_lambda_powertools.utilities.data_classes.appsync_resolver_event import (
     AppSyncResolverEvent,
diff --git a/examples/event_handler_graphql/src/getting_started_graphql_api_resolver.py b/examples/event_handler_graphql/src/getting_started_graphql_api_resolver.py
index e960a357d17..1e3925039ae 100644
--- a/examples/event_handler_graphql/src/getting_started_graphql_api_resolver.py
+++ b/examples/event_handler_graphql/src/getting_started_graphql_api_resolver.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, TypedDict
 
 import requests
 from requests import Response
@@ -6,7 +6,6 @@
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.data_classes.appsync import scalar_types_utils
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
diff --git a/examples/event_handler_graphql/src/graphql_transformer_merchant_info.py b/examples/event_handler_graphql/src/graphql_transformer_merchant_info.py
index 017528e3481..ec751882fe3 100644
--- a/examples/event_handler_graphql/src/graphql_transformer_merchant_info.py
+++ b/examples/event_handler_graphql/src/graphql_transformer_merchant_info.py
@@ -1,9 +1,8 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.data_classes.appsync import scalar_types_utils
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
diff --git a/examples/event_handler_graphql/src/graphql_transformer_search_merchant.py b/examples/event_handler_graphql/src/graphql_transformer_search_merchant.py
index 9b685a280dd..895b1f539e2 100644
--- a/examples/event_handler_graphql/src/graphql_transformer_search_merchant.py
+++ b/examples/event_handler_graphql/src/graphql_transformer_search_merchant.py
@@ -1,9 +1,8 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.data_classes.appsync import scalar_types_utils
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
diff --git a/examples/event_handler_graphql/src/nested_mappings.py b/examples/event_handler_graphql/src/nested_mappings.py
index a7cb58c1d98..60c005c95f5 100644
--- a/examples/event_handler_graphql/src/nested_mappings.py
+++ b/examples/event_handler_graphql/src/nested_mappings.py
@@ -1,9 +1,8 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import AppSyncResolver
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import TypedDict
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_graphql/src/split_operation_append_context_module.py b/examples/event_handler_graphql/src/split_operation_append_context_module.py
index 15ed7af1b9e..7b81241e8fe 100644
--- a/examples/event_handler_graphql/src/split_operation_append_context_module.py
+++ b/examples/event_handler_graphql/src/split_operation_append_context_module.py
@@ -1,8 +1,7 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler.appsync import Router
-from aws_lambda_powertools.shared.types import TypedDict
 
 tracer = Tracer()
 logger = Logger()
diff --git a/examples/event_handler_graphql/src/split_operation_module.py b/examples/event_handler_graphql/src/split_operation_module.py
index e4c7f978b73..5a97128b1e2 100644
--- a/examples/event_handler_graphql/src/split_operation_module.py
+++ b/examples/event_handler_graphql/src/split_operation_module.py
@@ -1,8 +1,7 @@
-from typing import List
+from typing import List, TypedDict
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler.appsync import Router
-from aws_lambda_powertools.shared.types import TypedDict
 
 tracer = Tracer()
 logger = Logger()
diff --git a/examples/event_handler_rest/src/validating_headers.py b/examples/event_handler_rest/src/validating_headers.py
index e830a49c38c..f92f0ead463 100644
--- a/examples/event_handler_rest/src/validating_headers.py
+++ b/examples/event_handler_rest/src/validating_headers.py
@@ -2,12 +2,12 @@
 
 import requests
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated  # (1)!
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Header  # (2)!
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import Annotated  # (1)!
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_rest/src/validating_path.py b/examples/event_handler_rest/src/validating_path.py
index e892e1c8597..c5ddbba4f37 100644
--- a/examples/event_handler_rest/src/validating_path.py
+++ b/examples/event_handler_rest/src/validating_path.py
@@ -2,12 +2,12 @@
 
 import requests
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Path
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_rest/src/validating_payload_subset.py b/examples/event_handler_rest/src/validating_payload_subset.py
index ac4ee603853..ebd0cf0c20f 100644
--- a/examples/event_handler_rest/src/validating_payload_subset.py
+++ b/examples/event_handler_rest/src/validating_payload_subset.py
@@ -2,10 +2,10 @@
 
 import requests
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Body  # (1)!
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 app = APIGatewayRestResolver(enable_validation=True)
diff --git a/examples/event_handler_rest/src/validating_query_strings.py b/examples/event_handler_rest/src/validating_query_strings.py
index 21d34dbd25a..047e9973b63 100644
--- a/examples/event_handler_rest/src/validating_query_strings.py
+++ b/examples/event_handler_rest/src/validating_query_strings.py
@@ -2,12 +2,12 @@
 
 import requests
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated  # (1)!
 
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Query  # (2)!
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.shared.types import Annotated  # (1)!
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
diff --git a/examples/event_handler_rest/src/working_with_headers_multi_value.py b/examples/event_handler_rest/src/working_with_headers_multi_value.py
index 956fd58b14d..6cb2e17f46d 100644
--- a/examples/event_handler_rest/src/working_with_headers_multi_value.py
+++ b/examples/event_handler_rest/src/working_with_headers_multi_value.py
@@ -1,9 +1,10 @@
 from enum import Enum
 from typing import List
 
+from typing_extensions import Annotated
+
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Header
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 app = APIGatewayRestResolver(enable_validation=True)
diff --git a/examples/event_handler_rest/src/working_with_multi_query_values.py b/examples/event_handler_rest/src/working_with_multi_query_values.py
index 7f6049dad46..6f60447e890 100644
--- a/examples/event_handler_rest/src/working_with_multi_query_values.py
+++ b/examples/event_handler_rest/src/working_with_multi_query_values.py
@@ -1,9 +1,10 @@
 from enum import Enum
 from typing import List
 
+from typing_extensions import Annotated
+
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.params import Query
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 app = APIGatewayRestResolver(enable_validation=True)
diff --git a/examples/parser/src/multiple_model_parsing.py b/examples/parser/src/multiple_model_parsing.py
index adbde35e4d0..556848bbff6 100644
--- a/examples/parser/src/multiple_model_parsing.py
+++ b/examples/parser/src/multiple_model_parsing.py
@@ -1,8 +1,8 @@
 from typing import Any, Literal, Union
 
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.parser import event_parser
 
 
diff --git a/tests/functional/event_handler/test_bedrock_agent.py b/tests/functional/event_handler/test_bedrock_agent.py
index c43f4cc4ea9..7d2da8c0486 100644
--- a/tests/functional/event_handler/test_bedrock_agent.py
+++ b/tests/functional/event_handler/test_bedrock_agent.py
@@ -1,9 +1,10 @@
 import json
 from typing import Any, Dict
 
+from typing_extensions import Annotated
+
 from aws_lambda_powertools.event_handler import BedrockAgentResolver, Response, content_types
 from aws_lambda_powertools.event_handler.openapi.params import Body
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
 from tests.functional.utils import load_event
 
diff --git a/tests/functional/event_handler/test_openapi_params.py b/tests/functional/event_handler/test_openapi_params.py
index 2ac9c036f3f..3d8bb73e4bf 100644
--- a/tests/functional/event_handler/test_openapi_params.py
+++ b/tests/functional/event_handler/test_openapi_params.py
@@ -3,6 +3,7 @@
 from typing import List
 
 from pydantic import BaseModel
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler.api_gateway import APIGatewayRestResolver, Response, Router
 from aws_lambda_powertools.event_handler.openapi.models import (
@@ -19,7 +20,6 @@
     Query,
     _create_model_field,
 )
-from aws_lambda_powertools.shared.types import Annotated
 
 JSON_CONTENT_TYPE = "application/json"
 
diff --git a/tests/functional/event_handler/test_openapi_schema_pydantic_v1.py b/tests/functional/event_handler/test_openapi_schema_pydantic_v1.py
index f2f80c51bc2..cb09562acfb 100644
--- a/tests/functional/event_handler/test_openapi_schema_pydantic_v1.py
+++ b/tests/functional/event_handler/test_openapi_schema_pydantic_v1.py
@@ -1,15 +1,15 @@
 import json
 import warnings
-from typing import Optional
+from typing import Literal, Optional
 
 import pytest
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.models import Contact, License, Server
 from aws_lambda_powertools.event_handler.openapi.params import Query
 from aws_lambda_powertools.event_handler.openapi.types import OpenAPIResponse
-from aws_lambda_powertools.shared.types import Annotated, Literal
 
 
 @pytest.mark.usefixtures("pydanticv1_only")
diff --git a/tests/functional/event_handler/test_openapi_schema_pydantic_v2.py b/tests/functional/event_handler/test_openapi_schema_pydantic_v2.py
index e52b8279912..0df8f6a22c5 100644
--- a/tests/functional/event_handler/test_openapi_schema_pydantic_v2.py
+++ b/tests/functional/event_handler/test_openapi_schema_pydantic_v2.py
@@ -1,15 +1,15 @@
 import json
 import warnings
-from typing import Optional
+from typing import Literal, Optional
 
 import pytest
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.openapi.models import Contact, License, Server
 from aws_lambda_powertools.event_handler.openapi.params import Query
 from aws_lambda_powertools.event_handler.openapi.types import OpenAPIResponse
-from aws_lambda_powertools.shared.types import Annotated, Literal
 
 
 @pytest.mark.usefixtures("pydanticv2_only")
diff --git a/tests/functional/event_handler/test_openapi_validation_middleware.py b/tests/functional/event_handler/test_openapi_validation_middleware.py
index da83f6f92f1..54425f34986 100644
--- a/tests/functional/event_handler/test_openapi_validation_middleware.py
+++ b/tests/functional/event_handler/test_openapi_validation_middleware.py
@@ -6,6 +6,7 @@
 
 import pytest
 from pydantic import BaseModel
+from typing_extensions import Annotated
 
 from aws_lambda_powertools.event_handler import (
     ALBResolver,
@@ -17,7 +18,6 @@
     VPCLatticeV2Resolver,
 )
 from aws_lambda_powertools.event_handler.openapi.params import Body, Header, Query
-from aws_lambda_powertools.shared.types import Annotated
 
 
 def test_validate_scalars(gw_event):
diff --git a/tests/functional/parser/test_parser.py b/tests/functional/parser/test_parser.py
index fdcfffe0c38..6b34d9d664a 100644
--- a/tests/functional/parser/test_parser.py
+++ b/tests/functional/parser/test_parser.py
@@ -3,8 +3,8 @@
 
 import pydantic
 import pytest
+from typing_extensions import Annotated
 
-from aws_lambda_powertools.shared.types import Annotated
 from aws_lambda_powertools.utilities.parser import (
     event_parser,
     exceptions,
diff --git a/tests/performance/parser/test_parser_performance.py b/tests/performance/parser/test_parser_performance.py
index 724368dbe2a..c27fa10fbdf 100644
--- a/tests/performance/parser/test_parser_performance.py
+++ b/tests/performance/parser/test_parser_performance.py
@@ -1,11 +1,11 @@
 import time
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Literal, Union
 
 import pytest
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
-from aws_lambda_powertools.shared.types import Annotated, Literal, Union
 from aws_lambda_powertools.utilities.parser import parse
 
 # adjusted for slower machines in CI too

From 64254115184ba7aac275a22328535c15e14440a9 Mon Sep 17 00:00:00 2001
From: Leandro Damascena <lcdama@amazon.pt>
Date: Tue, 6 Aug 2024 19:33:21 +0100
Subject: [PATCH 2/2] Docmentation fix + small changes

---
 aws_lambda_powertools/shared/types.py          |  4 +---
 .../utilities/feature_flags/__init__.py        | 10 +++++-----
 .../utilities/feature_flags/appconfig.py       |  7 +++----
 .../utilities/feature_flags/comparators.py     |  2 +-
 .../utilities/feature_flags/feature_flags.py   | 12 ++++++------
 .../utilities/feature_flags/schema.py          |  6 +++---
 .../utilities/feature_flags/types.py           |  4 ++++
 docs/core/event_handler/api_gateway.md         |  6 +++---
 docs/core/event_handler/appsync.md             | 18 +++++++++---------
 .../combining_powertools_utilities_function.py |  2 +-
 .../feature_flags/test_time_based_actions.py   |  2 +-
 11 files changed, 37 insertions(+), 36 deletions(-)
 create mode 100644 aws_lambda_powertools/utilities/feature_flags/types.py

diff --git a/aws_lambda_powertools/shared/types.py b/aws_lambda_powertools/shared/types.py
index e4e10192e55..c5c91535bd3 100644
--- a/aws_lambda_powertools/shared/types.py
+++ b/aws_lambda_powertools/shared/types.py
@@ -1,5 +1,3 @@
-from typing import Any, Callable, Dict, List, TypeVar, Union
+from typing import Any, Callable, TypeVar
 
 AnyCallableT = TypeVar("AnyCallableT", bound=Callable[..., Any])  # noqa: VNE001
-# JSON primitives only, mypy doesn't support recursive tho
-JSONType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
diff --git a/aws_lambda_powertools/utilities/feature_flags/__init__.py b/aws_lambda_powertools/utilities/feature_flags/__init__.py
index e8d8229c9dc..4514e92e991 100644
--- a/aws_lambda_powertools/utilities/feature_flags/__init__.py
+++ b/aws_lambda_powertools/utilities/feature_flags/__init__.py
@@ -1,10 +1,10 @@
 """Advanced feature flags utility"""
 
-from .appconfig import AppConfigStore
-from .base import StoreProvider
-from .exceptions import ConfigurationStoreError
-from .feature_flags import FeatureFlags
-from .schema import RuleAction, SchemaValidator
+from aws_lambda_powertools.utilities.feature_flags.appconfig import AppConfigStore
+from aws_lambda_powertools.utilities.feature_flags.base import StoreProvider
+from aws_lambda_powertools.utilities.feature_flags.exceptions import ConfigurationStoreError
+from aws_lambda_powertools.utilities.feature_flags.feature_flags import FeatureFlags
+from aws_lambda_powertools.utilities.feature_flags.schema import RuleAction, SchemaValidator
 
 __all__ = [
     "ConfigurationStoreError",
diff --git a/aws_lambda_powertools/utilities/feature_flags/appconfig.py b/aws_lambda_powertools/utilities/feature_flags/appconfig.py
index dd1750f2118..b979c32c9a8 100644
--- a/aws_lambda_powertools/utilities/feature_flags/appconfig.py
+++ b/aws_lambda_powertools/utilities/feature_flags/appconfig.py
@@ -4,17 +4,16 @@
 
 from botocore.config import Config
 
+from aws_lambda_powertools.logging import Logger
 from aws_lambda_powertools.utilities import jmespath_utils
+from aws_lambda_powertools.utilities.feature_flags.base import StoreProvider
+from aws_lambda_powertools.utilities.feature_flags.exceptions import ConfigurationStoreError, StoreClientError
 from aws_lambda_powertools.utilities.parameters import (
     AppConfigProvider,
     GetParameterError,
     TransformParameterError,
 )
 
-from ... import Logger
-from .base import StoreProvider
-from .exceptions import ConfigurationStoreError, StoreClientError
-
 
 class AppConfigStore(StoreProvider):
     def __init__(
diff --git a/aws_lambda_powertools/utilities/feature_flags/comparators.py b/aws_lambda_powertools/utilities/feature_flags/comparators.py
index 03cb91e649a..035566cad4c 100644
--- a/aws_lambda_powertools/utilities/feature_flags/comparators.py
+++ b/aws_lambda_powertools/utilities/feature_flags/comparators.py
@@ -5,7 +5,7 @@
 
 from dateutil.tz import gettz
 
-from .schema import HOUR_MIN_SEPARATOR, ModuloRangeValues, TimeValues
+from aws_lambda_powertools.utilities.feature_flags.schema import HOUR_MIN_SEPARATOR, ModuloRangeValues, TimeValues
 
 
 def _get_now_from_timezone(timezone: Optional[tzinfo]) -> datetime:
diff --git a/aws_lambda_powertools/utilities/feature_flags/feature_flags.py b/aws_lambda_powertools/utilities/feature_flags/feature_flags.py
index bd7e19d0efe..2b93887138c 100644
--- a/aws_lambda_powertools/utilities/feature_flags/feature_flags.py
+++ b/aws_lambda_powertools/utilities/feature_flags/feature_flags.py
@@ -5,11 +5,10 @@
 
 from typing_extensions import ParamSpec
 
-from ... import Logger
-from ...shared.types import JSONType
-from . import schema
-from .base import StoreProvider
-from .comparators import (
+from aws_lambda_powertools.logging import Logger
+from aws_lambda_powertools.utilities.feature_flags import schema
+from aws_lambda_powertools.utilities.feature_flags.base import StoreProvider
+from aws_lambda_powertools.utilities.feature_flags.comparators import (
     compare_all_in_list,
     compare_any_in_list,
     compare_datetime_range,
@@ -18,7 +17,8 @@
     compare_none_in_list,
     compare_time_range,
 )
-from .exceptions import ConfigurationStoreError
+from aws_lambda_powertools.utilities.feature_flags.exceptions import ConfigurationStoreError
+from aws_lambda_powertools.utilities.feature_flags.types import JSONType
 
 T = TypeVar("T")
 P = ParamSpec("P")
diff --git a/aws_lambda_powertools/utilities/feature_flags/schema.py b/aws_lambda_powertools/utilities/feature_flags/schema.py
index 1df16677bd8..0fdb5e47810 100644
--- a/aws_lambda_powertools/utilities/feature_flags/schema.py
+++ b/aws_lambda_powertools/utilities/feature_flags/schema.py
@@ -9,9 +9,9 @@
 
 from dateutil import tz
 
-from ... import Logger
-from .base import BaseValidator
-from .exceptions import SchemaValidationError
+from aws_lambda_powertools.logging import Logger
+from aws_lambda_powertools.utilities.feature_flags.base import BaseValidator
+from aws_lambda_powertools.utilities.feature_flags.exceptions import SchemaValidationError
 
 RULES_KEY = "rules"
 FEATURE_DEFAULT_VAL_KEY = "default"
diff --git a/aws_lambda_powertools/utilities/feature_flags/types.py b/aws_lambda_powertools/utilities/feature_flags/types.py
new file mode 100644
index 00000000000..6df79e5d608
--- /dev/null
+++ b/aws_lambda_powertools/utilities/feature_flags/types.py
@@ -0,0 +1,4 @@
+from typing import Any, Dict, List, Union
+
+# JSON primitives only, mypy doesn't support recursive tho
+JSONType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
diff --git a/docs/core/event_handler/api_gateway.md b/docs/core/event_handler/api_gateway.md
index 3c182b30e4e..cb4b1ee61be 100644
--- a/docs/core/event_handler/api_gateway.md
+++ b/docs/core/event_handler/api_gateway.md
@@ -380,7 +380,7 @@ In the following example, we use a new `Query` OpenAPI type to add [one out of m
 
 === "validating_query_strings.py"
 
-    ```python hl_lines="8 10 27"
+    ```python hl_lines="8 9 27"
     --8<-- "examples/event_handler_rest/src/validating_query_strings.py"
     ```
 
@@ -418,7 +418,7 @@ Just like we learned in [query string validation](#validating-query-strings), we
 
 For example, we could validate that `<todo_id>` dynamic path should be no greater than three digits.
 
-```python hl_lines="8 10 27" title="validating_path.py"
+```python hl_lines="8 9 27" title="validating_path.py"
 --8<-- "examples/event_handler_rest/src/validating_path.py"
 ```
 
@@ -440,7 +440,7 @@ In the following example, we use a new `Header` OpenAPI type to add [one out of
 
 === "validating_headers.py"
 
-    ```python hl_lines="8 10 27"
+    ```python hl_lines="5 9 27"
     --8<-- "examples/event_handler_rest/src/validating_headers.py"
     ```
 
diff --git a/docs/core/event_handler/appsync.md b/docs/core/event_handler/appsync.md
index fcadc2a1f27..091de6fea64 100644
--- a/docs/core/event_handler/appsync.md
+++ b/docs/core/event_handler/appsync.md
@@ -61,7 +61,7 @@ Here's an example with two separate functions to resolve `getTodo` and `listTodo
 
 === "getting_started_graphql_api_resolver.py"
 
-    ```python hl_lines="7 15 25 27 28 37 39 47 49 60"
+    ```python hl_lines="7 14 24 26 27 36 38 46 48 59"
     --8<-- "examples/event_handler_graphql/src/getting_started_graphql_api_resolver.py"
     ```
 
@@ -123,7 +123,7 @@ You can nest `app.resolver()` decorator multiple times when resolving fields wit
 
 === "nested_mappings.py"
 
-    ```python hl_lines="4 11 21 22 24 31"
+    ```python hl_lines="4 10 20 21 23 30"
     --8<-- "examples/event_handler_graphql/src/nested_mappings.py"
     ```
 
@@ -137,7 +137,7 @@ You can nest `app.resolver()` decorator multiple times when resolving fields wit
 
 For Lambda Python3.8+ runtime, this utility supports async functions when you use in conjunction with `asyncio.run`.
 
-```python hl_lines="6 15 25 26 35 37" title="Resolving GraphQL resolvers async"
+```python hl_lines="7 14 24 25 34 36" title="Resolving GraphQL resolvers async"
 --8<-- "examples/event_handler_graphql/src/async_resolvers.py"
 ```
 
@@ -162,13 +162,13 @@ Use the following code for `merchantInfo` and `searchMerchant` functions respect
 
 === "graphql_transformer_merchant_info.py"
 
-    ```python hl_lines="4 7 23 24 29 30 37"
+    ```python hl_lines="4 6 23 24 29 30 36"
     --8<-- "examples/event_handler_graphql/src/graphql_transformer_merchant_info.py"
     ```
 
 === "graphql_transformer_search_merchant.py"
 
-    ```python hl_lines="4 7 22 23 37 43"
+    ```python hl_lines="4 6 21 22 36 42"
     --8<-- "examples/event_handler_graphql/src/graphql_transformer_search_merchant.py"
     ```
 
@@ -196,7 +196,7 @@ You can subclass [AppSyncResolverEvent](../../utilities/data_classes.md#appsync-
 
 === "custom_models.py.py"
 
-    ```python hl_lines="4 8-10 26-28 31 32 39 46"
+    ```python hl_lines="4 7-9 25-27 31 32 39 45"
     --8<-- "examples/event_handler_graphql/src/custom_models.py"
     ```
 
@@ -225,7 +225,7 @@ Let's assume you have `split_operation.py` as your Lambda function entrypoint an
 
     We import **Router** instead of **AppSyncResolver**; syntax wise is exactly the same.
 
-  	```python hl_lines="4 9 19 20"
+  	```python hl_lines="4 8 18 19"
     --8<-- "examples/event_handler_graphql/src/split_operation_module.py"
   	```
 
@@ -255,7 +255,7 @@ You can use `append_context` when you want to share data between your App and Ro
 
 === "split_route_append_context_module.py"
 
-	```python hl_lines="23"
+	```python hl_lines="22"
     --8<-- "examples/event_handler_graphql/src/split_operation_append_context_module.py"
 	```
 
@@ -298,7 +298,7 @@ And an example for testing asynchronous resolvers. Note that this requires the `
 
 === "assert_async_graphql_response_module.py"
 
-    ```python hl_lines="15"
+    ```python hl_lines="14"
     --8<-- "examples/event_handler_graphql/src/assert_async_graphql_response_module.py"
     ```
 
diff --git a/examples/middleware_factory/src/combining_powertools_utilities_function.py b/examples/middleware_factory/src/combining_powertools_utilities_function.py
index 32a59723ead..23c153db444 100644
--- a/examples/middleware_factory/src/combining_powertools_utilities_function.py
+++ b/examples/middleware_factory/src/combining_powertools_utilities_function.py
@@ -9,8 +9,8 @@
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver
 from aws_lambda_powertools.event_handler.exceptions import InternalServerError
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
-from aws_lambda_powertools.shared.types import JSONType
 from aws_lambda_powertools.utilities.feature_flags import AppConfigStore, FeatureFlags
+from aws_lambda_powertools.utilities.feature_flags.types import JSONType
 from aws_lambda_powertools.utilities.jmespath_utils import extract_data_from_envelope
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_lambda_powertools.utilities.validation import SchemaValidationError, validate
diff --git a/tests/functional/feature_flags/test_time_based_actions.py b/tests/functional/feature_flags/test_time_based_actions.py
index 8b850d52fc3..872f2ac2862 100644
--- a/tests/functional/feature_flags/test_time_based_actions.py
+++ b/tests/functional/feature_flags/test_time_based_actions.py
@@ -4,7 +4,6 @@
 from botocore.config import Config
 from dateutil.tz import gettz
 
-from aws_lambda_powertools.shared.types import JSONType
 from aws_lambda_powertools.utilities.feature_flags.appconfig import AppConfigStore
 from aws_lambda_powertools.utilities.feature_flags.feature_flags import FeatureFlags
 from aws_lambda_powertools.utilities.feature_flags.schema import (
@@ -19,6 +18,7 @@
     TimeKeys,
     TimeValues,
 )
+from aws_lambda_powertools.utilities.feature_flags.types import JSONType
 
 
 def evaluate_mocked_schema(