Wrap aioboto3 client methods

youtype · Dec 5, 2024 · 0346ac9 · 0346ac9
1 parent ff90cfb
commit 0346ac9
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/mypy_boto3_builder/parsers/parse_wrapper_package.py b/mypy_boto3_builder/parsers/parse_wrapper_package.py
@@ -6,12 +6,14 @@
 
 from collections.abc import Iterable
 
+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.service_name import ServiceName
 from mypy_boto3_builder.structures.types_aioboto3_package import TypesAioBoto3Package
 from mypy_boto3_builder.structures.types_aiobotocore_package import TypesAioBotocorePackage
 from mypy_boto3_builder.structures.types_boto3_package import TypesBoto3Package
+from mypy_boto3_builder.type_annotations.external_import import ExternalImport
 from mypy_boto3_builder.type_annotations.internal_import import InternalImport
 from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
 
@@ -86,7 +88,12 @@ def parse_types_aioboto3_package(
     """
     package = TypesAioBoto3Package(package_data, service_names, version)
     parser = WrapperPackageParser(package)
-    package.session_class.methods.extend(parser.get_session_client_methods())
+    for method in parser.get_session_client_methods():
+        method.return_type = TypeSubscript(
+            ExternalImport(ImportString("botocore", "session"), "ClientCreatorContext"),
+            [method.return_type],
+        )
+        package.session_class.methods.append(method)
     for method in parser.get_session_resource_methods():
         method.return_type = TypeSubscript(
             InternalImport("ResourceCreatorContext", stringify=False),