Find source docstrings for stub members declared under version checks

pyrefly/lib/state/lsp_attributes.rs

@@ -5,12 +5,15 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::iter;
+
 use pyrefly_python::ast::Ast;
 use pyrefly_python::docstring::Docstring;
 use pyrefly_python::module::Module;
 use ruff_python_ast::Expr;
 use ruff_python_ast::ModModule;
 use ruff_python_ast::Stmt;
+use ruff_python_ast::StmtIf;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 
@@ -55,6 +58,14 @@ impl AttributeContext {
                 Stmt::FunctionDef(function) if function.name.range == target_range => {
                     AttributeKind::Function
                 }
+                Stmt::If(if_stmt) => {
+                    if let Some(context) = if_branches(if_stmt).find_map(|branch| {
+                        Self::from_body(branch, &mut parents.clone(), target_range)
+                    }) {
+                        return Some(context);
+                    }
+                    continue;
+                }
                 Stmt::Assign(assign)
                     if assign.targets.iter().any(
                         |target| matches!(target, Expr::Name(name) if name.range() == target_range),
@@ -92,26 +103,49 @@ pub(crate) fn definition_from_ast(
     context: &AttributeContext,
     attr_name: &ruff_python_ast::name::Name,
 ) -> Option<(TextRange, Option<TextRange>)> {
-    let mut body = ast.body.as_slice();
+    let mut body = flatten_conditionals(&ast.body);
     for class_name in &context.parent_classes {
         let class_def = body.iter().find_map(|stmt| match stmt {
             Stmt::ClassDef(class_def) if &class_def.name.id == class_name => Some(class_def),
             _ => None,
         })?;
-        body = class_def.body.as_slice();
+        body = flatten_conditionals(&class_def.body);
     }
     match context.kind {
-        AttributeKind::Function => definition_from_function(body, attr_name),
-        AttributeKind::Class => definition_from_class(body, attr_name),
-        AttributeKind::Assignment => definition_from_assignment(body, attr_name),
+        AttributeKind::Function => definition_from_function(&body, attr_name),
+        AttributeKind::Class => definition_from_class(&body, attr_name),
+        AttributeKind::Assignment => definition_from_assignment(&body, attr_name),
     }
 }
 
+/// The bodies of the `if`, `elif`, and `else` branches of `if_stmt`.
+fn if_branches(if_stmt: &StmtIf) -> impl Iterator<Item = &[Stmt]> {
+    iter::once(if_stmt.body.as_slice()).chain(
+        if_stmt
+            .elif_else_clauses
+            .iter()
+            .map(|clause| clause.body.as_slice()),
+    )
+}
+
+/// Lists the statements of `body`, including those nested in `if` blocks such as
+/// `if sys.version_info >= ...:`, which do not introduce a scope.
+fn flatten_conditionals(body: &[Stmt]) -> Vec<&Stmt> {
+    body.iter()
+        .flat_map(|stmt| match stmt {
+            Stmt::If(if_stmt) => if_branches(if_stmt)
+                .flat_map(flatten_conditionals)
+                .collect(),
+            _ => vec![stmt],
+        })
+        .collect()
+}
+
 fn definition_from_function(
-    body: &[Stmt],
+    body: &[&Stmt],
     attr_name: &ruff_python_ast::name::Name,
 ) -> Option<(TextRange, Option<TextRange>)> {
-    for stmt in body {
+    for &stmt in body {
         if let Stmt::FunctionDef(func_def) = stmt
             && &func_def.name.id == attr_name
         {
@@ -123,10 +157,10 @@ fn definition_from_function(
 }
 
 fn definition_from_class(
-    body: &[Stmt],
+    body: &[&Stmt],
     attr_name: &ruff_python_ast::name::Name,
 ) -> Option<(TextRange, Option<TextRange>)> {
-    for stmt in body {
+    for &stmt in body {
         if let Stmt::ClassDef(class_def) = stmt
             && &class_def.name.id == attr_name
         {
@@ -138,10 +172,10 @@ fn definition_from_class(
 }
 
 fn definition_from_assignment(
-    body: &[Stmt],
+    body: &[&Stmt],
     attr_name: &ruff_python_ast::name::Name,
 ) -> Option<(TextRange, Option<TextRange>)> {
-    for (idx, stmt) in body.iter().enumerate() {
+    for (idx, &stmt) in body.iter().enumerate() {
         let (matches, def_range) = match stmt {
             Stmt::Assign(assign) => {
                 let target_opt = assign
@@ -157,7 +191,7 @@ fn definition_from_assignment(
             _ => (false, None),
         };
         if matches {
-            let docstring_range = if let Some(Stmt::Expr(expr_stmt)) = body.get(idx + 1)
+            let docstring_range = if let Some(Stmt::Expr(expr_stmt)) = body.get(idx + 1).copied()
                 && matches!(expr_stmt.value.as_ref(), Expr::StringLiteral(_))
             {
                 Some(expr_stmt.range())
@@ -169,3 +203,33 @@ fn definition_from_assignment(
     }
     None
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use pyrefly_python::module_name::ModuleName;
+    use pyrefly_python::module_path::ModulePath;
+    use ruff_python_ast::name::Name;
+
+    use super::*;
+
+    /// A method declared under `if sys.version_info` in a stub maps to the unconditional
+    /// method in the source, and the source docstring is found.
+    #[test]
+    fn test_method_in_conditional_block() {
+        let stub = "class C:\n    if sys.version_info >= (3, 12):\n        def f(self): ...\n    else:\n        def f(self, x): ...\n";
+        let source = "class C:\n    def f(self):\n        \"\"\"Doc.\"\"\"\n";
+        let stub_module = Module::new(
+            ModuleName::from_str("m"),
+            ModulePath::filesystem("m.pyi".into()),
+            Arc::new(stub.to_owned()),
+        );
+        let offset = stub.find("f(self)").unwrap();
+        let range = TextRange::at(u32::try_from(offset).unwrap().into(), 1.into());
+        let context = AttributeContext::from_module(&stub_module, range).unwrap();
+        let ast = Ast::parse(source, stub_module.source_type()).0;
+        let (_, docstring) = definition_from_ast(&ast, &context, &Name::new("f")).unwrap();
+        assert!(docstring.is_some());
+    }
+}