Fix wrong unused-import when referenced in __all__ (#2125)

Summary:
Currently if a file contains an import that is only referenced in `__all__` it is incorrectly flagged as unused. This addresses that.

Fixes #2128

Pull Request resolved: #2125

Test Plan: Added `test_diagnostic_import_used_in_all`.

Reviewed By: kinto0

Differential Revision: D91080237

Pulled By: stroxler

fbshipit-source-id: 79b9a6d114cc94d4e5519fcbee34ac6bb5eb65e3

Author: maldoinc

pyrefly/lib/binding/bindings.rs

@@ -454,6 +454,10 @@ impl Bindings {
             );
         }
 
+        if let Some(exported_names) = exports.get_explicit_dunder_all_names_iter() {
+            builder.record_used_imports_from_dunder_all_names(exported_names);
+        }
+
         let unused_imports = builder.scopes.collect_module_unused_imports();
         builder.record_unused_imports(unused_imports);
         let scope_trace = builder.scopes.finish();
@@ -708,6 +712,17 @@ impl<'a> BindingsBuilder<'a> {
         self.unused_variables.extend(unused);
     }
 
+    pub fn record_used_imports_from_dunder_all_names<T>(&mut self, dunder_all_names: T)
+    where
+        T: Iterator<Item = &'a Name>,
+    {
+        for name in dunder_all_names {
+            if self.scopes.has_import_name(name) {
+                self.scopes.mark_import_used(name);
+            }
+        }
+    }
+
     pub(crate) fn with_await_context<R>(
         &mut self,
         ctx: AwaitContext,

pyrefly/lib/binding/scope.rs

@@ -1296,6 +1296,15 @@ impl Scopes {
         ScopeTrace(b)
     }
 
+    pub fn has_import_name(&self, name: &Name) -> bool {
+        let module_scope = self.scopes.first();
+
+        match module_scope.scope.kind {
+            ScopeKind::Module => module_scope.scope.imports.contains_key(name),
+            _ => false,
+        }
+    }
+
     pub fn collect_module_unused_imports(&self) -> Vec<UnusedImport> {
         let module_scope = self.scopes.first();
         if !matches!(module_scope.scope.kind, ScopeKind::Module) {

pyrefly/lib/export/exports.rs

@@ -168,6 +168,24 @@ impl Exports {
             .contains(name)
     }
 
+    /// Return an iterator with entries in `__all__` that are user-defined or None if `__all__` was not present.
+    pub fn get_explicit_dunder_all_names_iter(&self) -> Option<impl Iterator<Item = &Name>> {
+        match self.0.definitions.dunder_all.kind {
+            DunderAllKind::Specified => Some(
+                self.0
+                    .definitions
+                    .dunder_all
+                    .entries
+                    .iter()
+                    .filter_map(|entry| match entry {
+                        DunderAllEntry::Name(_, name) => Some(name),
+                        _ => None,
+                    }),
+            ),
+            _ => None,
+        }
+    }
+
     /// Returns entries in `__all__` that don't exist in the module's definitions.
     /// Only validates explicitly user-defined `__all__` entries, not synthesized ones.
     /// Returns a vector of (range, name) tuples for invalid entries.

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

@@ -407,6 +407,32 @@ fn test_unused_from_import_diagnostic() {
     interaction.shutdown().unwrap();
 }
 
+#[test]
+fn test_diagnostic_import_used_in_all() {
+    let test_files_root = get_test_files_root();
+    let mut interaction = LspInteraction::new();
+    interaction.set_root(test_files_root.path().to_path_buf());
+    interaction
+        .initialize(InitializeSettings {
+            configuration: Some(Some(json!([
+                {"pyrefly": {"displayTypeErrors": "force-on"}}
+            ]))),
+            ..Default::default()
+        })
+        .unwrap();
+
+    interaction.client.did_open("unused_import_all/__init__.py");
+    interaction
+        .client
+        .diagnostic("unused_import_all/__init__.py")
+        .expect_response(json!({
+            "items": [],
+            "kind": "full"
+        }))
+        .unwrap();
+    interaction.shutdown().unwrap();
+}
+
 #[test]
 fn test_unused_variable_diagnostic() {
     let test_files_root = get_test_files_root();

pyrefly/lib/test/lsp/lsp_interaction/test_files/unused_import_all/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from foo import Bar as Baz, Foo
+
+__all__ = ["Foo", "Baz"]

pyrefly/lib/test/lsp/lsp_interaction/test_files/unused_import_all/foo.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+class Foo: ...
+
+
+class Bar: ...