feat: support python package diagnostics

GreasySlug · Jul 7, 2024 · fcfa50d · fcfa50d
1 parent 1e23c6b
commit fcfa50d
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 23 deletions.
diff --git a/crates/els/diagnostics.rs b/crates/els/diagnostics.rs
@@ -8,7 +8,7 @@ use std::time::Duration;
 
 use erg_common::consts::PYTHON_MODE;
 use erg_common::dict::Dict;
-use erg_common::pathutil::{project_entry_file_of, project_root_dir_of};
+use erg_common::pathutil::{project_entry_dir_of, project_entry_file_of};
 use erg_common::spawn::{safe_yield, spawn_new_thread};
 use erg_common::style::*;
 use erg_common::{fn_name, lsp_log};
@@ -374,7 +374,9 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
             let Ok(entry) = entry else {
                 continue;
             };
-            if entry.path().extension() == Some(OsStr::new("er")) {
+            if entry.path().extension() == Some(OsStr::new("er"))
+                || (PYTHON_MODE && entry.path().extension() == Some(OsStr::new("py")))
+            {
                 if let Ok(uri) = NormalizedUrl::from_file_path(entry.path()) {
                     uris.push(uri);
                 }
@@ -443,14 +445,9 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
                 let Ok(current_dir) = current_dir() else {
                     work_done!(token)
                 };
-                let Some(project_root) = project_root_dir_of(&current_dir) else {
+                let Some(src_dir) = project_entry_dir_of(&current_dir) else {
                     work_done!(token);
                 };
-                let src_dir = if project_root.join("src").is_dir() {
-                    project_root.join("src")
-                } else {
-                    project_root
-                };
                 let Some(main_path) = project_entry_file_of(&current_dir) else {
                     work_done!(token);
                 };

diff --git a/crates/els/server.rs b/crates/els/server.rs
@@ -12,7 +12,7 @@ use erg_common::config::ErgConfig;
 use erg_common::consts::PYTHON_MODE;
 use erg_common::dict::Dict;
 use erg_common::env::erg_path;
-use erg_common::pathutil::{project_root_dir_of, NormalizedPathBuf};
+use erg_common::pathutil::{project_entry_dir_of, NormalizedPathBuf};
 use erg_common::shared::{MappedRwLockReadGuard, Shared};
 use erg_common::spawn::{safe_yield, spawn_new_thread};
 use erg_common::traits::Stream;
@@ -1043,7 +1043,7 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
     }
 
     pub(crate) fn get_workspace_ctxs(&self) -> Vec<&Context> {
-        let project_root = project_root_dir_of(&self.home).unwrap_or(self.home.clone());
+        let project_root = project_entry_dir_of(&self.home).unwrap_or(self.home.clone());
         let mut ctxs = vec![];
         for (path, ent) in self.shared.raw_path_and_modules() {
             if path.starts_with(&project_root) {

diff --git a/crates/erg_common/pathutil.rs b/crates/erg_common/pathutil.rs
@@ -4,6 +4,7 @@ use std::fmt;
 use std::ops::Deref;
 use std::path::{Component, Path, PathBuf};
 
+use crate::consts::PYTHON_MODE;
 use crate::env::erg_pkgs_path;
 use crate::{normalize_path, Str};
 
@@ -337,29 +338,86 @@ pub fn mod_name(path: &Path) -> Str {
     Str::from(name)
 }
 
-pub fn project_root_dir_of(path: &Path) -> Option<PathBuf> {
+fn erg_project_entry_dir_of(path: &Path) -> Option<PathBuf> {
     if path.is_dir() && path.join("package.er").exists() {
-        return Some(path.to_path_buf());
+        if path.join("src").exists() {
+            return Some(path.join("src"));
+        } else {
+            return Some(path.to_path_buf());
+        }
     }
     let mut path = path.to_path_buf();
     while let Some(parent) = path.parent() {
         if parent.join("package.er").exists() {
-            return Some(parent.to_path_buf());
+            if parent.join("src").exists() {
+                return Some(parent.join("src"));
+            } else {
+                return Some(parent.to_path_buf());
+            }
         }
         path = parent.to_path_buf();
     }
     None
 }
 
-pub fn project_entry_file_of(path: &Path) -> Option<PathBuf> {
-    let project_root = project_root_dir_of(path)?;
-    if project_root.join("src/lib.er").exists() {
-        Some(project_root.join("src/lib.er"))
-    } else if project_root.join("src/main.er").exists() {
-        Some(project_root.join("src/main.er"))
-    } else if project_root.join("src/lib.d.er").exists() {
-        Some(project_root.join("src/lib.d.er"))
+fn py_project_entry_dir_of(path: &Path) -> Option<PathBuf> {
+    let dir_name = path.file_name()?;
+    if path.is_dir() && path.join("pyproject.toml").exists() {
+        if path.join(dir_name).exists() {
+            return Some(path.join(dir_name));
+        } else if path.join("src").join(dir_name).exists() {
+            return Some(path.join("src").join(dir_name));
+        }
+    }
+    let mut path = path.to_path_buf();
+    while let Some(parent) = path.parent() {
+        let dir_name = parent.file_name()?;
+        if parent.join("pyproject.toml").exists() {
+            if parent.join(dir_name).exists() {
+                return Some(parent.join(dir_name));
+            } else if parent.join("src").join(dir_name).exists() {
+                return Some(parent.join("src").join(dir_name));
+            }
+        }
+        path = parent.to_path_buf();
+    }
+    None
+}
+
+pub fn project_entry_dir_of(path: &Path) -> Option<PathBuf> {
+    if PYTHON_MODE {
+        py_project_entry_dir_of(path)
+    } else {
+        erg_project_entry_dir_of(path)
+    }
+}
+
+fn erg_project_entry_file_of(path: &Path) -> Option<PathBuf> {
+    let entry = erg_project_entry_dir_of(path)?;
+    if entry.join("lib.er").exists() {
+        Some(entry.join("lib.er"))
+    } else if entry.join("main.er").exists() {
+        Some(entry.join("main.er"))
+    } else if entry.join("lib.d.er").exists() {
+        Some(entry.join("lib.d.er"))
     } else {
         None
     }
 }
+
+fn py_project_entry_file_of(path: &Path) -> Option<PathBuf> {
+    let entry = py_project_entry_dir_of(path)?;
+    if entry.join("__init__.py").exists() {
+        Some(entry.join("__init__.py"))
+    } else {
+        None
+    }
+}
+
+pub fn project_entry_file_of(path: &Path) -> Option<PathBuf> {
+    if PYTHON_MODE {
+        py_project_entry_file_of(path)
+    } else {
+        erg_project_entry_file_of(path)
+    }
+}
diff --git a/crates/erg_compiler/build_package.rs b/crates/erg_compiler/build_package.rs
@@ -21,7 +21,7 @@ use erg_common::error::MultiErrorDisplay;
 use erg_common::io::Input;
 #[allow(unused)]
 use erg_common::log;
-use erg_common::pathutil::{mod_name, project_root_dir_of, NormalizedPathBuf};
+use erg_common::pathutil::{mod_name, project_entry_dir_of, NormalizedPathBuf};
 use erg_common::spawn::spawn_new_thread;
 use erg_common::str::Str;
 use erg_common::traits::{ExitStatus, New, Runnable, Stream};
@@ -712,7 +712,7 @@ impl<ASTBuilder: ASTBuildable, HIRBuilder: Buildable>
         };
         let root_import_path = root_path.and_then(|path| cfg.input.resolve_path(path, cfg));
         if let Some(root_import_path) = root_import_path.map(NormalizedPathBuf::from) {
-            if project_root_dir_of(&root_import_path) != project_root_dir_of(&from_path) {
+            if project_entry_dir_of(&root_import_path) != project_entry_dir_of(&from_path) {
                 let mut root_import_cfg = cfg.inherit(root_import_path.to_path_buf());
                 self.shared.graph.add_node_if_none(&root_import_path);
                 let _ = self