Replace import with aio imports in wrapper packages

mypy_boto3_builder/import_helpers/import_record_group.py

@@ -8,15 +8,15 @@
 from typing import Final
 
 from mypy_boto3_builder.import_helpers.import_record import ImportRecord
-from mypy_boto3_builder.import_helpers.import_string import Import, ImportString
+from mypy_boto3_builder.import_helpers.import_string import Import
 
 
 class ImportRecordGroup:
     """
     Group tool for ImportRecord sets.
     """
 
-    _SYS_IMPORT_RECORD: Final[ImportRecord] = ImportRecord(ImportString("sys"))
+    _SYS_IMPORT_RECORD: Final[ImportRecord] = ImportRecord(Import.sys)
 
     def __init__(self, records: Iterable[ImportRecord] = ()) -> None:
         self.records: set[ImportRecord] = set()

mypy_boto3_builder/import_helpers/import_string.py

@@ -178,6 +178,8 @@ class Import:
     aiobotocore: Final = ImportString("aiobotocore")
     aioboto3: Final = ImportString("aioboto3")
     typing_extensions: Final = ImportString("typing_extensions")
+    types: Final = ImportString("types")
+    sys: Final = ImportString("sys")
 
     _THIRD_PARTY: Final[set[str]] = {
         boto3.parent,

mypy_boto3_builder/parsers/parse_wrapper_package.py

@@ -9,6 +9,7 @@
 from mypy_boto3_builder.import_helpers.import_string import ImportString
 from mypy_boto3_builder.package_data import BasePackageData
 from mypy_boto3_builder.parsers.wrapper_package_parser import WrapperPackageParser
+from mypy_boto3_builder.postprocessors.aio_imports import replace_imports_with_aio
 from mypy_boto3_builder.service_name import ServiceName
 from mypy_boto3_builder.structures.types_aioboto3_package import TypesAioBoto3Package
 from mypy_boto3_builder.structures.types_aiobotocore_package import TypesAioBotocorePackage
@@ -68,6 +69,8 @@ def parse_types_aiobotocore_package(
         )
         method.type_ignore = "override"
         package.session_class.methods.append(method)
+
+    replace_imports_with_aio(package.session_class.iterate_types())
     return package
 
 
@@ -101,4 +104,6 @@ def parse_types_aioboto3_package(
             [method.return_type],
         )
         package.session_class.methods.append(method)
+
+    replace_imports_with_aio(package.session_class.iterate_types())
     return package

mypy_boto3_builder/postprocessors/aio_imports.py

@@ -0,0 +1,63 @@
+"""
+Import replacement map for aio libraries.
+
+Copyright 2024 Vlad Emelianov
+"""
+
+from collections.abc import Iterator, Mapping
+from typing import Final
+
+from botocore.client import BaseClient
+from botocore.config import Config
+from botocore.eventstream import EventStream
+from botocore.paginate import Paginator
+from botocore.response import StreamingBody
+from botocore.waiter import Waiter
+
+from mypy_boto3_builder.import_helpers.import_string import Import
+from mypy_boto3_builder.type_annotations.external_import import ExternalImport
+from mypy_boto3_builder.type_annotations.fake_annotation import FakeAnnotation
+from mypy_boto3_builder.utils.type_checks import is_external_import
+
+AIO_IMPORT_MAP: Final[Mapping[ExternalImport, ExternalImport]] = {
+    ExternalImport.from_class(StreamingBody): ExternalImport(
+        Import.aiobotocore + "response", "StreamingBody"
+    ),
+    ExternalImport.from_class(EventStream): ExternalImport(
+        Import.aiobotocore + "eventstream", "AioEventStream"
+    ),
+    ExternalImport.from_class(Config): ExternalImport(Import.aiobotocore + "config", "AioConfig"),
+    ExternalImport.from_class(Waiter): ExternalImport(Import.aiobotocore + "waiter", "AIOWaiter"),
+    ExternalImport.from_class(Paginator): ExternalImport(
+        Import.aiobotocore + "paginate", "AioPaginator"
+    ),
+    ExternalImport.from_class(BaseClient): ExternalImport(
+        Import.aiobotocore + "client", "AioBaseClient"
+    ),
+    ExternalImport(Import.boto3 + "resources" + "base", "ServiceResource"): ExternalImport(
+        Import.aioboto3 + "resources" + "base", "AIOBoto3ServiceResource"
+    ),
+    ExternalImport(Import.boto3 + "resources" + "collection", "ResourceCollection"): ExternalImport(
+        Import.aioboto3 + "resources" + "collection", "AIOResourceCollection"
+    ),
+    ExternalImport(Import.boto3 + "dynamodb" + "table", "BatchWriter"): ExternalImport(
+        Import.aioboto3 + "dynamodb" + "table", "BatchWriter"
+    ),
+    ExternalImport(Import.boto3 + "dynamodb" + "table", "TableResource"): ExternalImport(
+        Import.aioboto3 + "dynamodb" + "table", "CustomTableResource"
+    ),
+}
+
+
+def replace_imports_with_aio(type_annotations: Iterator[FakeAnnotation]) -> None:
+    """
+    Replace botocore/boto3 imports in type annotations with aioboto3/aiobotocore.
+    """
+    for type_annotation in type_annotations:
+        if not is_external_import(type_annotation):
+            continue
+
+        if type_annotation in AIO_IMPORT_MAP:
+            new_type_annotation = AIO_IMPORT_MAP[type_annotation]
+            type_annotation.copy_from(new_type_annotation)
+            continue

mypy_boto3_builder/postprocessors/aiobotocore.py

@@ -4,17 +4,11 @@
 Copyright 2024 Vlad Emelianov
 """
 
-from collections.abc import Iterable, Iterator, Mapping
+from collections.abc import Iterable, Iterator
 from typing import Final
 
-from botocore.client import BaseClient
-from botocore.config import Config
-from botocore.eventstream import EventStream
-from botocore.paginate import Paginator
-from botocore.response import StreamingBody
-from botocore.waiter import Waiter
-
-from mypy_boto3_builder.import_helpers.import_string import ImportString
+from mypy_boto3_builder.import_helpers.import_string import Import
+from mypy_boto3_builder.postprocessors.aio_imports import replace_imports_with_aio
 from mypy_boto3_builder.postprocessors.base import BasePostprocessor
 from mypy_boto3_builder.structures.argument import Argument
 from mypy_boto3_builder.structures.collection import Collection
@@ -49,53 +43,6 @@ class AioBotocorePostprocessor(BasePostprocessor):
         "can_paginate",
     }
 
-    EXTERNAL_IMPORTS_MAP: Final[Mapping[ExternalImport, ExternalImport]] = {
-        ExternalImport.from_class(StreamingBody): ExternalImport(
-            ImportString("aiobotocore", "response"),
-            "StreamingBody",
-        ),
-        ExternalImport.from_class(EventStream): ExternalImport(
-            ImportString("aiobotocore", "eventstream"),
-            "AioEventStream",
-        ),
-        ExternalImport.from_class(Config): ExternalImport(
-            ImportString("aiobotocore", "config"),
-            "AioConfig",
-        ),
-        ExternalImport.from_class(Waiter): ExternalImport(
-            ImportString("aiobotocore", "waiter"),
-            "AIOWaiter",
-        ),
-        ExternalImport.from_class(Paginator): ExternalImport(
-            ImportString("aiobotocore", "paginate"),
-            "AioPaginator",
-        ),
-        ExternalImport.from_class(BaseClient): ExternalImport(
-            ImportString("aiobotocore", "client"),
-            "AioBaseClient",
-        ),
-        ExternalImport(
-            ImportString("boto3", "resources", "base"), "ServiceResource"
-        ): ExternalImport(
-            ImportString("aioboto3", "resources", "base"),
-            "AIOBoto3ServiceResource",
-        ),
-        ExternalImport(
-            ImportString("boto3", "resources", "collection"), "ResourceCollection"
-        ): ExternalImport(
-            ImportString("aioboto3", "resources", "collection"),
-            "AIOResourceCollection",
-        ),
-        ExternalImport(ImportString("boto3", "dynamodb", "table"), "BatchWriter"): ExternalImport(
-            ImportString("aioboto3", "dynamodb", "table"),
-            "BatchWriter",
-        ),
-        ExternalImport(ImportString("boto3", "dynamodb", "table"), "TableResource"): ExternalImport(
-            ImportString("aioboto3", "dynamodb", "table"),
-            "CustomTableResource",
-        ),
-    }
-
     def process_package(self) -> None:
         """
         Convert all methods to asynchronous.
@@ -220,7 +167,7 @@ def _add_contextmanager_methods(self) -> None:
                         Argument("exc_val", get_optional(ExternalImport.from_class(BaseException))),
                         Argument(
                             "exc_tb",
-                            get_optional(ExternalImport(ImportString("types"), "TracebackType")),
+                            get_optional(ExternalImport(Import.types, "TracebackType")),
                         ),
                     ),
                     return_type=Type.none,
