Add documentation for CompletionTreeName and CompletionDirEntry

Author: davidhalter

crates/zuban_python/src/completion.rs

@@ -12,12 +12,15 @@ use crate::{
     database::{ClassKind, Database, ParentScope, PointKind},
     debug,
     file::{ClassNodeRef, File as _, FuncNodeRef, PythonFile, is_reexport_issue},
-    goto::{PositionalDocument, unpack_union_types, with_i_s_non_self},
+    goto::{
+        FollowImportResult, PositionalDocument, try_to_follow_imports, unpack_union_types,
+        with_i_s_non_self,
+    },
     imports::{ImportResult, global_import},
     inference_state::InferenceState,
     inferred::Inferred,
     lines::BytePositionInfos,
-    name::{Range, TreeName, process_docstring},
+    name::{ModuleName, Range, TreeName, process_docstring},
     node_ref::NodeRef,
     recoverable_error,
     type_::{CallableParam, Enum, EnumMemberDefinition, FunctionKind, Namespace, Type},
@@ -732,22 +735,31 @@ impl<'db> Completion for CompletionTreeName<'db> {
     }
 
     fn documentation(&self) -> Option<Cow<'db, str>> {
-        Some(process_docstring(
-            self.file,
-            self.name.clean_docstring(),
-            || TreeName::with_unknown_parent_scope(self.db, self.file, self.name).goto_non_stub(),
-        ))
+        let doc = self.name.clean_docstring();
+        if doc.is_empty()
+            && let Some(r) = try_to_follow_imports(self.db, self.file, self.name)
+        {
+            return Some(match r {
+                FollowImportResult::File(file_index) => {
+                    let file = self.db.loaded_python_file(file_index);
+                    ModuleName { db: self.db, file }.documentation()
+                }
+                FollowImportResult::TreeName(tree_name) => tree_name.documentation(),
+            });
+        }
+        Some(process_docstring(self.file, doc, || {
+            TreeName::with_unknown_parent_scope(self.db, self.file, self.name).goto_non_stub()
+        }))
     }
 }
 
-#[expect(dead_code)]
 struct CompletionDirEntry<'db, 'x> {
     db: &'db Database,
     name: &'db str,
     entry: &'x DirectoryEntry,
 }
 
-impl Completion for CompletionDirEntry<'_, '_> {
+impl<'db> Completion for CompletionDirEntry<'db, '_> {
     fn label(&self) -> &str {
         self.name
     }
@@ -771,6 +783,22 @@ impl Completion for CompletionDirEntry<'_, '_> {
         */
         None
     }
+
+    fn documentation(&self) -> Option<Cow<'db, str>> {
+        match self.entry {
+            DirectoryEntry::File(entry) => {
+                let file_index = self.db.load_file_from_workspace(entry, false)?;
+                Some(
+                    ModuleName {
+                        db: self.db,
+                        file: self.db.loaded_python_file(file_index),
+                    }
+                    .documentation(),
+                )
+            }
+            _ => None,
+        }
+    }
 }
 
 struct KeywordCompletion {

crates/zuban_python/src/goto.rs

@@ -386,63 +386,27 @@ impl<'db, C: for<'a> FnMut(Name<'db, 'a>) -> T, T> GotoResolver<'db, C> {
         self.calculate_return(Name::ModuleName(ModuleName { db, file }))
     }
 
