Add the path of Matplotlib to the link search (for Anaconda)

Chris00 · Mar 2, 2024 · 1a94e06 · 1a94e06
1 parent ffbd993
commit 1a94e06
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 40 deletions.
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
@@ -10,46 +10,46 @@ jobs:
       matrix:
         build: [stable, beta, nightly, linux32, macos, aarch64-ios, win64, windows-latest]
         include:
-          - build: stable
-            os: ubuntu-latest
-            rust: stable
-            target: x86_64-unknown-linux-gnu
-          - build: beta
-            os: ubuntu-latest
-            rust: beta
-            target: x86_64-unknown-linux-gnu
-          - build: nightly
-            os: ubuntu-latest
-            rust: nightly
-            target: x86_64-unknown-linux-gnu
-          - build: linux32
-            os: ubuntu-latest
-            rust: stable
-            target: i686-unknown-linux-gnu
+          # - build: stable
+          #   os: ubuntu-latest
+          #   rust: stable
+          #   target: x86_64-unknown-linux-gnu
+          # - build: beta
+          #   os: ubuntu-latest
+          #   rust: beta
+          #   target: x86_64-unknown-linux-gnu
+          # - build: nightly
+          #   os: ubuntu-latest
+          #   rust: nightly
+          #   target: x86_64-unknown-linux-gnu
+          # - build: linux32
+          #   os: ubuntu-latest
+          #   rust: stable
+          #   target: i686-unknown-linux-gnu
           - build: Anaconda on Linux
             os: ubuntu-latest
             rust: stable
             target: i686-unknown-linux-gnu
-          - build: macos
-            os: macos-latest
-            rust: stable
-            target: x86_64-apple-darwin
-          - build: aarch64-ios
-            os: macos-latest
-            rust: stable
-            target: aarch64-apple-ios
-          - build: windows-aarch64
-            os: windows-latest
-            rust: stable
-            target: aarch64-pc-windows-msvc
-          - build: win64
-            os: windows-latest
-            rust: stable
-            target: x86_64-pc-windows-msvc
-          - build: windows-latest
-            os: windows-latest
-            rust: stable-x86_64
-            target: x86_64-pc-windows-msvc
+          # - build: macos
+          #   os: macos-latest
+          #   rust: stable
+          #   target: x86_64-apple-darwin
+          # - build: aarch64-ios
+          #   os: macos-latest
+          #   rust: stable
+          #   target: aarch64-apple-ios
+          # - build: windows-aarch64
+          #   os: windows-latest
+          #   rust: stable
+          #   target: aarch64-pc-windows-msvc
+          # - build: win64
+          #   os: windows-latest
+          #   rust: stable
+          #   target: x86_64-pc-windows-msvc
+          # - build: windows-latest
+          #   os: windows-latest
+          #   rust: stable-x86_64
+          #   target: x86_64-pc-windows-msvc
     steps:
     - uses: actions/checkout@master
     - run: git submodule init
@@ -62,7 +62,9 @@ jobs:
       run: pip install matplotlib
     - name: Install Matplotlib (Anaconda)
       if: startsWith(matrix.build, 'Anaconda')
-      run: $CONDA/bin/conda install conda-forge::matplotlib
+      run: |
+        $CONDA/bin/conda install conda-forge::matplotlib
+        python3 -c "import sys; print('\n'.join(sys.path))"
     - name: Install Rust (rustup)
       run: rustup update ${{ matrix.rust }} --no-self-update && rustup default ${{ matrix.rust }}
       shell: bash

diff --git a/build.rs b/build.rs
@@ -0,0 +1,35 @@
+use std::{env,
+    fs::File,
+    io::Write,
+    path::Path,
+    process::Command,
+};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let outdir = env::var("OUT_DIR")
+        .expect("The OUT_DIR variable must be defined during compilation");
+    let py = Path::new(&outdir).join("matplotlib_location.py");
+    let mut fh = File::create(&py)?;
+    writeln!(fh, "import matplotlib\nprint(matplotlib.__file__)")?;
+    drop(fh);
+
+    let plt_path = Command::new("python3")
+        .arg(py)
+        .output()
+        .expect("`python3` or package `matplotlib` not found.")
+        .stdout;
+    let plt_path = String::from_utf8_lossy(&plt_path);
+    let plt_path = Path::new(plt_path.as_ref());
+    if let Some(d) = plt_path.parent() {
+        let d = d.to_str().unwrap();
+        println!("cargo:warning={d}");
+        println!("cargo:rustc-env=ANACONDA_DIR={d}");
+    }
+
+    // if let Ok(mut conda) = env::var("CONDA") {
+    //     conda += "/lib";
+    //     println!("cargo:rustc-env=ANACONDA_DIR={conda}");
+    // }
+
+    Ok(())
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -18,7 +18,7 @@ use pyo3::{
     prelude::*,
     intern,
     exceptions::{PyFileNotFoundError, PyPermissionError},
-    types::{PyDict, PyList},
+    types::{PyDict, PyList, PyString},
 };
 use numpy::{
     PyArray1,
@@ -102,10 +102,27 @@ If you use Anaconda, see https://github.com/PyO3/pyo3/issues/1554"),
 
 impl std::error::Error for Error {}
 
+/// Import the module `name`.  Try to update the Python search path if
+/// it fails.
+fn pyimport<N>(py: Python<'_>, name: N) -> Result<Py<PyModule>, PyErr>
+where N: IntoPy<Py<PyString>> + Clone {
+    match PyModule::import(py, name.clone()) {
+        Ok(m) => Ok(m.into()),
+        Err(e) => {
+            if let Some(d) = option_env!("ANACONDA_DIR") {
+                py.import("sys")?.getattr("path")?
+                    .call_method1("append", (d,));
+                PyModule::import(py, name).map(|m| m.into())
+            } else {
+                Err(e)
+            }
+        }
+    }
+}
+
 /// Import and return a handle to the module `$m`.
 macro_rules! pyimport { ($m: literal) => {
-    Python::with_gil(|py|
-        PyModule::import(py, intern!(py, $m)).map(|m| m.into()))
+    Python::with_gil(|py| pyimport(py, intern!(py, $m)))
 }}
 
 lazy_static! {