Respect __all__ imports when determining definition visibility (#4357)

crates/ruff/resources/test/fixtures/pydocstyle/all.py

@@ -0,0 +1,18 @@
+def public_func():
+    pass
+
+
+def private_func():
+    pass
+
+
+class PublicClass:
+    class PublicNestedClass:
+        pass
+
+
+class PrivateClass:
+    pass
+
+
+__all__ = ("public_func", "PublicClass")

crates/ruff/src/checkers/ast/mod.rs

@@ -28,7 +28,7 @@ use ruff_python_semantic::binding::{
     Importation, StarImportation, SubmoduleImportation,
 };
 use ruff_python_semantic::context::Context;
-use ruff_python_semantic::definition::{Module, ModuleKind};
+use ruff_python_semantic::definition::{ContextualizedDefinition, Module, ModuleKind};
 use ruff_python_semantic::node::NodeId;
 use ruff_python_semantic::scope::{ClassDef, FunctionDef, Lambda, Scope, ScopeId, ScopeKind};
 use ruff_python_stdlib::builtins::{BUILTINS, MAGIC_GLOBALS};
@@ -5393,10 +5393,23 @@ impl<'a> Checker<'a> {
             return;
         }
 
-        let mut overloaded_name: Option<String> = None;
-
+        // Compute visibility of all definitions.
+        let global_scope = self.ctx.global_scope();
+        let exports: Option<&[&str]> = global_scope
+            .get("__all__")
+            .map(|index| &self.ctx.bindings[*index])
+            .and_then(|binding| match &binding.kind {
+                BindingKind::Export(Export { names }) => Some(names.as_slice()),
+                _ => None,
+            });
         let definitions = std::mem::take(&mut self.ctx.definitions);
-        for (definition, visibility) in definitions.iter() {
+
+        let mut overloaded_name: Option<String> = None;
+        for ContextualizedDefinition {
+            definition,
+            visibility,
+        } in definitions.resolve(exports).iter()
+        {
             let docstring = docstrings::extraction::extract_docstring(definition);
 
             // flake8-annotations
@@ -5415,7 +5428,9 @@ impl<'a> Checker<'a> {
                 }) {
                     self.diagnostics
                         .extend(flake8_annotations::rules::definition(
-                            self, definition, visibility,
+                            self,
+                            definition,
+                            *visibility,
                         ));
                 }
                 overloaded_name = flake8_annotations::helpers::overloaded_name(self, definition);
@@ -5442,7 +5457,7 @@ impl<'a> Checker<'a> {
 
                 // Extract a `Docstring` from a `Definition`.
                 let Some(expr) = docstring else {
-                    pydocstyle::rules::not_missing(self, definition, visibility);
+                    pydocstyle::rules::not_missing(self, definition, *visibility);
                     continue;
                 };
 

crates/ruff/src/rules/pydocstyle/mod.rs

@@ -168,4 +168,23 @@ mod tests {
         assert_messages!(diagnostics);
         Ok(())
     }
+
+    #[test]
+    fn all() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("pydocstyle/all.py"),
+            &settings::Settings::for_rules([
+                Rule::UndocumentedPublicModule,
+                Rule::UndocumentedPublicClass,
+                Rule::UndocumentedPublicMethod,
+                Rule::UndocumentedPublicFunction,
+                Rule::UndocumentedPublicPackage,
+                Rule::UndocumentedMagicMethod,
+                Rule::UndocumentedPublicNestedClass,
+                Rule::UndocumentedPublicInit,
+            ]),
+        )?;
+        assert_messages!(diagnostics);
+        Ok(())
+    }
 }

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__all.snap

@@ -0,0 +1,34 @@
+---
+source: crates/ruff/src/rules/pydocstyle/mod.rs
+---
+all.py:1:1: D100 Missing docstring in public module
+  |
+1 | def public_func():
+  |  D100
+2 |     pass
+  |
+
+all.py:1:5: D103 Missing docstring in public function
+  |
+1 | def public_func():
+  |     ^^^^^^^^^^^ D103
+2 |     pass
+  |
+
+all.py:9:7: D101 Missing docstring in public class
+   |
+ 9 | class PublicClass:
+   |       ^^^^^^^^^^^ D101
+10 |     class PublicNestedClass:
+11 |         pass
+   |
+
+all.py:10:11: D106 Missing docstring in public nested class
+   |
+10 | class PublicClass:
+11 |     class PublicNestedClass:
+   |           ^^^^^^^^^^^^^^^^^ D106
+12 |         pass
+   |
+
+

crates/ruff_python_semantic/src/definition.rs

@@ -2,11 +2,10 @@
 //! can be documented, such as a module, class, or function.
 
 use std::fmt::Debug;
-use std::iter::FusedIterator;
 use std::num::TryFromIntError;
 use std::ops::{Deref, Index};
 
-use rustpython_parser::ast::Stmt;
+use rustpython_parser::ast::{self, Stmt, StmtKind};
 
 use crate::analyze::visibility::{
     class_visibility, function_visibility, method_visibility, ModuleSource, Visibility,
@@ -86,6 +85,17 @@ pub struct Member<'a> {
     pub stmt: &'a Stmt,
 }
 
+impl<'a> Member<'a> {
+    fn name(&self) -> &'a str {
+        match &self.stmt.node {
+            StmtKind::FunctionDef(ast::StmtFunctionDef { name, .. }) => name,
+            StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef { name, .. }) => name,
+            StmtKind::ClassDef(ast::StmtClassDef { name, .. }) => name,
+            _ => unreachable!("Unexpected member kind: {:?}", self.kind),
+        }
+    }
+}
+
 /// A definition within a Python program.
 #[derive(Debug)]
 pub enum Definition<'a> {
