Use dedicated AST nodes on MemberKind (#6374)

## Summary

This PR leverages the unified function definition node to add precise
AST node types to `MemberKind`, which is used to power our docstring
definition tracking (e.g., classes and functions, whether they're
methods or functions or nested functions and so on, whether they have a
docstring, etc.). It was painful to do this in the past because the
function variants needed to support a union anyway, but storing precise
nodes removes like a dozen panics.

No behavior changes -- purely a refactor.

## Test Plan

`cargo test`

crates/ruff_python_semantic/src/analyze/visibility.rs

@@ -1,6 +1,6 @@
 use std::path::Path;
 
-use ruff_python_ast::{self as ast, Decorator, Stmt};
+use ruff_python_ast::{self as ast, Decorator};
 
 use ruff_python_ast::call_path::{collect_call_path, CallPath};
 use ruff_python_ast::helpers::map_callable;
@@ -176,57 +176,40 @@ impl ModuleSource<'_> {
     }
 }
 
-pub(crate) fn function_visibility(stmt: &Stmt) -> Visibility {
-    match stmt {
-        Stmt::FunctionDef(ast::StmtFunctionDef { name, .. }) => {
-            if name.starts_with('_') {
-                Visibility::Private
-            } else {
-                Visibility::Public
-            }
-        }
-        _ => panic!("Found non-FunctionDef in function_visibility"),
+pub(crate) fn function_visibility(function: &ast::StmtFunctionDef) -> Visibility {
+    if function.name.starts_with('_') {
+        Visibility::Private
+    } else {
+        Visibility::Public
     }
 }
 
-pub(crate) fn method_visibility(stmt: &Stmt) -> Visibility {
-    let Stmt::FunctionDef(ast::StmtFunctionDef {
-        name,
-        decorator_list,
-        ..
-    }) = stmt
-    else {
-        panic!("Found non-FunctionDef in method_visibility")
-    };
-
+pub(crate) fn method_visibility(function: &ast::StmtFunctionDef) -> Visibility {
     // Is this a setter or deleter?
-    if decorator_list.iter().any(|decorator| {
+    if function.decorator_list.iter().any(|decorator| {
         collect_call_path(&decorator.expression).is_some_and(|call_path| {
-            call_path.as_slice() == [name, "setter"] || call_path.as_slice() == [name, "deleter"]
+            call_path.as_slice() == [function.name.as_str(), "setter"]
+                || call_path.as_slice() == [function.name.as_str(), "deleter"]
         })
     }) {
         return Visibility::Private;
     }
 
     // Is the method non-private?
-    if !name.starts_with('_') {
+    if !function.name.starts_with('_') {
         return Visibility::Public;
     }
 
     // Is this a magic method?
-    if name.starts_with("__") && name.ends_with("__") {
+    if function.name.starts_with("__") && function.name.ends_with("__") {
         return Visibility::Public;
     }
 
     Visibility::Private
 }
 
-pub(crate) fn class_visibility(stmt: &Stmt) -> Visibility {
-    let Stmt::ClassDef(ast::StmtClassDef { name, .. }) = stmt else {
-        panic!("Found non-ClassDef in class_visibility")
-    };
-
-    if name.starts_with('_') {
+pub(crate) fn class_visibility(class: &ast::StmtClassDef) -> Visibility {
+    if class.name.starts_with('_') {
         Visibility::Private
     } else {
         Visibility::Public

crates/ruff_python_semantic/src/definition.rs

@@ -5,7 +5,8 @@ use std::fmt::Debug;
 use std::ops::Deref;
 
 use ruff_index::{newtype_index, IndexSlice, IndexVec};
-use ruff_python_ast::{self as ast, Stmt};
+use ruff_python_ast::{self as ast, Ranged, Stmt};
+use ruff_text_size::TextRange;
 
 use crate::analyze::visibility::{
     class_visibility, function_visibility, method_visibility, ModuleSource, Visibility,
@@ -23,7 +24,7 @@ impl DefinitionId {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, is_macro::Is)]
 pub enum ModuleKind {
     /// A Python file that represents a module within a package.
     Module,
@@ -57,35 +58,60 @@ impl<'a> Module<'a> {
     }
 }
 
-#[derive(Debug, Copy, Clone)]
-pub enum MemberKind {
+#[derive(Debug, Copy, Clone, is_macro::Is)]
+pub enum MemberKind<'a> {
     /// A class definition within a program.
-    Class,
+    Class(&'a ast::StmtClassDef),
     /// A nested class definition within a program.
-    NestedClass,
+    NestedClass(&'a ast::StmtClassDef),
     /// A function definition within a program.
-    Function,
+    Function(&'a ast::StmtFunctionDef),
     /// A nested function definition within a program.
-    NestedFunction,
+    NestedFunction(&'a ast::StmtFunctionDef),
     /// A method definition within a program.
-    Method,
+    Method(&'a ast::StmtFunctionDef),
 }
 
 /// A member of a Python module.
 #[derive(Debug)]
 pub struct Member<'a> {
-    pub kind: MemberKind,
+    pub kind: MemberKind<'a>,
     pub parent: DefinitionId,
-    pub stmt: &'a Stmt,
 }
 
 impl<'a> Member<'a> {
     /// Return the name of the member.
-    pub fn name(&self) -> Option<&'a str> {
-        match &self.stmt {
-            Stmt::FunctionDef(ast::StmtFunctionDef { name, .. })
-            | Stmt::ClassDef(ast::StmtClassDef { name, .. }) => Some(name),
-            _ => None,
+    pub fn name(&self) -> &str {
+        match self.kind {
+            MemberKind::Class(class) => &class.name,
+            MemberKind::NestedClass(class) => &class.name,
+            MemberKind::Function(function) => &function.name,
+            MemberKind::NestedFunction(function) => &function.name,
+            MemberKind::Method(method) => &method.name,
+        }
+    }
+
+    /// Return the body of the member.
+    pub fn body(&self) -> &[Stmt] {
+        match self.kind {
+            MemberKind::Class(class) => &class.body,
+            MemberKind::NestedClass(class) => &class.body,
+            MemberKind::Function(function) => &function.body,
+            MemberKind::NestedFunction(function) => &function.body,
+            MemberKind::Method(method) => &method.body,
+        }
+    }
+}
+
+impl Ranged for Member<'_> {
+    /// Return the range of the member.
+    fn range(&self) -> TextRange {
+        match self.kind {
+            MemberKind::Class(class) => class.range(),
+            MemberKind::NestedClass(class) => class.range(),
+            MemberKind::Function(function) => function.range(),
+            MemberKind::NestedFunction(function) => function.range(),
+            MemberKind::Method(method) => method.range(),
         }
     }
 }
@@ -103,16 +129,42 @@ impl Definition<'_> {
         matches!(
             self,
             Definition::Member(Member {
-                kind: MemberKind::Method,
+                kind: MemberKind::Method(_),
                 ..
             })
         )
     }
 
+    /// Return the name of the definition.
     pub fn name(&self) -> Option<&str> {
         match self {
             Definition::Module(module) => module.name(),
-            Definition::Member(member) => member.name(),
+            Definition::Member(member) => Some(member.name()),
+        }
+    }
+
+    /// Return the [`ast::StmtFunctionDef`] of the definition, if it's a function definition.
+    pub fn as_function_def(&self) -> Option<&ast::StmtFunctionDef> {
+        match self {
+            Definition::Member(Member {
+                kind:
+                    MemberKind::Function(function)
+                    | MemberKind::NestedFunction(function)
+                    | MemberKind::Method(function),
+                ..
+            }) => Some(function),
+            _ => None,
+        }
+    }
+
+    /// Return the [`ast::StmtClassDef`] of the definition, if it's a class definition.
+    pub fn as_class_def(&self) -> Option<&ast::StmtClassDef> {
+        match self {
+            Definition::Member(Member {
+                kind: MemberKind::Class(class) | MemberKind::NestedClass(class),
+                ..
+            }) => Some(class),
+            _ => None,
         }
     }
 }
@@ -146,55 +198,43 @@ impl<'a> Definitions<'a> {
                 match &definition {
                     Definition::Module(module) => module.source.to_visibility(),
                     Definition::Member(member) => match member.kind {
-                        MemberKind::Class => {
+                        MemberKind::Class(class) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private()
-                                || exports.is_some_and(|exports| {
-                                    member.name().is_some_and(|name| !exports.contains(&name))
-                                })
+                                || exports.is_some_and(|exports| !exports.contains(&member.name()))
                             {
                                 Visibility::Private
                             } else {
-                                class_visibility(member.stmt)
+                                class_visibility(class)
                             }
                         }
-                        MemberKind::NestedClass => {
+                        MemberKind::NestedClass(class) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private()
-                                || matches!(
-                                    parent.definition,
-                                    Definition::Member(Member {
-                                        kind: MemberKind::Function
-                                            | MemberKind::NestedFunction
-                                            | MemberKind::Method,
-                                        ..
-                                    })
-                                )
+                                || parent.definition.as_function_def().is_some()
                             {
                                 Visibility::Private
                             } else {
-                                class_visibility(member.stmt)
+                                class_visibility(class)
                             }
                         }
-                        MemberKind::Function => {
+                        MemberKind::Function(function) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private()
-                                || exports.is_some_and(|exports| {
-                                    member.name().is_some_and(|name| !exports.contains(&name))
-                                })
+                                || exports.is_some_and(|exports| !exports.contains(&member.name()))
                             {
                                 Visibility::Private
                             } else {
-                                function_visibility(member.stmt)
+                                function_visibility(function)
                             }
                         }
-                        MemberKind::NestedFunction => Visibility::Private,
-                        MemberKind::Method => {
+                        MemberKind::NestedFunction(_) => Visibility::Private,
+                        MemberKind::Method(function) => {
                             let parent = &definitions[member.parent];
                             if parent.visibility.is_private() {
                                 Visibility::Private
                             } else {
-                                method_visibility(member.stmt)
+                                method_visibility(function)
                             }
                         }
                     },