youtype/mypy_boto3_builder

Support `async with` on TableWriter

Closed this issue · 3 comments

Describe your idea

Support use of TableWriter in an async with block.

Code sample

from types_aiobotocore_dynamodb.service_resource import Table

async def write(table: Table) -> None:
    async with table.batch_writer() as batch:
        ...

Currently fails in pyright with:

/tmp/test.py
  /tmp/test.py:4:16 - error: Object of type "Coroutine[Any, Any, BatchWriter]" cannot be used with "with" because it does not implement __aenter__ (reportGeneralTypeIssues)
  /tmp/test.py:4:16 - error: Object of type "Coroutine[Any, Any, BatchWriter]" cannot be used with "with" because it does not implement __aexit__ (reportGeneralTypeIssues)

Additional context

❯ pip freeze  | grep boto
aioboto3==11.1.0
aiobotocore==2.5.0
boto3==1.26.76
botocore==1.29.76
botocore-stubs==1.29.125
types-aioboto3==11.1.0
types-aiobotocore==2.5.0.post2
types-aiobotocore-dynamodb==2.5.0.post1
vemel commented

Hello!

Good idea, resources should support AsyncContextManager interface.

vemel commented

Sorry for a late response.

THis is not directly related to aiobotocore, because you are using Table from aioboto3. So, try to provide aioboto3 type annotations:

from aioboto3.dynamodb.table import CustomTableResource


async def write(table: CustomTableResource) -> None:
    async with table.batch_writer() as batch:
        reveal_type(batch)

Hi 👋

If I use CustomTableResource for table then elsewhere I can no longer use await self.table.get_item(Key=key) without getting:

error: Cannot access member "get_item" for type "CustomTableResource"
    Member "get_item" is unknown (reportGeneralTypeIssues)