AttributeError when async with block has mixed return paths and non-return branch is taken

[Chainfire]

Conditions:

• inside a an async context manager block
• a branch must exist that explicitly returns
• another branch must exist that does not explicitly return
• the AttributeError occurs when the branch that does not explicitly return is executed

from types import TracebackType

import asyncio
from typing import Optional, Type, Literal


class DummyContext:
    async def __aenter__(self) -> None:
        pass

    async def __aexit__(
        self,
        t: Optional[Type[BaseException]],
        v: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> Literal[False]:
        return False


async def early_exit_bug() -> None:
    async with DummyContext():
        if False:
            return


asyncio.run(early_exit_bug())

Raises:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "test_mypyc_async_with_bug.py", line XXX, in <module>
    asyncio.run(early_exit_bug())
  File "/usr/lib/python3.12/asyncio/runners.py", line 195, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "test_mypyc_async_with_bug.py", line -1, in early_exit_bug
AttributeError: attribute '__mypyc_temp__4' of 'early_exit_bug_gen' undefined

[Chainfire]

Workarounds: either all or no branches should explicitly return; avoiding returns and using if statements instead works

History: it seems this problem does not occur with v1.15.0

[Chainfire]

Actually, it does exist in v1.15.0 as well, it just doesn't surface with this particular test script. I've tested back all the way to the introduction of 'async with' support, the problem exists in all versions.

[Chainfire]

Some test cases for mypyc/test-data/run-async.test:

[case testAsyncWithMixedReturn]
import asyncio
from typing import Optional, Type, Literal


class AsyncContextManager:
    async def __aenter__(self) -> "AsyncContextManager":
        return self

    async def __aexit__(
        self,
        t: Optional[Type[BaseException]],
        v: Optional[BaseException],
        tb: object,
    ) -> Literal[False]:
        return False


# Simple async functions (generator class)
async def test_gen_1(b: bool) -> bool:
    async with AsyncContextManager():
        if b:
            return b
    return b


async def test_gen_2(b: bool) -> bool:
    async with AsyncContextManager():
        if b:
            return b
        else:
            return b


async def test_gen_3(b: bool) -> bool:
    async with AsyncContextManager():
        if b:
            return b
        else:
            pass
    return b


async def test_gen_4(b: bool) -> bool:
    ret: bool
    async with AsyncContextManager():
        if b:
            ret = b
        else:
            ret = b
    return ret


async def test_gen_5(i: int) -> int:
    async with AsyncContextManager():
        if i == 1:
            return i
        elif i == 2:
            pass
        elif i == 3:
            return i
    return i


async def test_gen_6(i: int) -> int:
    async with AsyncContextManager():
        if i == 1:
            return i
        elif i == 2:
            return i
        elif i == 3:
            return i
    return i


async def test_gen_7(i: int) -> int:
    async with AsyncContextManager():
        if i == 1:
            return i
        elif i == 2:
            return i
        elif i == 3:
            return i
        else:
            return i


# Async functions with nested functions (environment class)
async def test_env_1(b: bool) -> bool:
    def helper() -> bool:
        return True

    async with AsyncContextManager():
        if b:
            return helper()
    return b


async def test_env_2(b: bool) -> bool:
    def helper() -> bool:
        return True

    async with AsyncContextManager():
        if b:
            return helper()
        else:
            return b


async def test_env_3(b: bool) -> bool:
    def helper() -> bool:
        return True

    async with AsyncContextManager():
        if b:
            return helper()
        else:
            pass
    return b


async def test_env_4(b: bool) -> bool:
    def helper() -> bool:
        return True

    ret: bool
    async with AsyncContextManager():
        if b:
            ret = helper()
        else:
            ret = b
    return ret


async def test_env_5(i: int) -> int:
    def helper() -> int:
        return 1

    async with AsyncContextManager():
        if i == 1:
            return helper()
        elif i == 2:
            pass
        elif i == 3:
            return i
    return i


async def test_env_6(i: int) -> int:
    def helper() -> int:
        return 1

    async with AsyncContextManager():
        if i == 1:
            return helper()
        elif i == 2:
            return i
        elif i == 3:
            return i
    return i


async def test_env_7(i: int) -> int:
    def helper() -> int:
        return 1

    async with AsyncContextManager():
        if i == 1:
            return helper()
        elif i == 2:
            return i
        elif i == 3:
            return i
        else:
            return i


async def run_all_tests() -> None:
    # Test simple async functions (generator class)
    # test_env_1: mixed return/no-return
    assert await test_gen_1(True) is True
    assert await test_gen_1(False) is False  # Bug: fails with AttributeError

    # test_gen_2: all branches return
    assert await test_gen_2(True) is True
    assert await test_gen_2(False) is False

    # test_gen_3: mixed return/pass
    assert await test_gen_3(True) is True
    assert await test_gen_3(False) is False  # Bug: fails with AttributeError

    # test_gen_4: no returns in async with
    assert await test_gen_4(True) is True
    assert await test_gen_4(False) is False

    # test_gen_5: multiple branches, some return
    assert await test_gen_5(0) == 0  # Bug: fails with AttributeError
    assert await test_gen_5(1) == 1
    assert await test_gen_5(2) == 2  # Bug: fails with AttributeError
    assert await test_gen_5(3) == 3

    # test_gen_6: all explicit branches return, implicit fallthrough
    assert await test_gen_6(0) == 0  # Bug: fails with AttributeError
    assert await test_gen_6(1) == 1
    assert await test_gen_6(2) == 2
    assert await test_gen_6(3) == 3

    # test_gen_7: all branches return including else
    assert await test_gen_7(0) == 0
    assert await test_gen_7(1) == 1
    assert await test_gen_7(2) == 2
    assert await test_gen_7(3) == 3

    # Test async functions with nested functions (environment class)
    # test_env_1: mixed return/no-return
    assert await test_env_1(True) is True
    assert await test_env_1(False) is False  # Bug: fails with AttributeError

    # test_env_2: all branches return
    assert await test_env_2(True) is True
    assert await test_env_2(False) is False

    # test_env_3: mixed return/pass
    assert await test_env_3(True) is True
    assert await test_env_3(False) is False  # Bug: fails with AttributeError

    # test_env_4: no returns in async with
    assert await test_env_4(True) is True
    assert await test_env_4(False) is False

    # test_env_5: multiple branches, some return
    assert await test_env_5(0) == 0  # Bug: fails with AttributeError
    assert await test_env_5(1) == 1
    assert await test_env_5(2) == 2  # Bug: fails with AttributeError
    assert await test_env_5(3) == 3

    # test_env_6: all explicit branches return, implicit fallthrough
    assert await test_env_6(0) == 0  # Bug: fails with AttributeError
    assert await test_env_6(1) == 1
    assert await test_env_6(2) == 2
    assert await test_env_6(3) == 3

    # test_env_7: all branches return including else
    assert await test_env_7(0) == 0
    assert await test_env_7(1) == 1
    assert await test_env_7(2) == 2
    assert await test_env_7(3) == 3


def test_async_with_mixed_return() -> None:
    """Test async with blocks that have mixed return/non-return branches.

    This test demonstrates a bug where mypyc fails to properly initialize
    spilled temporaries when:
    1. Inside an async with block
    2. Some branches have explicit returns
    3. Other branches don't have returns
    4. The non-return branch is executed
    """
    asyncio.run(run_all_tests())

[file driver.py]
from native import test_async_with_mixed_return
test_async_with_mixed_return()

[file asyncio/__init__.pyi]
def run(x: object) -> object: ...

[BobTheBuidler]

Not sure if this is appropriate, but I'm linking this to #1107 as both problems appear related to mypyc's handling of implicit return statements, albeit in very different contexts

[Chainfire]

After further investigation, this isn't even an async with problem specifically, those are implemented using try/finally. This also reproduces the error:

import asyncio


async def test_mixed_return(b: bool) -> bool:
    try:
        if b:
            return b
    finally:
        pass
    return b


async def main():
    print("")
    
    print("Testing with True (return branch):")
    result1 = await test_mixed_return(True)
    print(f"Result: {result1}")

    print("\nTesting with False (no-return branch):")
    result2 = await test_mixed_return(False)
    print(f"Result: {result2}")


asyncio.run(main())

Requirements:

• try/finally inside an async function (works fine with sync functions)
• at least one path inside the try block issues a return
• at least one path inside the try block does not issue a return
• the path that does not issue a return is taken at runtime
• there are no further await calls necessary for this problem to surface

[Chainfire]

See python/mypy#19361

[JukkaL]

@Chainfire Can we now close this issue?