@@ -125,13 +135,76 @@ impl<'a> Definitions<'a> {
         next_id
     }
 
-    /// Iterate over all definitions in a program, with their visibilities.
-    pub fn iter(&self) -> DefinitionsIter {
-        DefinitionsIter {
-            inner: self.0.iter(),
-            definitions: Vec::with_capacity(self.0.len()),
-            visibilities: Vec::with_capacity(self.0.len()),
+    /// Resolve the visibility of each definition in the collection.
+    pub fn resolve(self, exports: Option<&[&str]>) -> ContextualizedDefinitions<'a> {
+        let mut definitions: Vec<ContextualizedDefinition<'a>> = Vec::with_capacity(self.len());
+
+        for definition in self {
+            // Determine the visibility of the next definition, taking into account its parent's
+            // visibility.
+            let visibility = {
+                match &definition {
+                    Definition::Module(module) => module.source.to_visibility(),
+                    Definition::Member(member) => match member.kind {
+                        MemberKind::Class => {
+                            let parent = &definitions[usize::from(member.parent)];
+                            if parent.visibility.is_private()
+                                || exports
+                                    .map_or(false, |exports| !exports.contains(&member.name()))
+                            {
+                                Visibility::Private
+                            } else {
+                                class_visibility(member.stmt)
+                            }
+                        }
+                        MemberKind::NestedClass => {
+                            let parent = &definitions[usize::from(member.parent)];
+                            if parent.visibility.is_private()
+                                || matches!(
+                                    parent.definition,
+                                    Definition::Member(Member {
+                                        kind: MemberKind::Function
+                                            | MemberKind::NestedFunction
+                                            | MemberKind::Method,
+                                        ..
+                                    })
+                                )
+                            {
+                                Visibility::Private
+                            } else {
+                                class_visibility(member.stmt)
+                            }
+                        }
+                        MemberKind::Function => {
+                            let parent = &definitions[usize::from(member.parent)];
+                            if parent.visibility.is_private()
+                                || exports
+                                    .map_or(false, |exports| !exports.contains(&member.name()))
+                            {
+                                Visibility::Private
+                            } else {
+                                function_visibility(member.stmt)
+                            }
+                        }
+                        MemberKind::NestedFunction => Visibility::Private,
+                        MemberKind::Method => {
+                            let parent = &definitions[usize::from(member.parent)];
+                            if parent.visibility.is_private() {
+                                Visibility::Private
+                            } else {
+                                method_visibility(member.stmt)
+                            }
+                        }
+                    },
+                }
+            };
+            definitions.push(ContextualizedDefinition {
+                definition,
+                visibility,
+            });
         }
+
+        ContextualizedDefinitions(definitions)
     }
 }
 
@@ -150,76 +223,26 @@ impl<'a> Deref for Definitions<'a> {
     }
 }
 
-pub struct DefinitionsIter<'a> {
-    inner: std::slice::Iter<'a, Definition<'a>>,
-    definitions: Vec<&'a Definition<'a>>,
-    visibilities: Vec<Visibility>,
-}
+impl<'a> IntoIterator for Definitions<'a> {
+    type Item = Definition<'a>;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
 
-impl<'a> Iterator for DefinitionsIter<'a> {
-    type Item = (&'a Definition<'a>, Visibility);
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
 
-    fn next(&mut self) -> Option<Self::Item> {
-        let definition = self.inner.next()?;
+/// A [`Definition`] in a Python program with its resolved [`Visibility`].
+pub struct ContextualizedDefinition<'a> {
+    pub definition: Definition<'a>,
+    pub visibility: Visibility,
+}
 
-        // Determine the visibility of the next definition, taking into account its parent's
-        // visibility.
-        let visibility = {
-            match definition {
-                Definition::Module(module) => module.source.to_visibility(),
-                Definition::Member(member) => match member.kind {
-                    MemberKind::Class => {
-                        if self.visibilities[usize::from(member.parent)].is_private() {
-                            Visibility::Private
-                        } else {
-                            class_visibility(member.stmt)
-                        }
-                    }
-                    MemberKind::NestedClass => {
-                        if self.visibilities[usize::from(member.parent)].is_private()
-                            || matches!(
-                                self.definitions[usize::from(member.parent)],
-                                Definition::Member(Member {
-                                    kind: MemberKind::Function
-                                        | MemberKind::NestedFunction
-                                        | MemberKind::Method,
-                                    ..
-                                })
-                            )
-                        {
-                            Visibility::Private
-                        } else {
-                            class_visibility(member.stmt)
-                        }
-                    }
-                    MemberKind::Function => {
-                        if self.visibilities[usize::from(member.parent)].is_private() {
-                            Visibility::Private
-                        } else {
-                            function_visibility(member.stmt)
-                        }
-                    }
-                    MemberKind::NestedFunction => Visibility::Private,
-                    MemberKind::Method => {
-                        if self.visibilities[usize::from(member.parent)].is_private() {
-                            Visibility::Private
-                        } else {
-                            method_visibility(member.stmt)
-                        }
-                    }
-                },
-            }
-        };
-        self.definitions.push(definition);
-        self.visibilities.push(visibility);
-
-        Some((definition, visibility))
-    }
+/// A collection of [`Definition`] structs in a Python program with resolved [`Visibility`].
+pub struct ContextualizedDefinitions<'a>(Vec<ContextualizedDefinition<'a>>);
 
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        self.inner.size_hint()
+impl<'a> ContextualizedDefinitions<'a> {
+    pub fn iter(&self) -> impl Iterator<Item = &ContextualizedDefinition<'a>> {
+        self.0.iter()
     }
 }
-
-impl FusedIterator for DefinitionsIter<'_> {}
-impl ExactSizeIterator for DefinitionsIter<'_> {}