make datalayer a submodule for the namespacing benefits

Chia-Network · Jan 7, 2025 · ff365a2 · ff365a2
1 parent 4362b8c
commit ff365a2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/tests/test_datalayer.py b/tests/test_datalayer.py
@@ -1,4 +1,4 @@
-from chia_rs import LeafNode, MerkleBlob
+from chia_rs.datalayer import LeafNode, MerkleBlob
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import int64, uint8
 

diff --git a/wheel/src/api.rs b/wheel/src/api.rs
@@ -55,7 +55,7 @@ use pyo3::pybacked::PyBackedBytes;
 use pyo3::types::PyBytes;
 use pyo3::types::PyList;
 use pyo3::types::PyTuple;
-use pyo3::wrap_pyfunction;
+use pyo3::{py_run, wrap_pyfunction, wrap_pymodule};
 use std::collections::HashSet;
 use std::iter::zip;
 
@@ -452,7 +452,7 @@ pub fn py_get_flags_for_height_and_constants(height: u32, constants: &ConsensusC
 }
 
 #[pymodule]
-pub fn chia_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+pub fn chia_rs(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     // generator functions
     m.add_function(wrap_pyfunction!(run_block_generator, m)?)?;
     m.add_function(wrap_pyfunction!(run_block_generator2, m)?)?;
@@ -475,11 +475,6 @@ pub fn chia_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     // constants
     m.add_class::<ConsensusConstants>()?;
 
-    // datalayer
-    m.add_class::<MerkleBlob>()?;
-    m.add_class::<InternalNode>()?;
-    m.add_class::<LeafNode>()?;
-
     // merkle tree
     m.add_class::<MerkleSet>()?;
     m.add_function(wrap_pyfunction!(confirm_included_already_hashed, m)?)?;
@@ -647,5 +642,24 @@ pub fn chia_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<AugSchemeMPL>()?;
     m.add_class::<BlsCache>()?;
 
+    add_datalayer_submodule(py, m)?;
+
+    Ok(())
+}
+
+pub fn add_datalayer_submodule(py: Python<'_>, parent: &Bound<'_, PyModule>) -> PyResult<()> {
+    let datalayer = PyModule::new(py, "datalayer")?;
+    parent.add_submodule(&datalayer)?;
+
+    datalayer.add_class::<MerkleBlob>()?;
+    datalayer.add_class::<InternalNode>()?;
+    datalayer.add_class::<LeafNode>()?;
+
+    // https://github.com/PyO3/pyo3/issues/1517#issuecomment-808664021
+    // https://github.com/PyO3/pyo3/issues/759
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item("chia_rs.datalayer", datalayer)?;
+
     Ok(())
 }