use crate::Python;
use crate::lang_helpers::doc_comment::node_bytes;
use crate::node::{AncestorIter, Ancestors, Cursor, FoldLane, Node};
#[inline]
#[must_use]
pub fn python_is_lambda(node: &Node) -> bool {
matches!(node.kind_id().into(), Python::Lambda | Python::Lambda2)
}
#[inline]
#[must_use]
pub fn python_is_block(node: &Node) -> bool {
matches!(node.kind_id().into(), Python::Block | Python::Block2)
}
const MAX_PARENTHESES: u8 = 200;
#[inline]
#[must_use]
pub fn python_is_bare_statement_value(node: &Node, ancestors: Ancestors<'_, '_>) -> bool {
enclosing_statement(node, ancestors)
.is_some_and(|((statement, _), _)| is_bare_statement(&statement))
}
#[must_use]
pub fn python_is_bare_string_parenthesis(node: &Node, ancestors: Ancestors<'_, '_>) -> bool {
let token = matches!(node.kind_id().into(), Python::LPAREN | Python::RPAREN);
let Some((paren, paren_ancestors)) = (if token {
ancestors.iter(node).next()
} else {
Some((*node, ancestors))
}) else {
return false;
};
if paren.kind_id() != Python::ParenthesizedExpression {
return false;
}
if ancestors.memoises_folds() {
let climb = ancestors.fold_down(
node,
usize::from(token),
FoldLane::BareStringClimb,
0,
|climb, ancestor, _, _| bare_string_climb(climb, ancestor),
);
return (1..=MAX_PARENTHESES + 1).contains(&climb);
}
enclosing_statement(&paren, paren_ancestors)
.is_some_and(|((statement, _), _)| statement_string(&statement).is_some())
}
fn bare_string_climb(climb: u8, ancestor: &Node) -> u8 {
if ancestor.kind_id() != Python::ParenthesizedExpression {
u8::from(statement_string(ancestor).is_some())
} else if climb == 0 {
0
} else {
climb.saturating_add(1).min(MAX_PARENTHESES + 2)
}
}
fn enclosing_statement<'a, 'c>(
node: &Node<'a>,
ancestors: Ancestors<'a, 'c>,
) -> Option<((Node<'a>, Ancestors<'a, 'c>), AncestorIter<'a, 'c>)> {
let mut up = ancestors.iter(node);
let statement = up
.by_ref()
.take(usize::from(MAX_PARENTHESES) + 1)
.find(|(ancestor, _)| ancestor.kind_id() != Python::ParenthesizedExpression)?;
Some((statement, up))
}
fn unparenthesized(node: Node) -> Node {
let mut cursor = node.cursor();
let mut value = node;
for _ in 0..MAX_PARENTHESES {
let Some(inner) = unparenthesized_once(&value, &mut cursor) else {
break;
};
value = inner;
}
value
}
fn unparenthesized_once<'a>(node: &Node<'a>, cursor: &mut Cursor<'a>) -> Option<Node<'a>> {
if node.kind_id() != Python::ParenthesizedExpression {
return None;
}
node.children_with(cursor)
.find(|child| child.is_named() && child.kind_id() != Python::Comment)
}
fn statement_string<'a>(statement: &Node<'a>) -> Option<Node<'a>> {
if !is_bare_statement(statement) {
return None;
}
statement.child(0).map(unparenthesized).filter(|value| {
matches!(
value.kind_id().into(),
Python::String | Python::ConcatenatedString
)
})
}
#[inline]
#[must_use]
pub fn python_is_bare_string(node: &Node, ancestors: Ancestors<'_, '_>) -> bool {
let Some((parent, parent_ancestors)) = ancestors.iter(node).next() else {
return false;
};
if matches!(parent.kind_id().into(), Python::ConcatenatedString) {
return python_is_bare_statement_value(&parent, parent_ancestors);
}
python_is_bare_statement_value(node, ancestors)
}
fn documentation_run_start<'a>(
value: &Node<'a>,
code: &[u8],
ancestors: Ancestors<'a, '_>,
) -> Option<Node<'a>> {
if !is_str_literal(value, code) {
return None;
}
let ((statement, statement_ancestors), up) = enclosing_statement(value, ancestors)?;
let attributes = body_attributes(up, code)?;
let mut first = statement;
for sibling in statement_ancestors
.siblings_before(&statement)
.filter(|sibling| sibling.is_named() && sibling.kind_id() != Python::Comment)
{
if !is_string_statement(&sibling, code) {
return attributes.documented_by(&sibling, code).then_some(first);
}
first = sibling;
}
Some(first)
}
pub fn python_documentation_string_run(
value: &Node,
code: &[u8],
ancestors: Ancestors<'_, '_>,
mut credit: impl FnMut(usize, usize),
) {
let Some(((statement, outer), _)) = enclosing_statement(value, ancestors) else {
return;
};
if outer
.siblings_after(&statement)
.find(|sibling| sibling.is_named() && sibling.kind_id() != Python::Comment)
.is_some_and(|next| is_string_statement(&next, code))
{
return;
}
let Some(first) = documentation_run_start(value, code, ancestors) else {
return;
};
for member in std::iter::once(first)
.chain(outer.siblings_after(&first))
.take_while(|member| member.start_byte() <= statement.start_byte())
.filter(|member| is_string_statement(member, code))
{
credit_documentation_rows(member, &mut credit);
}
}
fn credit_documentation_rows(statement: Node, credit: &mut impl FnMut(usize, usize)) {
let Some(mut value) = statement.child(0) else {
return;
};
let mut cursor = value.cursor();
for _ in 0..MAX_PARENTHESES {
let Some(inner) = unparenthesized_once(&value, &mut cursor) else {
break;
};
for paren in value
.children_with(&mut cursor)
.filter(|child| !child.is_named())
{
credit(paren.start_row(), paren.end_row());
}
value = inner;
}
credit(value.start_row(), value.end_row());
}
enum Attributes<'code> {
Names,
OfSelf(&'code [u8]),
Nothing,
}
impl Attributes<'_> {
fn documented_by(&self, statement: &Node, code: &[u8]) -> bool {
if statement.kind_id() == Python::TypeAliasStatement {
return matches!(self, Self::Names);
}
let Some(target) = assignment_target(statement) else {
return false;
};
match self {
Self::Names => binds(&target, |leaf| leaf.kind_id() == Python::Identifier),
Self::OfSelf(name) => binds(&target, |leaf| {
leaf.kind_id() == Python::Attribute
&& leaf.child_by_field_name("object").is_some_and(|object| {
object.kind_id() == Python::Identifier && node_bytes(&object, code) == *name
})
}),
Self::Nothing => false,
}
}
}
fn body_attributes<'code>(
mut up: AncestorIter<'_, '_>,
code: &'code [u8],
) -> Option<Attributes<'code>> {
let (container, _) = up.next()?;
if container.kind_id() == Python::Module {
return Some(Attributes::Names);
}
let (owner, _) = up.next().filter(|_| python_is_block(&container))?;
owner_attributes(&owner, up, code)
}
fn owner_attributes<'code>(
owner: &Node,
up: AncestorIter<'_, '_>,
code: &'code [u8],
) -> Option<Attributes<'code>> {
if owner.kind_id() == Python::ClassDefinition {
Some(Attributes::Names)
} else if owner.kind_id() == Python::FunctionDefinition {
let receiver = is_class_init(owner, up, code)
.then(|| first_parameter_name(owner, code))
.flatten();
Some(receiver.map_or(Attributes::Nothing, Attributes::OfSelf))
} else {
None
}
}
fn is_class_init(function: &Node, up: AncestorIter<'_, '_>, code: &[u8]) -> bool {
function
.child_by_field_name("name")
.is_some_and(|name| node_bytes(&name, code) == b"__init__")
&& up
.map(|(node, _)| node)
.find(|node| node.kind_id() != Python::DecoratedDefinition && !python_is_block(node))
.is_some_and(|class| class.kind_id() == Python::ClassDefinition)
}
fn first_parameter_name<'code>(function: &Node, code: &'code [u8]) -> Option<&'code [u8]> {
let first = function
.child_by_field_name("parameters")?
.children()
.find(|child| child.is_named() && child.kind_id() != Python::Comment)?;
let name = if first.kind_id() == Python::TypedParameter {
first.children().find(Node::is_named)?
} else {
first
};
(name.kind_id() == Python::Identifier).then(|| node_bytes(&name, code))
}
fn is_str_literal(value: &Node, code: &[u8]) -> bool {
let is_str = |literal: &Node| {
!node_bytes(literal, code)
.iter()
.take_while(|byte| byte.is_ascii_alphabetic())
.any(|byte| matches!(byte.to_ascii_lowercase(), b'f' | b't' | b'b'))
};
if value.kind_id() == Python::ConcatenatedString {
value
.children()
.filter(|fragment| fragment.kind_id() == Python::String)
.all(|fragment| is_str(&fragment))
} else {
is_str(value)
}
}
fn is_bare_statement(statement: &Node) -> bool {
statement.kind_id() == Python::ExpressionStatement && statement.child_count() == 1
}
fn is_string_statement(statement: &Node, code: &[u8]) -> bool {
statement_string(statement).is_some_and(|value| is_str_literal(&value, code))
}
fn binds(target: &Node, leaf: impl Fn(&Node) -> bool) -> bool {
let mut pending = vec![*target];
while let Some(node) = pending.pop() {
match node.kind_id().into() {
Python::Identifier | Python::Subscript | Python::Attribute => {
if leaf(&node) {
return true;
}
}
_ => pending.extend(node.children().filter(Node::is_named)),
}
}
false
}
fn assignment_target<'a>(statement: &Node<'a>) -> Option<Node<'a>> {
(statement.kind_id() == Python::ExpressionStatement)
.then(|| statement.child(0))
.flatten()
.filter(|value| value.kind_id() == Python::Assignment)
.and_then(|assignment| assignment.child_by_field_name("left"))
}
pub fn python_sphinx_doc_comment<'a>(
comment: &Node<'a>,
code: &[u8],
ancestors: Ancestors<'a, '_>,
mut credit: impl FnMut(usize, usize),
) {
if !is_sphinx_doc_comment(comment, code) {
return;
}
let documents = match ancestors.siblings_before(comment).next() {
Some(statement) if statement.end_row() >= comment.start_row() => {
body_attributes(ancestors.iter(comment), code)
.is_some_and(|attributes| attributes.documented_by(&statement, code))
}
_ => ends_block_above_assignment(comment, code, ancestors),
};
if documents {
credit_block(comment, code, ancestors, &mut credit);
}
}
fn ends_block_above_assignment(comment: &Node, code: &[u8], ancestors: Ancestors<'_, '_>) -> bool {
let Some(next) = ancestors.siblings_after(comment).next() else {
return false;
};
let below = comment.end_row().saturating_add(1);
if next.start_row() != below || is_sphinx_doc_comment(&next, code) {
return false;
}
if !python_is_block(&next) {
return body_attributes(ancestors.iter(comment), code)
.is_some_and(|attributes| attributes.documented_by(&next, code));
}
let mut up = ancestors.iter(comment);
let Some((owner, _)) = up.next() else {
return false;
};
next.children()
.find(Node::is_named)
.filter(|first| first.start_row() == below)
.is_some_and(|first| {
owner_attributes(&owner, up, code)
.is_some_and(|attributes| attributes.documented_by(&first, code))
})
}
fn credit_block(
comment: &Node,
code: &[u8],
ancestors: Ancestors<'_, '_>,
credit: &mut impl FnMut(usize, usize),
) {
credit(comment.start_row(), comment.end_row());
let mut member = *comment;
let mut before = ancestors.siblings_before(comment).peekable();
while let Some(above) = before.next() {
let leads_its_row = before
.peek()
.is_none_or(|prev| prev.end_row() < above.start_row());
if !(is_sphinx_doc_comment(&above, code)
&& above.end_row().saturating_add(1) == member.start_row()
&& leads_its_row)
{
break;
}
credit(above.start_row(), above.end_row());
member = above;
}
}
fn is_sphinx_doc_comment(node: &Node, code: &[u8]) -> bool {
node.kind_id() == Python::Comment && node_bytes(node, code).starts_with(b"#:")
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "python")]
fn bare_string_parenthesis_answers(source: &str) -> Vec<[bool; 2]> {
use crate::langs::PythonCode;
use crate::node::{FoldCache, Tree};
let tree = Tree::new::<PythonCode>(source.as_bytes());
let cache = FoldCache::default();
let mut answers = Vec::new();
let mut stack = vec![(tree.get_root(), 0_usize)];
let mut chain: Vec<Node> = Vec::new();
while let Some((node, depth)) = stack.pop() {
chain.truncate(depth);
if matches!(
node.kind_id().into(),
Python::ParenthesizedExpression | Python::LPAREN | Python::RPAREN
) {
answers.push([
python_is_bare_string_parenthesis(
&node,
Ancestors::known(&chain).with_folds(&cache),
),
python_is_bare_string_parenthesis(&node, Ancestors::known(&chain)),
]);
}
chain.push(node);
let children: Vec<_> = node.children().collect();
stack.extend(children.into_iter().rev().map(|child| (child, depth + 1)));
}
answers
}
#[cfg(feature = "python")]
#[test]
fn memoised_climb_answers_as_the_climb_does() {
let nest = |depth: usize, inner: &str| {
format!("{}{inner}{}", "(".repeat(depth), ")".repeat(depth))
};
let cases = [
(nest(199, "\"s\""), 199 * 3),
(nest(200, "\"s\""), 200 * 3),
(nest(201, "\"s\""), 0),
(
format!("def f():\n {}\n", nest(3, "\"\"\"d\"\"\"")),
3 * 3,
),
("( # c\n \"s\"\n)\n".to_owned(), 3),
("(\"a\" \"b\")\n".to_owned(), 3),
(nest(3, "x"), 0),
("(\"a\", \"b\")\n".to_owned(), 0),
(nest(3, "f\"x\""), 3 * 3),
(format!("x = {}\n", nest(1_000, "\"s\"")), 0),
(format!("{}\n(\"s\"\n", nest(2, "\"s\"")), 2 * 3),
];
for (source, documenting) in cases {
let answers = bare_string_parenthesis_answers(&source);
assert!(!answers.is_empty(), "{source:.40}");
for [memoised, climbed] in &answers {
assert_eq!(memoised, climbed, "{source:.40}");
}
let found = answers.iter().filter(|[memoised, ..]| *memoised).count();
assert_eq!(found, documenting, "{source:.40}");
}
}
}