[ty] Resolve imported pytest fixture exposures

crates/ty_python_semantic/src/types/dedicated/pytest.rs

@@ -9,6 +9,7 @@ use ty_python_core::{ProgramFile, global_scope, place_table, semantic_index, use
 use crate::Db;
 use crate::types::Type;
 use crate::types::function::FunctionDecorators;
+use crate::types::ide_support::resolve_definition_targets;
 use crate::types::infer::{function_known_decorator_flags, function_known_decorators};
 
 /// Resolve the same-file pytest fixtures requested by `parameter`.
@@ -167,24 +168,26 @@ fn bindings_in_provider<'db>(
     for (symbol_id, definitions) in use_def.all_end_of_scope_symbol_bindings() {
         let symbol_name = table.symbol(symbol_id).name();
         for definition in definitions.filter_map(|binding| binding.binding.definition()) {
-            let Some(declaration) = fixture_declaration(db, definition) else {
-                continue;
-            };
-            let Some(exposure) = fixture_exposure(symbol_name, declaration) else {
-                continue;
-            };
-            if exposure.name != request.name
-                || exposure.declaration.definition == request.owner
-                || bindings
-                    .iter()
-                    .any(|binding: &FixtureBinding<'db>| binding.fixture == definition)
-            {
-                continue;
+            for definition in resolve_definition_targets(db, definition, symbol_name) {
+                let Some(declaration) = fixture_declaration(db, definition) else {
+                    continue;
+                };
+                let Some(exposure) = fixture_exposure(symbol_name, declaration) else {
+                    continue;
+                };
+                if exposure.name != request.name
+                    || exposure.declaration.definition == request.owner
+                    || bindings
+                        .iter()
+                        .any(|binding: &FixtureBinding<'db>| binding.fixture == definition)
+                {
+                    continue;
+                }
+                bindings.push(FixtureBinding {
+                    request: request.definition,
+                    fixture: definition,
+                });
             }
-            bindings.push(FixtureBinding {
-                request: request.definition,
-                fixture: definition,
-            });
         }
     }
 
@@ -612,6 +615,140 @@ def test_use(value): ...
         Ok(())
     }
 
+    #[test]
+    fn resolves_imported_fixture_exposures() -> Result<()> {
+        let db = pytest_db_with_files(&[
+            (
+                "/src/fixtures.py",
+                r#"
+import pytest
+
+@pytest.fixture
+def resource(): ...
+
+@pytest.fixture(name="public_name")
+def implementation(): ...
+"#,
+            ),
+            (
+                "/src/reexports.py",
+                r#"
+from fixtures import resource as middle
+"#,
+            ),
+            (
+                "/src/star_fixtures.py",
+                r#"
+import pytest
+
+@pytest.fixture
+def star_fixture(): ...
+"#,
+            ),
+            (
+                "/src/test_example.py",
+                r#"
+from fixtures import resource as direct_alias, implementation
+from reexports import middle as chained
+from star_fixtures import *
+
+def test_use(
+    direct_alias,
+    chained,
+    public_name,
+    implementation,
+    resource,
+    star_fixture,
+): ...
+"#,
+            ),
+        ])?;
+
+        assert_eq!(fixture_names(&db, "test_use", "direct_alias"), ["resource"]);
+        assert_eq!(fixture_names(&db, "test_use", "chained"), ["resource"]);
+        assert_eq!(
+            fixture_names(&db, "test_use", "public_name"),
+            ["implementation"]
+        );
+        assert_eq!(
+            fixture_names(&db, "test_use", "star_fixture"),
+            ["star_fixture"]
+        );
+        assert!(fixture_names(&db, "test_use", "implementation").is_empty());
+        assert!(fixture_names(&db, "test_use", "resource").is_empty());
+        Ok(())
+    }
+
+    #[test]
+    fn ignores_overwritten_imported_fixture_exposures() -> Result<()> {
+        let db = pytest_db_with_files(&[
+            (
+                "/src/fixtures.py",
+                r#"
+import pytest
+
+@pytest.fixture
+def resource(): ...
+"#,
+            ),
+            (
+                "/src/test_example.py",
+                r#"
+from fixtures import resource
+
+resource = object()
+
+def test_use(resource): ...
+"#,
+            ),
+        ])?;
+
+        assert!(fixture_names(&db, "test_use", "resource").is_empty());
+        Ok(())
+    }
+
+    #[test]
+    fn preserves_conditional_imported_fixture_definitions() -> Result<()> {
+        let db = pytest_db_with_files(&[
+            (
+                "/src/first.py",
+                r#"
+import pytest
+
+@pytest.fixture
+def first(): ...
+"#,
+            ),
+            (
+                "/src/second.py",
+                r#"
+import pytest
+
+@pytest.fixture
+def second(): ...
+"#,
+            ),
+            (
+                "/src/test_example.py",
+                r#"
+flag: bool
+
+if flag:
+    from first import first as resource
+else:
+    from second import second as resource
+
+def test_use(resource): ...
+"#,
+            ),
+        ])?;
+
+        let mut fixtures = fixture_names(&db, "test_use", "resource");
+        fixtures.sort();
+        assert_eq!(fixtures, ["first", "second"]);
+        Ok(())
+    }
+
     fn fixture_names(db: &TestDb, function: &str, parameter: &str) -> Vec<String> {
         let parameter = parameter_definition(db, function, parameter);
         fixture_bindings_for_parameter(db, parameter)
@@ -679,7 +816,11 @@ def test_use(value): ...
     }
 
     fn pytest_db(path: &'static str, source: &'static str) -> Result<TestDb> {
-        TestDbBuilder::new()
+        pytest_db_with_files(&[(path, source)])
+    }
+
+    fn pytest_db_with_files(files: &[(&'static str, &'static str)]) -> Result<TestDb> {
+        let mut builder = TestDbBuilder::new()
             .with_site_packages()
             .with_file("/site-packages/_pytest/__init__.pyi", "")
             .with_file(
@@ -713,8 +854,10 @@ class MarkGenerator:
 
 mark: MarkGenerator
 "#,
-            )
-            .with_file(path, source)
-            .build()
+            );
+        for (path, source) in files {
+            builder = builder.with_file(*path, source);
+        }
+        builder.build()
     }
 }

crates/ty_python_semantic/src/types/ide_support.rs

@@ -1289,6 +1289,25 @@ pub fn definitions_for_imported_symbol<'db>(
     )
 }
 
+/// Resolve a binding to its canonical definitions, following imports recursively.
+pub(crate) fn resolve_definition_targets<'db>(
+    db: &'db dyn Db,
+    definition: Definition<'db>,
+    symbol_name: &str,
+) -> Vec<Definition<'db>> {
+    let env = ProgramEnvironment::from_definition(definition);
+    resolve_definition(
+        db,
+        &env,
+        definition,
+        Some(symbol_name),
+        ImportAliasResolution::ResolveAliases,
+    )
+    .into_iter()
+    .filter_map(|resolved| resolved.definition())
+    .collect()
+}
+
 /// Returns the definition and overload co-definitions for a function declaration.
 ///
 /// For overloaded functions this includes sibling overload declarations and the