Track ranges of names inside __all__ definitions

AlexWaygood · AlexWaygood · commit eefcdd4702c3 · 2024-03-22T13:17:16.000Z
diff --git a/crates/ruff_linter/src/checkers/ast/analyze/definitions.rs b/crates/ruff_linter/src/checkers/ast/analyze/definitions.rs
@@ -1,4 +1,4 @@
-use ruff_python_ast::str::raw_contents_range;
+use ruff_python_ast::{str::raw_contents_range, ExprStringLiteral};
 use ruff_text_size::{Ranged, TextRange};
 
 use ruff_python_semantic::{
@@ -93,7 +93,7 @@ pub(crate) fn definitions(checker: &mut Checker) {
     }
 
     // Compute visibility of all definitions.
-    let exports: Option<Vec<&str>> = {
+    let exports: Option<Vec<&ExprStringLiteral>> = {
         checker
             .semantic
             .global_scope()
diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs
@@ -2097,45 +2097,41 @@ impl<'a> Checker<'a> {
     fn visit_exports(&mut self) {
         let snapshot = self.semantic.snapshot();
 
-        let exports: Vec<(&str, TextRange)> = self
+        let exports: Vec<&ast::ExprStringLiteral> = self
             .semantic
             .global_scope()
             .get_all("__all__")
             .map(|binding_id| &self.semantic.bindings[binding_id])
             .filter_map(|binding| match &binding.kind {
-                BindingKind::Export(Export { names }) => {
-                    Some(names.iter().map(|name| (*name, binding.range())))
-                }
+                BindingKind::Export(Export { names }) => Some(names.iter().copied()),
                 _ => None,
             })
             .flatten()
             .collect();
 
-        for (name, range) in exports {
-            if let Some(binding_id) = self.semantic.global_scope().get(name) {
+        for ast::ExprStringLiteral { value, range } in exports {
+            if let Some(binding_id) = self.semantic.global_scope().get(value.to_str()) {
                 // Mark anything referenced in `__all__` as used.
-                // TODO(charlie): `range` here should be the range of the name in `__all__`, not
-                // the range of `__all__` itself.
                 self.semantic
-                    .add_global_reference(binding_id, ExprContext::Load, range);
+                    .add_global_reference(binding_id, ExprContext::Load, *range);
             } else {
                 if self.semantic.global_scope().uses_star_imports() {
                     if self.enabled(Rule::UndefinedLocalWithImportStarUsage) {
                         self.diagnostics.push(Diagnostic::new(
                             pyflakes::rules::UndefinedLocalWithImportStarUsage {
-                                name: (*name).to_string(),
+                                name: value.to_string(),
                             },
-                            range,
+                            *range,
                         ));
                     }
                 } else {
                     if self.enabled(Rule::UndefinedExport) {
                         if !self.path.ends_with("__init__.py") {
                             self.diagnostics.push(Diagnostic::new(
                                 pyflakes::rules::UndefinedExport {
-                                    name: (*name).to_string(),
+                                    name: value.to_string(),
                                 },
-                                range,
+                                *range,
                             ));
                         }
                     }
diff --git a/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F405_F405.py.snap b/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F405_F405.py.snap
@@ -8,12 +8,10 @@ F405.py:5:11: F405 `name` may be undefined, or defined from star imports
   |           ^^^^ F405
   |
 
-F405.py:11:1: F405 `a` may be undefined, or defined from star imports
+F405.py:11:12: F405 `a` may be undefined, or defined from star imports
    |
  9 |     print(name)
 10 | 
 11 | __all__ = ['a']
-   | ^^^^^^^ F405
+   |            ^^^ F405
    |
-
-
diff --git a/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_0.py.snap b/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_0.py.snap
@@ -1,12 +1,10 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-F822_0.py:3:1: F822 Undefined name `b` in `__all__`
+F822_0.py:3:17: F822 Undefined name `b` in `__all__`
   |
 1 | a = 1
 2 | 
 3 | __all__ = ["a", "b"]
-  | ^^^^^^^ F822
+  |                 ^^^ F822
   |
-
-
diff --git a/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_0.pyi.snap b/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_0.pyi.snap
@@ -1,10 +1,10 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-F822_0.pyi:4:1: F822 Undefined name `c` in `__all__`
+F822_0.pyi:4:22: F822 Undefined name `c` in `__all__`
   |
 2 | b: int  # Considered a binding in a `.pyi` stub file, not in a `.py` runtime file
 3 | 
 4 | __all__ = ["a", "b", "c"]  # c is flagged as missing; b is not
-  | ^^^^^^^ F822
+  |                      ^^^ F822
   |
diff --git a/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_1.py.snap b/crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F822_F822_1.py.snap
@@ -1,12 +1,10 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
 ---
-F822_1.py:3:1: F822 Undefined name `b` in `__all__`
+F822_1.py:3:22: F822 Undefined name `b` in `__all__`
   |
 1 | a = 1
 2 | 
 3 | __all__ = list(["a", "b"])
-  | ^^^^^^^ F822
+  |                      ^^^ F822
   |
-
-
diff --git a/crates/ruff_python_ast/src/all.rs b/crates/ruff_python_ast/src/all.rs
@@ -1,7 +1,7 @@
 use bitflags::bitflags;
 
 use crate::helpers::map_subscript;
-use crate::{self as ast, Expr, Stmt};
+use crate::{self as ast, Expr, ExprStringLiteral, Stmt};
 
 bitflags! {
     #[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
@@ -18,14 +18,18 @@ bitflags! {
 /// Extract the names bound to a given __all__ assignment.
 ///
 /// Accepts a closure that determines whether a given name (e.g., `"list"`) is a Python builtin.
-pub fn extract_all_names<F>(stmt: &Stmt, is_builtin: F) -> (Vec<&str>, DunderAllFlags)
+pub fn extract_all_names<F>(stmt: &Stmt, is_builtin: F) -> (Vec<&ExprStringLiteral>, DunderAllFlags)
 where
     F: Fn(&str) -> bool,
 {
-    fn add_to_names<'a>(elts: &'a [Expr], names: &mut Vec<&'a str>, flags: &mut DunderAllFlags) {
+    fn add_to_names<'a>(
+        elts: &'a [Expr],
+        names: &mut Vec<&'a ExprStringLiteral>,
+        flags: &mut DunderAllFlags,
+    ) {
         for elt in elts {
-            if let Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) = elt {
-                names.push(value.to_str());
+            if let Expr::StringLiteral(string_literal) = elt {
+                names.push(string_literal);
             } else {
                 *flags |= DunderAllFlags::INVALID_OBJECT;
             }
@@ -96,7 +100,7 @@ where
         (None, DunderAllFlags::INVALID_FORMAT)
     }
 
-    let mut names: Vec<&str> = vec![];
+    let mut names: Vec<&ExprStringLiteral> = vec![];
     let mut flags = DunderAllFlags::empty();
 
     if let Some(value) = match stmt {
diff --git a/crates/ruff_python_semantic/src/binding.rs b/crates/ruff_python_semantic/src/binding.rs
@@ -5,6 +5,7 @@ use bitflags::bitflags;
 
 use ruff_index::{newtype_index, IndexSlice, IndexVec};
 use ruff_python_ast::name::QualifiedName;
+use ruff_python_ast::ExprStringLiteral;
 use ruff_python_ast::Stmt;
 use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
@@ -370,7 +371,7 @@ impl<'a> FromIterator<Binding<'a>> for Bindings<'a> {
 #[derive(Debug, Clone)]
 pub struct Export<'a> {
     /// The names of the bindings exported via `__all__`.
-    pub names: Box<[&'a str]>,
+    pub names: Box<[&'a ExprStringLiteral]>,
 }
 
 /// A binding for an `import`, keyed on the name to which the import is bound.
diff --git a/crates/ruff_python_semantic/src/definition.rs b/crates/ruff_python_semantic/src/definition.rs
@@ -187,7 +187,10 @@ impl<'a> Definitions<'a> {
     }
 
     /// Resolve the visibility of each definition in the collection.
-    pub fn resolve(self, exports: Option<&[&str]>) -> ContextualizedDefinitions<'a> {
+    pub fn resolve(
+        self,
+        exports: Option<&[&ast::ExprStringLiteral]>,
+    ) -> ContextualizedDefinitions<'a> {
         let mut definitions: IndexVec<DefinitionId, ContextualizedDefinition<'a>> =
             IndexVec::with_capacity(self.len());
 
@@ -201,7 +204,9 @@ impl<'a> Definitions<'a> {
                         MemberKind::Class(class) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private()
-                                || exports.is_some_and(|exports| !exports.contains(&member.name()))
+                                || exports.is_some_and(|exports| {
+                                    !exports.iter().any(|export| export.value == *member.name())
+                                })
                             {
                                 Visibility::Private
                             } else {
@@ -221,7 +226,9 @@ impl<'a> Definitions<'a> {
                         MemberKind::Function(function) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private()
-                                || exports.is_some_and(|exports| !exports.contains(&member.name()))
+                                || exports.is_some_and(|exports| {
+                                    !exports.iter().any(|export| export.value == *member.name())
+                                })
                             {
                                 Visibility::Private
                             } else {

Original file line number	Diff line number	Diff line change
@@ -8,12 +8,10 @@ F405.py:5:11: F405 `name` may be undefined, or defined from star imports
`8`	`8`	`\| ^^^^ F405`
`9`	`9`	`\|`
`10`	`10`
`11`		-F405.py:11:1: F405 `a` may be undefined, or defined from star imports
	`11`	+F405.py:11:12: F405 `a` may be undefined, or defined from star imports
`12`	`12`	`\|`
`13`	`13`	`9 \| print(name)`
`14`	`14`	`10 \|`
`15`	`15`	`11 \| __all__ = ['a']`
`16`		`- \| ^^^^^^^ F405`
	`16`	`+ \| ^^^ F405`
`17`	`17`	`\|`
`18`		`-`
`19`		`-`
Original file line number	Diff line number	Diff line change
`@@ -1,10 +1,10 @@`
`1`	`1`	`---`
`2`	`2`	`source: crates/ruff_linter/src/rules/pyflakes/mod.rs`
`3`	`3`	`---`
`4`		-F822_0.pyi:4:1: F822 Undefined name `c` in `__all__`
	`4`	+F822_0.pyi:4:22: F822 Undefined name `c` in `__all__`
`5`	`5`	`\|`
`6`	`6`	2 \| b: int # Considered a binding in a `.pyi` stub file, not in a `.py` runtime file
`7`	`7`	`3 \|`
`8`	`8`	`4 \| __all__ = ["a", "b", "c"] # c is flagged as missing; b is not`
`9`		`- \| ^^^^^^^ F822`
	`9`	`+ \| ^^^ F822`
`10`	`10`	`\|`