diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 0f7cc7e3b3c5..9f3c7fc6f270 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -842,6 +842,10 @@ def adjust_generator_classes_of_methods(mapper: Mapper) -> None: if subcls is None: # Override could be of a different type, so we can't make assumptions. precise_ret_type = False + elif class_ir.is_trait: + # Give up on traits. We could possibly have an abstract base class + # for generator return types to make this use precise types. + precise_ret_type = False else: for s in subcls: if name in s.method_decls: diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index cf063310fd89..718325b8b7b6 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -1297,6 +1297,8 @@ from typing import final, Coroutine, Any, TypeVar import asyncio +from mypy_extensions import trait + class Base1: async def foo(self) -> int: return 1 @@ -1363,5 +1365,22 @@ def test_override_non_async() -> None: assert asyncio.run(base3_foo(Base3())) == 7 assert asyncio.run(base3_foo(Derived3())) == 8 +class Base4: pass + +@trait +class TraitBase: + async def foo(self, value: int) -> int: + raise NotImplementedError() + +class DerivedFromTrait(Base4, TraitBase): + async def foo(self, value: int) -> int: + return value + 3 + +async def trait_foo(o: TraitBase, x: int) -> int: + return await o.foo(x) + +def test_override_trait() -> None: + assert asyncio.run(trait_foo(DerivedFromTrait(), 7)) == 10 + [file asyncio/__init__.pyi] def run(x: object) -> object: ...