@@ -252,11 +199,4 @@ def _iterate_types(
 
     def _replace_external_imports(self) -> None:
         shallow_type_annotations = self._iterate_types_shallow()
-        for type_annotation in self._iterate_types(shallow_type_annotations):
-            if not isinstance(type_annotation, ExternalImport):
-                continue
-
-            if type_annotation in self.EXTERNAL_IMPORTS_MAP:
-                new_type_annotation = self.EXTERNAL_IMPORTS_MAP[type_annotation]
-                type_annotation.copy_from(new_type_annotation)
-                continue
+        replace_imports_with_aio(self._iterate_types(shallow_type_annotations))

mypy_boto3_builder/utils/type_checks.py

@@ -6,6 +6,7 @@
 
 from typing import TypeGuard
 
+from mypy_boto3_builder.type_annotations.external_import import ExternalImport
 from mypy_boto3_builder.type_annotations.fake_annotation import FakeAnnotation
 from mypy_boto3_builder.type_annotations.type import Type
 from mypy_boto3_builder.type_annotations.type_def_sortable import TypeDefSortable
@@ -57,6 +58,13 @@ def is_type_parent(annotation: FakeAnnotation) -> TypeGuard[TypeParent]:
     return isinstance(annotation, TypeParent)
 
 
+def is_external_import(annotation: FakeAnnotation) -> TypeGuard[ExternalImport]:
+    """
+    Whether type annotation is a ExternalImport.
+    """
+    return isinstance(annotation, ExternalImport)
+
+
 def get_optional(wrapped: FakeAnnotation) -> TypeSubscript:
     """
     Get Optional type annotation.