Skip to content

Commit 3d097f1

Browse files
committed
Find source docstrings for stub members declared under version checks
Hover on a stdlib method such as Path.read_text showed no docstring. Typeshed stubs declare many members inside `if sys.version_info` blocks, and the lookup that maps a stub declaration to its source declaration only searched the top level of a class body, so it gave up. Treat statements nested in `if` blocks as part of the enclosing scope when locating a declaration in the stub and in the source. Fixes #5568
1 parent c7a8e3a commit 3d097f1

1 file changed

Lines changed: 76 additions & 12 deletions

File tree

‎pyrefly/lib/state/lsp_attributes.rs‎

Lines changed: 76 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@
55
* LICENSE file in the root directory of this source tree.
66
*/
77

8+
use std::iter;
9+
810
use pyrefly_python::ast::Ast;
911
use pyrefly_python::docstring::Docstring;
1012
use pyrefly_python::module::Module;
1113
use ruff_python_ast::Expr;
1214
use ruff_python_ast::ModModule;
1315
use ruff_python_ast::Stmt;
16+
use ruff_python_ast::StmtIf;
1417
use ruff_text_size::Ranged;
1518
use ruff_text_size::TextRange;
1619

@@ -55,6 +58,14 @@ impl AttributeContext {
5558
Stmt::FunctionDef(function) if function.name.range == target_range => {
5659
AttributeKind::Function
5760
}
61+
Stmt::If(if_stmt) => {
62+
if let Some(context) = if_branches(if_stmt).find_map(|branch| {
63+
Self::from_body(branch, &mut parents.clone(), target_range)
64+
}) {
65+
return Some(context);
66+
}
67+
continue;
68+
}
5869
Stmt::Assign(assign)
5970
if assign.targets.iter().any(
6071
|target| matches!(target, Expr::Name(name) if name.range() == target_range),
@@ -92,26 +103,49 @@ pub(crate) fn definition_from_ast(
92103
context: &AttributeContext,
93104
attr_name: &ruff_python_ast::name::Name,
94105
) -> Option<(TextRange, Option<TextRange>)> {
95-
let mut body = ast.body.as_slice();
106+
let mut body = flatten_conditionals(&ast.body);
96107
for class_name in &context.parent_classes {
97108
let class_def = body.iter().find_map(|stmt| match stmt {
98109
Stmt::ClassDef(class_def) if &class_def.name.id == class_name => Some(class_def),
99110
_ => None,
100111
})?;
101-
body = class_def.body.as_slice();
112+
body = flatten_conditionals(&class_def.body);
102113
}
103114
match context.kind {
104-
AttributeKind::Function => definition_from_function(body, attr_name),
105-
AttributeKind::Class => definition_from_class(body, attr_name),
106-
AttributeKind::Assignment => definition_from_assignment(body, attr_name),
115+
AttributeKind::Function => definition_from_function(&body, attr_name),
116+
AttributeKind::Class => definition_from_class(&body, attr_name),
117+
AttributeKind::Assignment => definition_from_assignment(&body, attr_name),
107118
}
108119
}
109120

121+
/// The bodies of the `if`, `elif`, and `else` branches of `if_stmt`.
122+
fn if_branches(if_stmt: &StmtIf) -> impl Iterator<Item = &[Stmt]> {
123+
iter::once(if_stmt.body.as_slice()).chain(
124+
if_stmt
125+
.elif_else_clauses
126+
.iter()
127+
.map(|clause| clause.body.as_slice()),
128+
)
129+
}
130+
131+
/// Lists the statements of `body`, including those nested in `if` blocks such as
132+
/// `if sys.version_info >= ...:`, which do not introduce a scope.
133+
fn flatten_conditionals(body: &[Stmt]) -> Vec<&Stmt> {
134+
body.iter()
135+
.flat_map(|stmt| match stmt {
136+
Stmt::If(if_stmt) => if_branches(if_stmt)
137+
.flat_map(flatten_conditionals)
138+
.collect(),
139+
_ => vec![stmt],
140+
})
141+
.collect()
142+
}
143+
110144
fn definition_from_function(
111-
body: &[Stmt],
145+
body: &[&Stmt],
112146
attr_name: &ruff_python_ast::name::Name,
113147
) -> Option<(TextRange, Option<TextRange>)> {
114-
for stmt in body {
148+
for &stmt in body {
115149
if let Stmt::FunctionDef(func_def) = stmt
116150
&& &func_def.name.id == attr_name
117151
{
@@ -123,10 +157,10 @@ fn definition_from_function(
123157
}
124158

125159
fn definition_from_class(
126-
body: &[Stmt],
160+
body: &[&Stmt],
127161
attr_name: &ruff_python_ast::name::Name,
128162
) -> Option<(TextRange, Option<TextRange>)> {
129-
for stmt in body {
163+
for &stmt in body {
130164
if let Stmt::ClassDef(class_def) = stmt
131165
&& &class_def.name.id == attr_name
132166
{
@@ -138,10 +172,10 @@ fn definition_from_class(
138172
}
139173

140174
fn definition_from_assignment(
141-
body: &[Stmt],
175+
body: &[&Stmt],
142176
attr_name: &ruff_python_ast::name::Name,
143177
) -> Option<(TextRange, Option<TextRange>)> {
144-
for (idx, stmt) in body.iter().enumerate() {
178+
for (idx, &stmt) in body.iter().enumerate() {
145179
let (matches, def_range) = match stmt {
146180
Stmt::Assign(assign) => {
147181
let target_opt = assign
@@ -157,7 +191,7 @@ fn definition_from_assignment(
157191
_ => (false, None),
158192
};
159193
if matches {
160-
let docstring_range = if let Some(Stmt::Expr(expr_stmt)) = body.get(idx + 1)
194+
let docstring_range = if let Some(Stmt::Expr(expr_stmt)) = body.get(idx + 1).copied()
161195
&& matches!(expr_stmt.value.as_ref(), Expr::StringLiteral(_))
162196
{
163197
Some(expr_stmt.range())
@@ -169,3 +203,33 @@ fn definition_from_assignment(
169203
}
170204
None
171205
}
206+
207+
#[cfg(test)]
208+
mod tests {
209+
use std::sync::Arc;
210+
211+
use pyrefly_python::module_name::ModuleName;
212+
use pyrefly_python::module_path::ModulePath;
213+
use ruff_python_ast::name::Name;
214+
215+
use super::*;
216+
217+
/// A method declared under `if sys.version_info` in a stub maps to the unconditional
218+
/// method in the source, and the source docstring is found.
219+
#[test]
220+
fn test_method_in_conditional_block() {
221+
let stub = "class C:\n if sys.version_info >= (3, 12):\n def f(self): ...\n else:\n def f(self, x): ...\n";
222+
let source = "class C:\n def f(self):\n \"\"\"Doc.\"\"\"\n";
223+
let stub_module = Module::new(
224+
ModuleName::from_str("m"),
225+
ModulePath::filesystem("m.pyi".into()),
226+
Arc::new(stub.to_owned()),
227+
);
228+
let offset = stub.find("f(self)").unwrap();
229+
let range = TextRange::at(u32::try_from(offset).unwrap().into(), 1.into());
230+
let context = AttributeContext::from_module(&stub_module, range).unwrap();
231+
let ast = Ast::parse(source, stub_module.source_type()).0;
232+
let (_, docstring) = definition_from_ast(&ast, &context, &Name::new("f")).unwrap();
233+
assert!(docstring.is_some());
234+
}
235+
}

0 commit comments

Comments
 (0)