Hack to make submodules importable

via `from spk.solve import ...` PyO3/pyo3#1517 (comment) Signed-off-by: J Robert Ray <[email protected]>
spkenv · Aug 8, 2021 · 11a928c · 11a928c
1 parent d1e602d
commit 11a928c
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 30 deletions.
diff --git a/spk/__init__.py b/spk/__init__.py
@@ -9,7 +9,8 @@
 
 from . import pysolve
 
-solve = spkrs.solve
+from spk import solve
+
 solve.legacy = pysolve.legacy
 
 from . import storage, build, exec, test

diff --git a/src/lib.rs b/src/lib.rs
@@ -17,7 +17,7 @@ pub use error::{Error, Result};
 
 // -- begin python wrappers --
 
-use pyo3::prelude::*;
+use pyo3::{prelude::*, py_run};
 use spfs::{self, prelude::*};
 
 #[pyclass]
@@ -76,6 +76,12 @@ fn spkrs(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_submodule(build_mod)?;
 
     let solve_mod = PyModule::new(py, "solve")?;
+    // Hack to make `from spk.solve import ...` work
+    py_run!(
+        py,
+        solve_mod,
+        "import sys; sys.modules['spk.solve'] = solve_mod"
+    );
     solve::init_module(&py, solve_mod)?;
     m.add_submodule(solve_mod)?;
 

diff --git a/src/solve/python.rs b/src/solve/python.rs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // https://github.com/imageworks/spk
 use pyo3::prelude::*;
+use pyo3::py_run;
 
 use super::errors::SolverError;
 use super::graph::{
@@ -17,7 +18,7 @@ fn init_submodule_errors(py: &Python, module: &PyModule) -> PyResult<()> {
     Ok(())
 }
 
-fn init_submodule_graph(module: &PyModule) -> PyResult<()> {
+fn init_submodule_graph(_py: &Python, module: &PyModule) -> PyResult<()> {
     module.add_class::<Change>()?;
     module.add_class::<Decision>()?;
     module.add_class::<Graph>()?;
@@ -33,7 +34,7 @@ fn init_submodule_graph(module: &PyModule) -> PyResult<()> {
     Ok(())
 }
 
-fn init_submodule_solution(module: &PyModule) -> PyResult<()> {
+fn init_submodule_solution(_py: &Python, module: &PyModule) -> PyResult<()> {
     module.add_class::<Solution>()?;
     Ok(())
 }
@@ -44,37 +45,34 @@ fn init_submodule_solver(py: &Python, module: &PyModule) -> PyResult<()> {
     Ok(())
 }
 
-fn init_submodule_validation(module: &PyModule) -> PyResult<()> {
+fn init_submodule_validation(_py: &Python, module: &PyModule) -> PyResult<()> {
     module.add_class::<Validator>()?;
     Ok(())
 }
 
+macro_rules! add_submodule {
+    ($m:ident, $py:ident, $mod_name:expr, $init_fn:ident) => {
+        let submod = PyModule::new(*$py, $mod_name)?;
+        // Hack to make `from spk.solve.foo import ...` work
+        py_run!(
+            *$py,
+            submod,
+            &format!(
+                "import sys; sys.modules['spk.solve.{}'] = submod",
+                $mod_name
+            )
+        );
+        $init_fn($py, submod)?;
+        $m.add_submodule(submod)?;
+    };
+}
+
 pub fn init_module(py: &Python, m: &PyModule) -> PyResult<()> {
-    {
-        let submod_errors = PyModule::new(*py, "_errors")?;
-        init_submodule_errors(py, submod_errors)?;
-        m.add_submodule(submod_errors)?;
-    }
-    {
-        let submod_graph = PyModule::new(*py, "graph")?;
-        init_submodule_graph(submod_graph)?;
-        m.add_submodule(submod_graph)?;
-    }
-    {
-        let submod_solver = PyModule::new(*py, "_solver")?;
-        init_submodule_solver(py, submod_solver)?;
-        m.add_submodule(submod_solver)?;
-    }
-    {
-        let submod_solution = PyModule::new(*py, "_solution")?;
-        init_submodule_solution(submod_solution)?;
-        m.add_submodule(submod_solution)?;
-    }
-    {
-        let submod_validation = PyModule::new(*py, "validation")?;
-        init_submodule_validation(submod_validation)?;
-        m.add_submodule(submod_validation)?;
-    }
+    add_submodule!(m, py, "_errors", init_submodule_errors);
+    add_submodule!(m, py, "graph", init_submodule_graph);
+    add_submodule!(m, py, "_solver", init_submodule_solver);
+    add_submodule!(m, py, "_solution", init_submodule_solution);
+    add_submodule!(m, py, "validation", init_submodule_validation);
 
     m.add_class::<Graph>()?;
     m.add_class::<Solution>()?;