fix Find-references on library/third-party code #1399

pyrefly/lib/lsp/non_wasm/server.rs

@@ -5750,7 +5750,10 @@ impl Server {
                     import_behavior: ImportBehavior::StopAtRenamedImports,
                     ..Default::default()
                 },
-                options: ReferenceOptions::all(params.context.include_declaration),
+                options: ReferenceOptions {
+                    include_interface_declaration: false,
+                    ..ReferenceOptions::all(params.context.include_declaration)
+                },
                 activity_key,
             },
             move |results| {

pyrefly/lib/state/lsp.rs

@@ -321,6 +321,8 @@ impl Default for FindPreference {
 pub struct ReferenceOptions {
     /// Include the definition itself in the results.
     pub include_declaration: bool,
+    /// Include the corresponding stub declaration when declarations are requested.
+    pub include_interface_declaration: bool,
     /// Include call sites that reach the definition implicitly through the constructor
     /// protocol, e.g. `Foo()` as a reference to `Foo.__init__`.
     ///
@@ -336,6 +338,7 @@ impl ReferenceOptions {
     pub fn all(include_declaration: bool) -> Self {
         Self {
             include_declaration,
+            include_interface_declaration: true,
             include_constructor_call_sites: true,
         }
     }
@@ -345,6 +348,7 @@ impl ReferenceOptions {
     pub fn textual_only(include_declaration: bool) -> Self {
         Self {
             include_declaration,
+            include_interface_declaration: true,
             include_constructor_call_sites: false,
         }
     }
@@ -2063,6 +2067,39 @@ impl<'a> Transaction<'a> {
         Some((executable_module, def_range, docstring_range))
     }
 
+    /// Find the definition in the interface used to type-check an executable module.
+    fn find_corresponding_interface_definition(
+        &self,
+        sys_info: SysInfo,
+        definition: &TextRangeWithModule,
+    ) -> Option<TextRangeWithModule> {
+        if definition.module.path().style() != ModuleStyle::Executable {
+            return None;
+        }
+        let source_handle = Handle::new(
+            definition.module.name(),
+            definition.module.path().dupe(),
+            sys_info,
+        );
+        let interface_handle = self
+            .import_handle(&source_handle, definition.module.name(), None)
+            .finding()?;
+        if interface_handle.path().style() != ModuleStyle::Interface {
+            return None;
+        }
+        let context = AttributeContext::from_module(&definition.module, definition.range)?;
+        let name = Name::new(definition.module.code_at(definition.range));
+        let _ = self.get_exports(&interface_handle);
+        let interface_module = self.get_module_info(&interface_handle)?;
+        let ast = self.get_ast(&interface_handle).unwrap_or_else(|| {
+            Ast::parse(interface_module.contents(), interface_module.source_type())
+                .0
+                .into()
+        });
+        let (range, _) = definition_from_executable_ast(ast.as_ref(), &context, &name)?;
+        Some(TextRangeWithModule::new(interface_module, range))
+    }
+
     pub fn key_to_export(
         &self,
         handle: &Handle,
@@ -5259,6 +5296,11 @@ trait RdepTransaction {
     fn module_info(&self, handle: &Handle) -> Option<Module>;
     fn transitive_rdeps(&self, handle: Handle) -> HashSet<Handle>;
     fn run_for_handles(&mut self, handles: &[Handle], require: Require) -> Result<(), Cancelled>;
+    fn find_corresponding_interface_definition(
+        &self,
+        sys_info: SysInfo,
+        definition: &TextRangeWithModule,
+    ) -> Option<TextRangeWithModule>;
     fn local_references_from_definition(
         &self,
         handle: &Handle,
@@ -5288,6 +5330,14 @@ impl<'a> RdepTransaction for Transaction<'a> {
         Ok(())
     }
 
+    fn find_corresponding_interface_definition(
+        &self,
+        sys_info: SysInfo,
+        definition: &TextRangeWithModule,
+    ) -> Option<TextRangeWithModule> {
+        self.find_corresponding_interface_definition(sys_info, definition)
+    }
+
     fn local_references_from_definition(
         &self,
         handle: &Handle,
@@ -5319,6 +5369,15 @@ impl<'a> RdepTransaction for CancellableTransaction<'a> {
         self.run(handles, require, None)
     }
 
+    fn find_corresponding_interface_definition(
+        &self,
+        sys_info: SysInfo,
+        definition: &TextRangeWithModule,
+    ) -> Option<TextRangeWithModule> {
+        self.as_ref()
+            .find_corresponding_interface_definition(sys_info, definition)
+    }
+
     fn local_references_from_definition(
         &self,
         handle: &Handle,
@@ -5489,57 +5548,75 @@ fn find_global_references_from_definition_impl<T: RdepTransaction>(
     definition: TextRangeWithModule,
     options: ReferenceOptions,
 ) -> Result<Vec<(Module, Vec<TextRange>)>, Cancelled> {
-    let candidate_handles =
-        compute_transitive_rdeps_for_definition_impl(transaction, sys_info, &definition)?;
-    if definition_kind.symbol_kind() == Some(SymbolKind::Parameter) {
-        // Keyword argument references require each candidate's AST and bindings to resolve the
-        // callee and refine the argument back to this parameter.
-        transaction.run_for_handles(&candidate_handles, Require::Everything)?;
-    }
-    let results = process_candidate_handles_with_definition_impl(
-        transaction,
-        candidate_handles,
-        &definition,
-        |transaction, handle, patched_definition| {
-            let mut module_refs: Vec<(Module, Vec<TextRange>)> = Vec::new();
+    let interface_definition =
+        transaction.find_corresponding_interface_definition(sys_info, &definition);
+    let definitions = interface_definition
+        .into_iter()
+        .map(|definition| {
+            (
+                definition,
+                ReferenceOptions {
+                    include_declaration: options.include_declaration
+                        && options.include_interface_declaration,
+                    ..options
+                },
+            )
+        })
+        .chain(iter::once((definition, options)));
+    let mut results = Vec::new();
+    for (definition, options) in definitions {
+        let candidate_handles =
+            compute_transitive_rdeps_for_definition_impl(transaction, sys_info, &definition)?;
+        if definition_kind.symbol_kind() == Some(SymbolKind::Parameter) {
+            // Keyword argument references require each candidate's AST and bindings to resolve the
+            // callee and refine the argument back to this parameter.
+            transaction.run_for_handles(&candidate_handles, Require::Everything)?;
+        }
+        results.extend(process_candidate_handles_with_definition_impl(
+            transaction,
+            candidate_handles,
+            &definition,
+            |transaction, handle, patched_definition| {
+                let mut module_refs: Vec<(Module, Vec<TextRange>)> = Vec::new();
 
-            let references = transaction
-                .local_references_from_definition(
-                    handle,
-                    definition_kind.clone(),
-                    patched_definition.range,
-                    &patched_definition.module,
-                    options,
-                )
-                .unwrap_or_default();
-            if !references.is_empty()
-                && let Some(module_info) = transaction.module_info(handle)
-            {
-                module_refs.push((module_info, references));
-            }
+                let references = transaction
+                    .local_references_from_definition(
+                        handle,
+                        definition_kind.clone(),
+                        patched_definition.range,
+                        &patched_definition.module,
+                        options,
+                    )
+                    .unwrap_or_default();
+                if !references.is_empty()
+                    && let Some(module_info) = transaction.module_info(handle)
+                {
+                    module_refs.push((module_info, references));
+                }
 
-            let child_implementations =
-                find_child_implementations_impl(transaction, handle, patched_definition);
-            if !child_implementations.is_empty()
-                && let Some(module_info) = transaction.module_info(handle)
-            {
-                if let Some((_, ranges)) = module_refs
-                    .iter_mut()
-                    .find(|(m, _)| m.path() == module_info.path())
+                let child_implementations =
+                    find_child_implementations_impl(transaction, handle, patched_definition);
+                if !child_implementations.is_empty()
+                    && let Some(module_info) = transaction.module_info(handle)
                 {
-                    ranges.extend(child_implementations);
-                } else {
-                    module_refs.push((module_info, child_implementations));
+                    if let Some((_, ranges)) = module_refs
+                        .iter_mut()
+                        .find(|(m, _)| m.path() == module_info.path())
+                    {
+                        ranges.extend(child_implementations);
+                    } else {
+                        module_refs.push((module_info, child_implementations));
+                    }
                 }
-            }
 
-            if module_refs.is_empty() {
-                None
-            } else {
-                Some(module_refs)
-            }
-        },
-    );
+                if module_refs.is_empty() {
+                    None
+                } else {
+                    Some(module_refs)
+                }
+            },
+        ));
+    }
 
     let mut global_references: Vec<(Module, Vec<TextRange>)> = Vec::new();
     for module_refs in results {

pyrefly/lib/state/lsp_attributes.rs

@@ -28,6 +28,7 @@ pub(crate) struct AttributeContext {
 }
 
 impl AttributeContext {
+    /// Finds a module-level or class-level declaration with this exact name range.
     pub(crate) fn from_module(
         module: &Module,
         target_range: TextRange,
@@ -43,36 +44,36 @@ impl AttributeContext {
         target_range: TextRange,
     ) -> Option<AttributeContext> {
         for stmt in body {
-            match stmt {
-                Stmt::ClassDef(class_def) => {
+            let kind = match stmt {
+                Stmt::ClassDef(class_def) if class_def.name.range == target_range => {
+                    AttributeKind::Class
+                }
+                Stmt::ClassDef(class_def) if class_def.range().contains_range(target_range) => {
                     parents.push(class_def.name.id.clone());
-                    if let Some(ctx) = Self::from_body(&class_def.body, parents, target_range) {
-                        return Some(ctx);
-                    }
-                    parents.pop();
-                    if class_def.range().contains_range(target_range) {
-                        return Some(AttributeContext {
-                            parent_classes: parents.clone(),
-                            kind: AttributeKind::Class,
-                        });
-                    }
+                    return Self::from_body(&class_def.body, parents, target_range);
                 }
-                Stmt::FunctionDef(_) if stmt.range().contains_range(target_range) => {
-                    return Some(AttributeContext {
-                        parent_classes: parents.clone(),
-                        kind: AttributeKind::Function,
-                    });
+                Stmt::FunctionDef(function) if function.name.range == target_range => {
+                    AttributeKind::Function
                 }
-                Stmt::Assign(_) | Stmt::AnnAssign(_)
-                    if stmt.range().contains_range(target_range) =>
+                Stmt::Assign(assign)
+                    if assign.targets.iter().any(
+                        |target| matches!(target, Expr::Name(name) if name.range() == target_range),
+                    ) =>
                 {
-                    return Some(AttributeContext {
-                        parent_classes: parents.clone(),
-                        kind: AttributeKind::Assignment,
-                    });
+                    AttributeKind::Assignment
                 }
-                _ => {}
-            }
+                Stmt::AnnAssign(assign)
+                    if assign.target.range() == target_range
+                        && matches!(assign.target.as_ref(), Expr::Name(_)) =>
+                {
+                    AttributeKind::Assignment
+                }
+                _ => continue,
+            };
+            return Some(AttributeContext {
+                parent_classes: parents.clone(),
+                kind,
+            });
         }
         None
     }

pyrefly/lib/test/lsp/lsp_interaction/references.rs

@@ -683,6 +683,85 @@ fn test_references_for_aliased_import_with_config() {
     interaction.shutdown().unwrap();
 }
 
+#[test]
+fn test_references_from_third_party_source_with_stub() {
+    let root = get_test_files_root();
+    let root_path = root.path().join("references_third_party");
+    let scope_uri = Uri::from_file_path(root_path.clone()).unwrap();
+    let mut interaction = LspInteraction::new_with_args(LspInteractionArgs {
+        args: LspArgs {
+            indexing_mode: IndexingMode::LazyBlocking,
+            ..LspInteractionArgs::default().args
+        },
+        ..Default::default()
+    });
+    interaction.set_root(root_path.clone());
+    interaction
+        .initialize(InitializeSettings {
+            workspace_folders: Some(vec![("test".to_owned(), scope_uri)]),
+            ..Default::default()
+        })
+        .unwrap();
+
+    let library = root_path.join("venv/lib/python3.12/site-packages/third_party_lib.py");
+    let user_code = root_path.join("user.py");
+
+    interaction.client.did_open("user.py");
+    interaction
+        .client
+        .did_open("venv/lib/python3.12/site-packages/third_party_lib.py");
+
+    interaction
+        .client
+        .references(
+            "venv/lib/python3.12/site-packages/third_party_lib.py",
+            0,
+            6,
+            true,
+        )
+        .expect_response(json!([
+            {
+                "range": {"start":{"line":2,"character":11},"end":{"line":2,"character":16}},
+                "uri": Uri::from_file_path(&user_code).unwrap().to_string()
+            },
+            {
+                "range": {"start":{"line":2,"character":23},"end":{"line":2,"character":28}},
+                "uri": Uri::from_file_path(&user_code).unwrap().to_string()
+            },
+            {
+                "range": {"start":{"line":3,"character":12},"end":{"line":3,"character":17}},
+                "uri": Uri::from_file_path(&user_code).unwrap().to_string()
+            },
+            {
+                "range": {"start":{"line":0,"character":6},"end":{"line":0,"character":11}},
+                "uri": Uri::from_file_path(&library).unwrap().to_string()
+            },
+        ]))
+        .unwrap();
+
+    interaction
+        .client
+        .references(
+            "venv/lib/python3.12/site-packages/third_party_lib.py",
+            1,
+            8,
+            true,
+        )
+        .expect_response(json!([
+            {
+                "range": {"start":{"line":4,"character":6},"end":{"line":4,"character":13}},
+                "uri": Uri::from_file_path(&user_code).unwrap().to_string()
+            },
+            {
+                "range": {"start":{"line":1,"character":8},"end":{"line":1,"character":15}},
+                "uri": Uri::from_file_path(&library).unwrap().to_string()
+            },
+        ]))
+        .unwrap();
+
+    interaction.shutdown().unwrap();
+}
+
 #[test]
 fn test_references_after_file_modification_with_config() {
     let root = get_test_files_root();

pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/pyrefly.toml

@@ -0,0 +1 @@
+site_package_path = ["venv/lib/python3.12/site-packages"]