diff --git a/spk/__init__.py b/spk/__init__.py index 3a991b4da3..8497a101c0 100644 --- 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 index 211a06499b..bb26a473cc 100644 --- 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 index 3d110c59d8..421108b9d7 100644 --- 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::()?; module.add_class::()?; module.add_class::()?; @@ -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::()?; 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::()?; 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::()?; m.add_class::()?;