Delay computation of Definition visibility (#4339)

crates/ruff_python_semantic/src/analyze/visibility.rs

@@ -7,25 +7,12 @@ use ruff_python_ast::helpers::map_callable;
 
 use crate::context::Context;
 
-#[derive(Debug, Clone, Copy)]
-pub enum Modifier {
-    Module,
-    Class,
-    Function,
-}
-
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, is_macro::Is)]
 pub enum Visibility {
     Public,
     Private,
 }
 
-#[derive(Debug, Clone, Copy)]
-pub struct VisibleScope {
-    pub modifier: Modifier,
-    pub visibility: Visibility,
-}
-
 /// Returns `true` if a function is a "static method".
 pub fn is_staticmethod(ctx: &Context, decorator_list: &[Expr]) -> bool {
     decorator_list.iter().any(|expr| {
@@ -138,29 +125,45 @@ fn stem(path: &str) -> &str {
     }
 }
 
-/// Return the `Visibility` of the Python file at `Path` based on its name.
-pub fn module_visibility(module_path: Option<&[String]>, path: &Path) -> Visibility {
-    if let Some(module_path) = module_path {
-        if module_path.iter().any(|m| is_private_module(m)) {
-            return Visibility::Private;
-        }
-    } else {
-        // When module_path is None, path is a script outside a package, so just
-        // check to see if the module name itself is private.
-        // Ex) `_foo.py` (but not `__init__.py`)
-        let mut components = path.iter().rev();
-        if let Some(filename) = components.next() {
-            let module_name = filename.to_string_lossy();
-            let module_name = stem(&module_name);
-            if is_private_module(module_name) {
-                return Visibility::Private;
-            }
-        }
-    }
-    Visibility::Public
+/// A Python module can either be defined as a module path (i.e., the dot-separated path to the
+/// module) or, if the module can't be resolved, as a file path (i.e., the path to the file defining
+/// the module).
+#[derive(Debug)]
+pub enum ModuleSource<'a> {
+    /// A module path is a dot-separated path to the module.
+    Path(&'a [String]),
+    /// A file path is the path to the file defining the module, often a script outside of a
+    /// package.
+    File(&'a Path),
 }
 
-pub fn function_visibility(stmt: &Stmt) -> Visibility {
+impl ModuleSource<'_> {
+    /// Return the `Visibility` of the module.
+    pub(crate) fn to_visibility(&self) -> Visibility {
+        match self {
+            Self::Path(path) => {
+                if path.iter().any(|m| is_private_module(m)) {
+                    return Visibility::Private;
+                }
+            }
+            Self::File(path) => {
+                // Check to see if the filename itself indicates private visibility.
+                // Ex) `_foo.py` (but not `__init__.py`)
+                let mut components = path.iter().rev();
+                if let Some(filename) = components.next() {
+                    let module_name = filename.to_string_lossy();
+                    let module_name = stem(&module_name);
+                    if is_private_module(module_name) {
+                        return Visibility::Private;
+                    }
+                }
+            }
+        }
+        Visibility::Public
+    }
+}
+
+pub(crate) fn function_visibility(stmt: &Stmt) -> Visibility {
     match &stmt.node {
         StmtKind::FunctionDef(ast::StmtFunctionDef { name, .. })
         | StmtKind::AsyncFunctionDef(ast::StmtAsyncFunctionDef { name, .. }) => {
@@ -174,7 +177,7 @@ pub fn function_visibility(stmt: &Stmt) -> Visibility {
     }
 }
 
-pub fn method_visibility(stmt: &Stmt) -> Visibility {
+pub(crate) fn method_visibility(stmt: &Stmt) -> Visibility {
     match &stmt.node {
         StmtKind::FunctionDef(ast::StmtFunctionDef {
             name,
@@ -212,7 +215,7 @@ pub fn method_visibility(stmt: &Stmt) -> Visibility {
     }
 }
 
-pub fn class_visibility(stmt: &Stmt) -> Visibility {
+pub(crate) fn class_visibility(stmt: &Stmt) -> Visibility {
     match &stmt.node {
         StmtKind::ClassDef(ast::StmtClassDef { name, .. }) => {
             if name.starts_with('_') {

crates/ruff_python_semantic/src/context.rs

@@ -11,11 +11,11 @@ use ruff_python_ast::typing::AnnotationKind;
 use ruff_python_stdlib::path::is_python_stub_file;
 use ruff_python_stdlib::typing::TYPING_EXTENSIONS;
 
-use crate::analyze::visibility::{module_visibility, Modifier, VisibleScope};
 use crate::binding::{
     Binding, BindingId, BindingKind, Bindings, Exceptions, ExecutionContext, FromImportation,
     Importation, SubmoduleImportation,
 };
+use crate::definition::{Definition, DefinitionId, Definitions, Member, Module};
 use crate::node::{NodeId, Nodes};
 use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
 
@@ -24,6 +24,7 @@ use crate::scope::{Scope, ScopeId, ScopeKind, Scopes};
 pub struct Snapshot {
     scope_id: ScopeId,
     stmt_id: Option<NodeId>,
+    definition_id: DefinitionId,
     in_annotation: bool,
     in_type_checking_block: bool,
 }
@@ -31,7 +32,7 @@ pub struct Snapshot {
 #[allow(clippy::struct_excessive_bools)]
 pub struct Context<'a> {
     pub typing_modules: &'a [String],
-    pub module_path: Option<Vec<String>>,
+    pub module_path: Option<&'a [String]>,
     // Stack of all visited statements, along with the identifier of the current statement.
     pub stmts: Nodes<'a>,
     pub stmt_id: Option<NodeId>,
@@ -41,6 +42,10 @@ pub struct Context<'a> {
     pub scopes: Scopes<'a>,
     pub scope_id: ScopeId,
     pub dead_scopes: Vec<ScopeId>,
+    // Stack of all definitions created in any scope, at any point in execution, along with the
+    // identifier of the current definition.
+    pub definitions: Definitions<'a>,
+    pub definition_id: DefinitionId,
     // A stack of all bindings created in any scope, at any point in execution.
     pub bindings: Bindings<'a>,
     // Map from binding index to indexes of bindings that shadow it in other scopes.
@@ -49,7 +54,6 @@ pub struct Context<'a> {
     pub body: &'a [Stmt],
     pub body_index: usize,
     // Internal, derivative state.
-    pub visible_scope: VisibleScope,
     pub in_annotation: bool,
     pub in_type_definition: bool,
     pub in_deferred_string_type_definition: Option<AnnotationKind>,
@@ -67,29 +71,22 @@ pub struct Context<'a> {
 }
 
 impl<'a> Context<'a> {
-    pub fn new(
-        typing_modules: &'a [String],
-        path: &'a Path,
-        module_path: Option<Vec<String>>,
-    ) -> Self {
-        let visibility = module_visibility(module_path.as_deref(), path);
+    pub fn new(typing_modules: &'a [String], path: &'a Path, module: Module<'a>) -> Self {
         Self {
             typing_modules,
-            module_path,
+            module_path: module.path(),
             stmts: Nodes::default(),
             stmt_id: None,
+            exprs: Vec::default(),
             scopes: Scopes::default(),
             scope_id: ScopeId::global(),
             dead_scopes: Vec::default(),
+            definitions: Definitions::for_module(module),
+            definition_id: DefinitionId::module(),
             bindings: Bindings::default(),
             shadowed_bindings: IntMap::default(),
-            exprs: Vec::default(),
             body: &[],
             body_index: 0,
-            visible_scope: VisibleScope {
-                modifier: Modifier::Module,
-                visibility,
-            },
             in_annotation: false,
             in_type_definition: false,
             in_deferred_string_type_definition: None,
@@ -347,6 +344,19 @@ impl<'a> Context<'a> {
             .expect("Attempted to pop without scope");
     }
 
+    /// Push a [`Member`] onto the stack.
+    pub fn push_definition(&mut self, definition: Member<'a>) {
+        self.definition_id = self.definitions.push_member(definition);
+    }
+
+    /// Pop the current [`Member`] off the stack.
+    pub fn pop_definition(&mut self) {
+        let Definition::Member(member) = &self.definitions[self.definition_id] else {
+            panic!("Attempted to pop without member definition");
+        };
+        self.definition_id = member.parent;
+    }
+
     /// Return the current `Stmt`.
     pub fn stmt(&self) -> &'a Stmt {
         let node_id = self.stmt_id.expect("No current statement");
@@ -450,6 +460,7 @@ impl<'a> Context<'a> {
         Snapshot {
             scope_id: self.scope_id,
             stmt_id: self.stmt_id,
+            definition_id: self.definition_id,
             in_annotation: self.in_annotation,
             in_type_checking_block: self.in_type_checking_block,
         }
@@ -460,11 +471,13 @@ impl<'a> Context<'a> {
         let Snapshot {
             scope_id,
             stmt_id,
+            definition_id,
             in_annotation,
             in_type_checking_block,
         } = snapshot;
         self.scope_id = scope_id;
         self.stmt_id = stmt_id;
+        self.definition_id = definition_id;
         self.in_annotation = in_annotation;
         self.in_type_checking_block = in_type_checking_block;
     }

crates/ruff_python_semantic/src/definition.rs

@@ -0,0 +1,225 @@
+//! Definitions within a Python program. In this context, a "definition" is any named entity that
+//! 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 crate::analyze::visibility::{
+    class_visibility, function_visibility, method_visibility, ModuleSource, Visibility,
+};
+
+/// Id uniquely identifying a definition in a program.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
+pub struct DefinitionId(u32);
+
+impl DefinitionId {
+    /// Returns the ID for the module definition.
+    #[inline]
+    pub const fn module() -> Self {
+        DefinitionId(0)
+    }
+}
+
+impl TryFrom<usize> for DefinitionId {
+    type Error = TryFromIntError;
+
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        Ok(Self(u32::try_from(value)?))
+    }
+}
+
+impl From<DefinitionId> for usize {
+    fn from(value: DefinitionId) -> Self {
+        value.0 as usize
+    }
+}
+
+#[derive(Debug)]
+pub enum ModuleKind {
+    /// A Python file that represents a module within a package.
+    Module,
+    /// A Python file that represents the root of a package (i.e., an `__init__.py` file).
+    Package,
+}
+
+/// A Python module.
+#[derive(Debug)]
+pub struct Module<'a> {
+    pub kind: ModuleKind,
+    pub source: ModuleSource<'a>,
+    pub python_ast: &'a [Stmt],
+}
+
+impl<'a> Module<'a> {
+    pub fn path(&self) -> Option<&'a [String]> {
+        if let ModuleSource::Path(path) = self.source {
+            Some(path)
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+pub enum MemberKind {
+    /// A class definition within a program.
+    Class,
+    /// A nested class definition within a program.
+    NestedClass,
+    /// A function definition within a program.
+    Function,
+    /// A nested function definition within a program.
+    NestedFunction,
+    /// A method definition within a program.
+    Method,
+}
+
+/// A member of a Python module.
+#[derive(Debug)]
+pub struct Member<'a> {
+    pub kind: MemberKind,
+    pub parent: DefinitionId,
+    pub stmt: &'a Stmt,
+}
+
+/// A definition within a Python program.
+#[derive(Debug)]
+pub enum Definition<'a> {
+    Module(Module<'a>),
+    Member(Member<'a>),
+}
+
+impl Definition<'_> {
+    /// Returns `true` if the [`Definition`] is a method definition.
+    pub const fn is_method(&self) -> bool {
+        matches!(
+            self,
+            Definition::Member(Member {
+                kind: MemberKind::Method,
+                ..
+            })
+        )
+    }
+}
+
+/// The definitions within a Python program indexed by [`DefinitionId`].
+#[derive(Debug, Default)]
+pub struct Definitions<'a>(Vec<Definition<'a>>);
+
+impl<'a> Definitions<'a> {
+    pub fn for_module(definition: Module<'a>) -> Self {
+        Self(vec![Definition::Module(definition)])
+    }
+
+    /// Pushes a new member definition and returns its unique id.
+    ///
+    /// Members are assumed to be pushed in traversal order, such that parents are pushed before
+    /// their children.
+    pub fn push_member(&mut self, member: Member<'a>) -> DefinitionId {
+        let next_id = DefinitionId::try_from(self.0.len()).unwrap();
+        self.0.push(Definition::Member(member));
+        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()),
+        }
+    }
+}
+
+impl<'a> Index<DefinitionId> for Definitions<'a> {
+    type Output = Definition<'a>;
+
+    fn index(&self, index: DefinitionId) -> &Self::Output {
+        &self.0[usize::from(index)]
+    }
+}
+
+impl<'a> Deref for Definitions<'a> {
+    type Target = [Definition<'a>];
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+pub struct DefinitionsIter<'a> {
+    inner: std::slice::Iter<'a, Definition<'a>>,
+    definitions: Vec<&'a Definition<'a>>,
+    visibilities: Vec<Visibility>,
+}
+
+impl<'a> Iterator for DefinitionsIter<'a> {
+    type Item = (&'a Definition<'a>, Visibility);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let definition = self.inner.next()?;
+
+        // 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))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+impl FusedIterator for DefinitionsIter<'_> {}
+impl ExactSizeIterator for DefinitionsIter<'_> {}

crates/ruff_python_semantic/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod analyze;
 pub mod binding;
 pub mod context;
+pub mod definition;
 pub mod node;
 pub mod scope;