[mypyc] Fix async or generator methods in traits

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:

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: ...