diff --git a/mypy_boto3_builder/parsers/parse_wrapper_package.py b/mypy_boto3_builder/parsers/parse_wrapper_package.py index e17568f9..5d208acf 100644 --- 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),