-    fn try_to_follow(&mut self, n: NodeRef, follow_imports: bool) -> Option<Option<T>> {
-        let p = n.point();
-        if !p.calculated() {
-            return None;
-        }
-        match p.kind() {
-            PointKind::Redirect => Some(self.check_node_ref_and_maybe_follow_import(
-                p.as_redirected_node_ref(self.infos.db),
-                follow_imports,
-            )),
-            PointKind::FileReference => Some(Some(self.goto_on_file(p.file_index()))),
-            _ => None,
-        }
+    fn try_to_follow(&mut self, n: NodeRef<'db>, follow_imports: bool) -> Option<T> {
+        self.process_follow_import_result(try_to_follow(self.infos.db, n, follow_imports)?)
+    }
+
+    fn process_follow_import_result(&mut self, r: Option<FollowImportResult<'db>>) -> Option<T> {
+        Some(match r? {
+            FollowImportResult::File(file_index) => self.goto_on_file(file_index),
+            FollowImportResult::TreeName(n) => self.calculate_return(Name::TreeName(n)),
+        })
     }
 
     fn check_node_ref_and_maybe_follow_import(
         &mut self,
         node_ref: NodeRef<'db>,
         follow_imports: bool,
     ) -> Option<T> {
-        let n = node_ref.maybe_name()?;
-        let db = self.infos.db;
-        if follow_imports && let Some(name_def) = n.name_def() {
-            match name_def.maybe_import() {
-                Some(NameImportParent::ImportFromAsName(_)) => {
-                    let ref_ = NodeRef::new(node_ref.file, name_def.index());
-                    if let Some(result) = self.try_to_follow(ref_, follow_imports) {
-                        return result;
-                    }
-                }
-                Some(NameImportParent::DottedAsName(_)) => {
-                    let p = NodeRef::new(node_ref.file, name_def.index()).point();
-                    if p.kind() == PointKind::FileReference {
-                        return Some(self.goto_on_file(p.file_index()));
-                    }
-                }
-                None => {
-                    if matches!(
-                        name_def.parent(),
-                        NameDefParent::GlobalStmt | NameDefParent::NonlocalStmt
-                    ) {
-                        let ref_ = NodeRef::new(self.infos.file, name_def.index())
-                            .global_or_nonlocal_ref();
-                        if let Some(result) = self.try_to_follow(ref_, follow_imports) {
-                            return result;
-                        }
-                    }
-                }
-            }
-        }
-        Some(
-            self.calculate_return(Name::TreeName(TreeName::with_unknown_parent_scope(
-                db,
-                node_ref.file,
-                n,
-            ))),
-        )
+        self.process_follow_import_result(check_node_ref_and_maybe_follow_import(
+            self.infos.db,
+            node_ref,
+            follow_imports,
+        ))
     }
 
     fn goto_name(&mut self, follow_imports: bool, check_inferred_attrs: bool) -> Option<Vec<T>> {
@@ -526,10 +490,10 @@ impl<'db, C: for<'a> FnMut(Name<'db, 'a>) -> T, T> GotoResolver<'db, C> {
             GotoNode::ImportFromAsName { import_as_name, .. } => Some(vec![self.try_to_follow(
                 NodeRef::new(file, import_as_name.name_def().index()),
                 follow_imports,
-            )??]),
+            )?]),
             GotoNode::GlobalName(name_def) | GotoNode::NonlocalName(name_def) => {
                 let ref_ = NodeRef::new(file, name_def.index()).global_or_nonlocal_ref();
-                if let Some(result) = self.try_to_follow(ref_, follow_imports).flatten() {
+                if let Some(result) = self.try_to_follow(ref_, follow_imports) {
                     Some(vec![result])
                 } else {
                     // This essentially just returns the name of the global definition, because we
@@ -591,6 +555,79 @@ impl<'db, C: for<'a> FnMut(Name<'db, 'a>) -> T, T> GotoResolver<'db, C> {
     }
 }
 
+pub(crate) fn try_to_follow_imports<'db>(
+    db: &'db Database,
+    file: &'db PythonFile,
+    n: CSTName,
+) -> Option<FollowImportResult<'db>> {
+    let name_def = n.name_def()?;
+    match name_def.maybe_import() {
+        Some(NameImportParent::ImportFromAsName(_)) => {
+            let ref_ = NodeRef::new(file, name_def.index());
+            if let Some(result) = try_to_follow(db, ref_, true) {
+                return result;
+            }
+        }
+        Some(NameImportParent::DottedAsName(_)) => {
+            let p = NodeRef::new(file, name_def.index()).point();
+            if p.kind() == PointKind::FileReference {
+                return Some(FollowImportResult::File(p.file_index()));
+            }
+        }
+        None => {
+            if matches!(
+                name_def.parent(),
+                NameDefParent::GlobalStmt | NameDefParent::NonlocalStmt
+            ) {
+                let ref_ = NodeRef::new(file, name_def.index()).global_or_nonlocal_ref();
+                if let Some(result) = try_to_follow(db, ref_, true) {
+                    return result;
+                }
+            }
+        }
+    }
+    None
+}
+
+fn try_to_follow<'db>(
+    db: &'db Database,
+    n: NodeRef<'db>,
+    follow_imports: bool,
+) -> Option<Option<FollowImportResult<'db>>> {
+    let p = n.point();
+    if !p.calculated() {
+        return None;
+    }
+    match p.kind() {
+        PointKind::Redirect => Some(check_node_ref_and_maybe_follow_import(
+            db,
+            p.as_redirected_node_ref(db),
+            follow_imports,
+        )),
+        PointKind::FileReference => Some(Some(FollowImportResult::File(p.file_index()))),
+        _ => None,
+    }
+}
+
+fn check_node_ref_and_maybe_follow_import<'db>(
+    db: &'db Database,
+    node_ref: NodeRef<'db>,
+    follow_imports: bool,
+) -> Option<FollowImportResult<'db>> {
+    let n = node_ref.maybe_name()?;
+    if follow_imports && let result @ Some(_) = try_to_follow_imports(db, node_ref.file, n) {
+        return result;
+    }
+    Some(FollowImportResult::TreeName(
+        TreeName::with_unknown_parent_scope(db, node_ref.file, n),
+    ))
+}
+
+pub(crate) enum FollowImportResult<'db> {
+    File(FileIndex),
+    TreeName(TreeName<'db>),
+}
+
 pub(crate) struct ReferencesResolver<'db, C, T> {
     infos: PositionalDocument<'db, GotoNode<'db>>,
     definitions: FastHashSet<(FileIndex, usize)>,

