diff --git a/jedi/plugins/pytest.py b/jedi/plugins/pytest.py index 0e196c72..c78bdb4f 100644 --- a/jedi/plugins/pytest.py +++ b/jedi/plugins/pytest.py @@ -31,7 +31,15 @@ def execute(callback): def infer_anonymous_param(func): def get_returns(value): if value.tree_node.annotation is not None: - return value.execute_with_values() + result = value.execute_with_values() + if any(v.name.get_qualified_names(include_module_names=True) + == ('typing', 'Generator') + for v in result): + return ValueSet.from_sets( + v.py__getattribute__('__next__').execute_annotation() + for v in result + ) + return result # In pytest we need to differentiate between generators and normal # returns. diff --git a/test/completion/pytest.py b/test/completion/pytest.py index 3c648c6c..a900dcda 100644 --- a/test/completion/pytest.py +++ b/test/completion/pytest.py @@ -1,3 +1,5 @@ +from typing import Generator + import pytest from pytest import fixture @@ -169,3 +171,15 @@ def test_inheritance_fixture(inheritance_fixture, caplog): @pytest.fixture def caplog(caplog): yield caplog + +# ----------------- +# Generator with annotation +# ----------------- + +@pytest.fixture +def with_annot() -> Generator[float, None, None]: + pass + +def test_with_annot(inheritance_fixture, with_annot): + #? float() + with_annot