Wrap aiobotocore create_client return types in ClientCreatorContext

mypy_boto3_builder/parsers/parse_wrapper_package.py

@@ -59,7 +59,12 @@ def parse_aiobotocore_stubs_package(
     """
     package = TypesAioBotocorePackage(package_data, service_names, version)
     parser = WrapperPackageParser(package)
-    package.session_class.methods.extend(parser.get_session_client_methods("create_client"))
+    for method in parser.get_session_client_methods("create_client"):
+        method.return_type = TypeSubscript(
+            InternalImport("ClientCreatorContext", stringify=False),
+            [method.return_type],
+        )
+        package.session_class.methods.append(method)
     return package
 
 

mypy_boto3_builder/stubs_static/types-aiobotocore/session.pyi

@@ -1,5 +1,5 @@
 from types import TracebackType
-from typing import Any, List, Optional, Type, Union
+from typing import Any, Generic, List, Optional, Type, TypeVar, Union
 
 from aiobotocore.client import AioBaseClient as AioBaseClient
 from aiobotocore.client import AioClientCreator as AioClientCreator
@@ -12,9 +12,11 @@ from botocore.model import ServiceModel
 from botocore.session import EVENT_ALIASES as EVENT_ALIASES
 from botocore.session import Session
 
-class ClientCreatorContext:
+_AioBaseClient = TypeVar("_AioBaseClient", bound="AioBaseClient")
+
+class ClientCreatorContext(Generic[_AioBaseClient]):
     def __init__(self, coro: Any) -> None: ...
-    async def __aenter__(self) -> AioBaseClient: ...
+    async def __aenter__(self) -> _AioBaseClient: ...
     async def __aexit__(
         self,
         exc_type: Optional[Type[BaseException]],
@@ -49,7 +51,7 @@ class AioSession(Session):
         aws_secret_access_key: Optional[str] = ...,
         aws_session_token: Optional[str] = ...,
         config: Optional[Config] = ...,
-    ) -> AioBaseClient: ...
+    ) -> ClientCreatorContext[AioBaseClient]: ...
     async def get_credentials(self) -> Optional[AioCredentials]: ...  # type: ignore [override]
     def set_credentials(
         self, access_key: str, secret_key: str, token: Optional[Any] = ...

mypy_boto3_builder/templates/types-aiobotocore/aiobotocore-stubs/session.pyi.jinja2

@@ -1,5 +1,5 @@
 from types import TracebackType
-from typing import Any, List, Type, Optional
+from typing import Any, Generic, List, Optional, Type, TypeVar, Union
 
 from aiobotocore.client import AioBaseClient as AioBaseClient
 from aiobotocore.client import AioClientCreator as AioClientCreator
@@ -16,9 +16,11 @@ from botocore.session import Session as BotocoreSession
     {{ import_record -}}{{ "\n" -}}
 {% endfor -%}
 
-class ClientCreatorContext:
+_AioBaseClient = TypeVar("_AioBaseClient", bound="AioBaseClient")
+
+class ClientCreatorContext(Generic[_AioBaseClient]):
     def __init__(self, coro: Any) -> None: ...
-    async def __aenter__(self) -> AioBaseClient: ...
+    async def __aenter__(self) -> _AioBaseClient: ...
     async def __aexit__(
         self, 
         exc_type: Optional[Type[BaseException]],