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`
This commit is contained in:
Charlie Marsh 2023-08-07 13:17:58 -04:00 committed by GitHub
parent daefa74e9a
commit c439435615
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 316 additions and 414 deletions

View file

@ -1,7 +1,6 @@
//! Extract docstrings from an AST.
use ruff_python_ast::{self as ast, Constant, Expr, Stmt};
use ruff_python_semantic::{Definition, DefinitionId, Definitions, Member, MemberKind};
/// Extract a docstring from a function or class body.
@ -28,62 +27,48 @@ pub(crate) fn docstring_from(suite: &[Stmt]) -> Option<&Expr> {
pub(crate) fn extract_docstring<'a>(definition: &'a Definition<'a>) -> Option<&'a Expr> {
match definition {
Definition::Module(module) => docstring_from(module.python_ast),
Definition::Member(member) => {
if let Stmt::ClassDef(ast::StmtClassDef { body, .. })
| Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) = &member.stmt
{
docstring_from(body)
} else {
None
}
}
Definition::Member(member) => docstring_from(member.body()),
}
}
#[derive(Copy, Clone)]
pub(crate) enum ExtractionTarget {
Class,
Function,
pub(crate) enum ExtractionTarget<'a> {
Class(&'a ast::StmtClassDef),
Function(&'a ast::StmtFunctionDef),
}
/// Extract a `Definition` from the AST node defined by a `Stmt`.
pub(crate) fn extract_definition<'a>(
target: ExtractionTarget,
stmt: &'a Stmt,
target: ExtractionTarget<'a>,
parent: DefinitionId,
definitions: &Definitions<'a>,
) -> Member<'a> {
match target {
ExtractionTarget::Function => match &definitions[parent] {
ExtractionTarget::Function(function) => match &definitions[parent] {
Definition::Module(..) => Member {
parent,
kind: MemberKind::Function,
stmt,
kind: MemberKind::Function(function),
},
Definition::Member(Member {
kind: MemberKind::Class | MemberKind::NestedClass,
kind: MemberKind::Class(_) | MemberKind::NestedClass(_),
..
}) => Member {
parent,
kind: MemberKind::Method,
stmt,
kind: MemberKind::Method(function),
},
Definition::Member(..) => Member {
Definition::Member(_) => Member {
parent,
kind: MemberKind::NestedFunction,
stmt,
kind: MemberKind::NestedFunction(function),
},
},
ExtractionTarget::Class => match &definitions[parent] {
Definition::Module(..) => Member {
ExtractionTarget::Class(class) => match &definitions[parent] {
Definition::Module(_) => Member {
parent,
kind: MemberKind::Class,
stmt,
kind: MemberKind::Class(class),
},
Definition::Member(..) => Member {
Definition::Member(_) => Member {
parent,
kind: MemberKind::NestedClass,
stmt,
kind: MemberKind::NestedClass(class),
},
},
}