name-error when using `async with`

[dpcollins-google]

The following code:

async with (await conn_fut) as connection:
                        # Needs to happen prior to reinitialization to clear outstanding waiters.
                        if last_failure is not None:
                            while not self._write_queue.empty():
                                self._write_queue.get_nowait().response_future.set_exception(
                                    last_failure
                                )
                        self._read_queue = asyncio.Queue(maxsize=1)
                        self._write_queue = asyncio.Queue(maxsize=1)
                        await self._reinitializer.reinitialize(connection, last_failure)
                        self._initialized_once.set()
                        bad_retries = 0
                        await self._loop_connection(connection)

Triggers:

File ".../google/cloud/pubsublite/internal/wire/retrying_connection.py", line 97, in _run_loop: Name 'connection' is not defined [name-error]
File ".../google/cloud/pubsublite/internal/wire/retrying_connection.py", line 100, in _run_loop: Name 'connection' is not defined [name-error]

[rchen152]

Here's a self-contained example of the issue:

# connection.pyi
import abc
from typing import Any, AsyncContextManager, Coroutine

class Connection(AsyncContextManager): ...

class ConnectionFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def new(self) -> Coroutine[Any, Any, Connection]: ...

# retrying_connection.py
from typing import Any
from connection import ConnectionFactory

class RetryingConnection:

    _connection_factory: ConnectionFactory
    _reinitializer: Any

    async def _run_loop(self):
        conn_fut = self._connection_factory.new()
        async with (await conn_fut) as connection:
          await connection

Running pytype-single retrying_connection.py produces:

File "retrying_connection.py", line 12, in _run_loop: Name 'connection' is not defined [name-error]

[rchen152]

When I run with -v4, I see this suspicious debug output:

INFO:pytype.vm  >  ## 11: 12: YIELD_FROM 
[...]
INFO:pytype.vm  >  ## 11: 13: SETUP_ASYNC_WITH  30
[...]
INFO:pytype.vm  >  ## 11: 14: STORE_FAST  connection
--> DEBUG:pytype.attribute Setting connection to the 0 values in <Variable v604: 0 choices> <-- ZERO CHOICES?
[...]
INFO:pytype.vm  >  ## 12: 15: LOAD_FAST  connection
INFO:pytype.errors Added error to log: name-error