crates/zuban_python/src/name.rs

@@ -498,7 +498,7 @@ impl<'db> TreeName<'db> {
         )))
     }
 
-    fn documentation(&self) -> Cow<'db, str> {
+    pub(crate) fn documentation(&self) -> Cow<'db, str> {
         let doc = self.cst_name.clean_docstring();
         /*
         // If we don't have a result try to lookup the next assignment like
@@ -561,7 +561,7 @@ pub struct ModuleName<'db> {
 }
 
 impl<'db> ModuleName<'db> {
-    fn documentation(&self) -> Cow<'db, str> {
+    pub(crate) fn documentation(&self) -> Cow<'db, str> {
         let result = self.file.tree.root().clean_docstring();
         process_docstring(self.file, result, || self.goto_non_stub())
     }

crates/zuban_python/tests/mypylike/skipped

@@ -7,7 +7,6 @@ testLiteralAndGenericsRespectsUpperBound
 self_instance_variable1
 testNewRedefineForStatementIndexNarrowing
 goto_doc_contains_stdlib_compiled
-completion_wrong_ranges_after_non_empty_line
 
 testNarrowingTypeVarMultiple
 

crates/zuban_python/tests/mypylike/tests/completion.test

@@ -296,3 +296,104 @@ Asdf().fooba
 __main__.py:14:complete -> [Asdf:Some("hello")]
 __main__.py:16:complete -> [foobar:Some("world\n\n```python\n>>> asdf\nasdf\n```\n")]
 __main__.py:18:complete -> [foobar:Some("world\n\n```python\n>>> asdf\nasdf\n```\n")]
+
+[case completion_docstr_from_import]
+# flags: --no-typecheck
+from m import Asdf
+#? complete --show-documentation
+Asd
+#? complete --show-documentation
+Asdf.fooba
+#? complete --show-documentation
+Asdf().fooba
+
+[file m.py]
+# flags: --no-typecheck
+class Asdf:
+    """
+    hello
+    """
+    def foobar(self):
+        """
+        world
+        """
+[file m.pyi]
+class Asdf:
+    def foobar(self) -> None: ...
+
+[out]
+__main__.py:4:complete -> [Asdf:Some("hello")]
+__main__.py:6:complete -> [foobar:Some("world")]
+__main__.py:8:complete -> [foobar:Some("world")]
+
+[case completion_docstr_from_stubs]
+# flags: --no-typecheck
+from m import Asdf
+#? complete --show-documentation
+Asd
+#? complete --show-documentation
+Asdf.fooba
+#? complete --show-documentation
+Asdf().fooba
+
+[file m.py]
+# flags: --no-typecheck
+class Asdf:
+    """
+    hello
+    """
+    def foobar(self):
+        """
+        world
+        """
+[file m.pyi]
+class Asdf:
+    def foobar(self) -> None: ...
+
+[out]
+__main__.py:4:complete -> [Asdf:Some("hello")]
+__main__.py:6:complete -> [foobar:Some("world")]
+__main__.py:8:complete -> [foobar:Some("world")]
+
+[case completion_modules]
+#? complete --show-documentation
+import mod_m
+#? complete --show-documentation
+import mod_n
+#? complete --show-documentation
+mod_m
+#? complete --show-documentation
+mod_n
+
+#? complete --show-documentation
+mod_m.mod_o
+#? complete --show-documentation
+mod_m.mod_p
+
+#? complete --show-documentation
+from mod_m import mod_o
+#? complete --show-documentation
+from mod_m import mod_p
+
+[file mod_m.py]
+"m docstr"
+import mod_o
+import mod_p
+[file mod_n.py]
+"n docstr"
+[file mod_n.pyi]
+[file mod_o.py]
+"o docstr"
+[file mod_p.py]
+"p docstr"
+[file mod_p.pyi]
+
+[out]
+__main__.py:2:complete -> [mod_m:Some("m docstr")]
+__main__.py:4:complete -> [mod_n:Some("n docstr")]
+__main__.py:6:complete -> [mod_m:Some("m docstr")]
+__main__.py:8:complete -> [mod_n:Some("n docstr")]
+__main__.py:11:complete -> [mod_o:Some("o docstr")]
+__main__.py:13:complete -> [mod_p:Some("p docstr")]
+__main__.py:16:complete -> [mod_o:Some("o docstr")]
+__main__.py:18:complete -> [mod_p:Some("p docstr")]