feat: search site-packages with pyimport

crates/erg_common/config.rs

@@ -14,7 +14,7 @@ use crate::help_messages::{command_message, mode_message, OPTIONS};
 use crate::levenshtein::get_similar_name;
 use crate::normalize_path;
 use crate::pathutil::add_postfix_foreach;
-use crate::python_util::{detect_magic_number, get_python_version, PythonVersion};
+use crate::python_util::{detect_magic_number, get_python_version, get_sys_path, PythonVersion};
 use crate::random::random;
 use crate::serialize::{get_magic_num_from_bytes, get_ver_from_magic_num};
 use crate::stdin::GLOBAL_STDIN;
@@ -442,7 +442,7 @@ impl Input {
     /// resolution order:
     /// 1. `{path}.er`
     /// 2. `{path}/__init__.er`
-    pub fn local_resolve(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
+    pub fn resolve_local(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
         let mut dir = self.dir();
         dir.push(path);
         dir.set_extension("er"); // {path}.er
@@ -455,10 +455,10 @@ impl Input {
         Ok(normalize_path(path))
     }
 
-    pub fn local_decl_resolve(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
-        self._local_decl_resolve(path).or_else(|_| {
+    pub fn resolve_local_decl(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
+        self._resolve_local_decl(path).or_else(|_| {
             let path = add_postfix_foreach(path, ".d");
-            self._local_decl_resolve(&path)
+            self._resolve_local_decl(&path)
         })
     }
 
@@ -467,7 +467,7 @@ impl Input {
     /// 2. `{path}/__init__.d.er`
     /// 3. `__pycache__/{path}.d.er`
     /// 4. `{path}/__pycache__/__init__.d.er`
-    fn _local_decl_resolve(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
+    fn _resolve_local_decl(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
         let mut dir = self.dir();
         let mut comps = path.components();
         let last = comps
@@ -504,7 +504,7 @@ impl Input {
         Ok(normalize_path(path))
     }
 
-    pub fn local_py_resolve(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
+    fn resolve_local_py(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
         let mut dir = self.dir();
         dir.push(path);
         dir.set_extension("py");
@@ -516,6 +516,33 @@ impl Input {
         })?;
         Ok(normalize_path(path))
     }
+
+    pub fn resolve_py(&self, path: &Path) -> Result<PathBuf, std::io::Error> {
+        self.resolve_local_py(path).or_else(|_| {
+            // For now, only see `site-packages`
+            let site_packages = get_sys_path()
+                .into_iter()
+                .filter(|p| p.as_os_str().to_string_lossy().contains("site-packages"));
+            for sys_path in site_packages {
+                let mut dir = sys_path;
+                dir.push(path);
+                dir.set_extension("py");
+                if dir.exists() {
+                    return Ok(normalize_path(dir));
+                }
+                dir.pop();
+                dir.push(path);
+                dir.push("__init__.py");
+                if dir.exists() {
+                    return Ok(normalize_path(dir));
+                }
+            }
+            Err(std::io::Error::new(
+                std::io::ErrorKind::NotFound,
+                format!("cannot find module `{}`", path.display()),
+            ))
+        })
+    }
 }
 
 #[derive(Debug, Clone)]

crates/erg_common/env.rs

@@ -2,6 +2,7 @@ use std::env::var;
 use std::path::PathBuf;
 
 use crate::normalize_path;
+use crate::python_util::get_sys_path;
 use crate::style::colors::*;
 use crate::style::RESET;
 
@@ -52,6 +53,17 @@ fn _erg_external_lib_path() -> PathBuf {
             PathBuf::from("lib/external/")
         })
 }
+fn _python_site_packages() -> impl Iterator<Item = PathBuf> {
+    get_sys_path()
+        .into_iter()
+        .filter(|p| p.ends_with("site-packages"))
+        .map(|p| {
+            p.canonicalize().unwrap_or_else(|_| {
+                eprintln!("{RED}[ERR] ERG_PATH/lib/external not found {RESET}");
+                PathBuf::from("lib/external/")
+            })
+        })
+}
 
 thread_local! {
     pub static ERG_PATH: PathBuf = normalize_path(_erg_path());
@@ -59,6 +71,7 @@ thread_local! {
     pub static ERG_STD_DECL_PATH: PathBuf = normalize_path(_erg_std_decl_path());
     pub static ERG_PYSTD_PATH: PathBuf = normalize_path(_erg_pystd_path());
     pub static ERG_EXTERNAL_LIB_PATH: PathBuf = normalize_path(_erg_external_lib_path());
+    pub static PYTHON_SITE_PACKAGES: Vec<PathBuf> = _python_site_packages().collect();
 }
 
 pub fn erg_path() -> PathBuf {
@@ -80,3 +93,7 @@ pub fn erg_pystd_path() -> PathBuf {
 pub fn erg_py_external_lib_path() -> PathBuf {
     ERG_EXTERNAL_LIB_PATH.with(|s| s.clone())
 }
+
+pub fn python_site_packages() -> Vec<PathBuf> {
+    PYTHON_SITE_PACKAGES.with(|s| s.clone())
+}

crates/erg_common/python_util.rs

@@ -1,6 +1,7 @@
 //! utilities for calling CPython.
 //!
 //! CPythonを呼び出すためのユーティリティー
+use std::path::PathBuf;
 use std::process::Command;
 
 use crate::fn_name_full;
@@ -687,6 +688,32 @@ pub fn env_python_version() -> PythonVersion {
     get_python_version(&which_python())
 }
 
+pub fn get_sys_path() -> Vec<PathBuf> {
+    let py_command = which_python();
+    let code = "import sys; print('\\n'.join(sys.path))";
+    let out = if cfg!(windows) {
+        Command::new("cmd")
+            .arg("/C")
+            .arg(py_command)
+            .arg("-c")
+            .arg(code)
+            .output()
+            .expect("cannot get the sys.path")
+    } else {
+        let exec_command = format!("{py_command} -c \"{code}\"");
+        Command::new("sh")
+            .arg("-c")
+            .arg(exec_command)
+            .output()
+            .expect("cannot get the sys.path")
+    };
+    let s_sys_path = String::from_utf8(out.stdout).unwrap();
+    s_sys_path
+        .split('\n')
+        .map(|s| PathBuf::from(s.trim().to_string()))
+        .collect()
+}
+
 /// executes over a shell, cause `python` may not exist as an executable file (like pyenv)
 pub fn exec_pyc<S: Into<String>>(
     file: S,

crates/erg_compiler/context/inquire.rs

@@ -5,7 +5,9 @@ use std::path::{Path, PathBuf};
 use erg_common::config::{ErgConfig, Input};
 use erg_common::consts::{ERG_MODE, PYTHON_MODE};
 use erg_common::dict;
-use erg_common::env::{erg_py_external_lib_path, erg_pystd_path, erg_std_path};
+use erg_common::env::{
+    erg_py_external_lib_path, erg_pystd_path, erg_std_path, python_site_packages,
+};
 use erg_common::error::{ErrorCore, Location, SubMessage};
 use erg_common::levenshtein;
 use erg_common::pathutil::add_postfix_foreach;
@@ -2463,7 +2465,7 @@ impl Context {
     /// 3. `std/{path}.er`
     /// 4. `std/{path}/__init__.er`
     pub(crate) fn resolve_real_path(cfg: &ErgConfig, path: &Path) -> Option<PathBuf> {
-        if let Ok(path) = cfg.input.local_resolve(path) {
+        if let Ok(path) = cfg.input.resolve_local(path) {
             Some(path)
         } else if let Ok(path) = erg_std_path()
             .join(format!("{}.er", path.display()))
@@ -2482,16 +2484,18 @@ impl Context {
     }
 
     /// resolution order:
-    /// 1. `{path}.d.er`
-    /// 2. `{path}/__init__.d.er`
-    /// 3. `__pycache__/{path}.d.er`
-    /// 4. `{path}/__pycache__/__init__.d.er`
-    /// 5. `{path}.d/__init__.d.er`
-    /// 6. `{path}.d/__pycache__/__init__.d.er`
-    /// 7. `std/{path}.d.er`
-    /// 8. `std/{path}/__init__.d.er`
+    /// 1.  `{path}.d.er`
+    /// 2.  `{path}/__init__.d.er`
+    /// 3.  `__pycache__/{path}.d.er`
+    /// 4.  `{path}/__pycache__/__init__.d.er`
+    /// 5.  `{path}.d/__init__.d.er`
+    /// 6.  `{path}.d/__pycache__/__init__.d.er`
+    /// 7.  `std/{path}.d.er`
+    /// 8.  `std/{path}/__init__.d.er`
+    /// 9.  `site-packages/__pycache__/{path}.d.er`
+    /// 10. `site-packages/{path}/__pycache__/__init__.d.er`
     pub(crate) fn resolve_decl_path(cfg: &ErgConfig, path: &Path) -> Option<PathBuf> {
-        if let Ok(path) = cfg.input.local_decl_resolve(path) {
+        if let Ok(path) = cfg.input.resolve_local_decl(path) {
             Some(path)
         } else {
             let py_roots = [erg_pystd_path, erg_py_external_lib_path];
@@ -2500,11 +2504,16 @@ impl Context {
                     return Some(path);
                 }
             }
+            for site_packages in python_site_packages() {
+                if let Some(path) = Self::resolve_site_pkgs_decl_path(site_packages, path) {
+                    return Some(path);
+                }
+            }
             None
         }
     }
 
-    pub(crate) fn resolve_std_decl_path(root: PathBuf, path: &Path) -> Option<PathBuf> {
+    fn resolve_std_decl_path(root: PathBuf, path: &Path) -> Option<PathBuf> {
         let mut path = add_postfix_foreach(path, ".d");
         path.set_extension("d.er"); // set_extension overrides the previous one
         if let Ok(path) = root.join(&path).canonicalize() {
@@ -2524,6 +2533,29 @@ impl Context {
         }
     }
 
+    /// 1. `site-packages/__pycache__/{path}.d.er`
+    /// 2. `site-packages/{path}/__pycache__/__init__.d.er`
+    fn resolve_site_pkgs_decl_path(site_packages: PathBuf, path: &Path) -> Option<PathBuf> {
+        let mut path_buf = path.to_path_buf();
+        path_buf.set_extension("d.er"); // set_extension overrides the previous one
+        if let Ok(path) = site_packages
+            .join("__pycache__")
+            .join(&path_buf)
+            .canonicalize()
+        {
+            Some(normalize_path(path))
+        } else if let Ok(path) = site_packages
+            .join(path)
+            .join("__pycache__")
+            .join("__init__.d.er")
+            .canonicalize()
+        {
+            Some(normalize_path(path))
+        } else {
+            None
+        }
+    }
+
     pub(crate) fn try_push_path(&self, mut path: PathBuf, add: &Path) -> Result<PathBuf, String> {
         path.pop(); // __init__.d.er
         if let Ok(path) = path.join(add).canonicalize() {

crates/erg_compiler/context/register.rs

@@ -54,13 +54,43 @@ pub fn valid_mod_name(name: &str) -> bool {
     !name.is_empty() && !name.starts_with('/') && name.trim() == name
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum CheckStatus {
+    Succeed,
+    Failed,
+    Ongoing,
+}
+
+impl fmt::Display for CheckStatus {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            CheckStatus::Succeed => write!(f, "succeed"),
+            CheckStatus::Failed => write!(f, "failed"),
+            CheckStatus::Ongoing => write!(f, "ongoing"),
+        }
+    }
+}
+
+impl std::str::FromStr for CheckStatus {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "succeed" => Ok(CheckStatus::Succeed),
+            "failed" => Ok(CheckStatus::Failed),
+            "ongoing" => Ok(CheckStatus::Ongoing),
+            _ => Err(format!("invalid status: {s}")),
+        }
+    }
+}
+
 /// format:
 /// ```python
 /// #[pylyzer] succeed foo.py 1234567890
 /// ```
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct PylyzerStatus {
-    pub succeed: bool,
+    pub status: CheckStatus,
     pub file: PathBuf,
     pub timestamp: SystemTime,
 }
@@ -70,7 +100,7 @@ impl fmt::Display for PylyzerStatus {
         write!(
             f,
             "##[pylyzer] {} {} {}",
-            if self.succeed { "succeed" } else { "failed" },
+            self.status,
             self.file.display(),
             self.timestamp
                 .duration_since(SystemTime::UNIX_EPOCH)
@@ -89,8 +119,8 @@ impl std::str::FromStr for PylyzerStatus {
         if pylyzer != "##[pylyzer]" {
             return Err("not pylyzer".to_string());
         }
-        let succeed = iter.next().ok_or("no succeed")?;
-        let succeed = succeed == "succeed";
+        let status = iter.next().ok_or("no succeed")?;
+        let status = status.parse()?;
         let file = iter.next().ok_or("no file")?;
         let file = PathBuf::from(file);
         let timestamp = iter.next().ok_or("no timestamp")?;
@@ -102,7 +132,7 @@ impl std::str::FromStr for PylyzerStatus {
             ))
             .ok_or("timestamp overflow")?;
         Ok(PylyzerStatus {
-            succeed,
+            status,
             file,
             timestamp,
         })
@@ -1901,7 +1931,7 @@ impl Context {
     }
 
     fn try_gen_py_decl_file(&self, __name__: &Str) -> Result<PathBuf, ()> {
-        if let Ok(path) = self.cfg.input.local_py_resolve(Path::new(&__name__[..])) {
+        if let Ok(path) = self.cfg.input.resolve_py(Path::new(&__name__[..])) {
             let (out, err) = if self.cfg.mode == ErgMode::LanguageServer || self.cfg.quiet_repl {
                 (Stdio::null(), Stdio::null())
             } else {