Refactor rust code to submodules (#776)

Author: edenhaus

deebot_client/commands/json/map.py

@@ -18,7 +18,7 @@
 from deebot_client.events.map import CachedMapInfoEvent
 from deebot_client.logging_filter import get_logger
 from deebot_client.message import HandlingResult, HandlingState, MessageBodyDataDict
-from deebot_client.rs import decompress_7z_base64_data
+from deebot_client.rs.util import decompress_7z_base64_data
 
 from .common import JsonCommandWithMessageHandling
 

deebot_client/map.py

@@ -34,7 +34,8 @@
 from .exceptions import MapError
 from .logging_filter import get_logger
 from .models import Room
-from .rs import TracePoint, decompress_7z_base64_data, extract_trace_points
+from .rs.map import TracePoint, extract_trace_points
+from .rs.util import decompress_7z_base64_data
 from .util import (
     OnChangedDict,
     OnChangedList,

deebot_client/rs.pyi renamed to deebot_client/rs/map.pyi

@@ -1,8 +1,5 @@
 from typing import Self
 
-def decompress_7z_base64_data(value: str) -> bytes:
-    """Decompress base64 decoded 7z compressed string."""
-
 class TracePoint:
     """Trace point."""
 

deebot_client/rs/util.pyi

@@ -0,0 +1,2 @@
+def decompress_7z_base64_data(value: str) -> bytes:
+    """Decompress base64 decoded 7z compressed string."""

pyproject.toml

@@ -148,7 +148,7 @@ required-imports = ["from __future__ import annotations"]
 "deebot_client/hardware/deebot/*" = [
     "N999", # Invalid module name
 ]
-"deebot_client/rs.pyi" = [
+"deebot_client/rs/*.pyi" = [
     "PYI021", #  docstring-in-stub
 ]
 
@@ -167,7 +167,11 @@ max-args = 7
 py-version = "3.13"
 ignore = ["tests"]
 fail-on = ["I"]
-extension-pkg-allow-list = ["deebot_client.rs"]
+extension-pkg-allow-list = [
+    "deebot_client.rs",
+    "deebot_client.rs.map",
+    "deebot_client.rs.util",
+]
 
 [tool.pylint.BASIC]
 good-names = ["i", "j", "k", "ex", "_", "T", "x", "y", "id", "tg"]

src/lib.rs

@@ -1,84 +1,31 @@
-use std::error::Error;
-
-use base64::{engine::general_purpose, Engine as _};
-use byteorder::{LittleEndian, ReadBytesExt};
-use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
-use std::io::Cursor;
-
-fn _decompress_7z_base64_data(value: String) -> Result<Vec<u8>, Box<dyn Error>> {
-    let mut bytes = general_purpose::STANDARD.decode(value)?;
-
-    if bytes.len() < 8 {
-        return Err("Invalid 7z compressed data".into());
-    }
-
-    for _ in 0..=3 {
-        bytes.insert(8, 0);
-    }
-
-    Ok(lzma::decompress(&bytes)?)
-}
-
-/// Decompress base64 decoded 7z compressed string.
-#[pyfunction]
-fn decompress_7z_base64_data(value: String) -> Result<Vec<u8>, PyErr> {
-    Ok(_decompress_7z_base64_data(value).map_err(|err| PyValueError::new_err(err.to_string()))?)
-}
-
-/// Trace point
-#[pyclass]
-struct TracePoint {
-    #[pyo3(get)]
-    x: i16,
-
-    #[pyo3(get)]
-    y: i16,
-
-    #[pyo3(get)]
-    connected: bool,
-}
-
-#[pymethods]
-impl TracePoint {
-    #[new]
-    fn new(x: i16, y: i16, connected: bool) -> Self {
-        TracePoint { x, y, connected }
-    }
-}
 
-fn process_trace_points(trace_points: &[u8]) -> Result<Vec<TracePoint>, Box<dyn Error>> {
-    let mut trace_values = Vec::new();
-    for i in (0..trace_points.len()).step_by(5) {
-        if i + 4 >= trace_points.len() {
-            return Err("Invalid trace points length".into());
-        }
-
-        let mut cursor = Cursor::new(&trace_points[i..i + 4]);
-        let x = cursor.read_i16::<LittleEndian>()?;
-        let y = cursor.read_i16::<LittleEndian>()?;
-
-        // Determine connection status
-        let connected = (trace_points[i + 4] >> 7 & 1) == 0;
-
-        trace_values.push(TracePoint { x, y, connected });
-    }
-    Ok(trace_values)
-}
-
-#[pyfunction]
-/// Extract trace points from 7z compressed data string.
-fn extract_trace_points(value: String) -> Result<Vec<TracePoint>, PyErr> {
-    let decompressed_data = decompress_7z_base64_data(value)?;
-    Ok(process_trace_points(&decompressed_data)
-        .map_err(|err| PyValueError::new_err(err.to_string()))?)
-}
+mod map;
+mod util;
 
 /// Deebot client written in Rust
 #[pymodule]
 fn rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(decompress_7z_base64_data, m)?)?;
-    m.add_function(wrap_pyfunction!(extract_trace_points, m)?)?;
-    m.add_class::<TracePoint>()?;
+    register_child_module(m, "map", map::init_module)?;
+    register_child_module(m, "util", util::init_module)?;
     Ok(())
 }
+
+fn register_child_module(
+    parent_module: &Bound<'_, PyModule>,
+    name: &str,
+    func: fn(&Bound<'_, PyModule>) -> PyResult<()>,
+) -> PyResult<()> {
+    let child_module = PyModule::new(parent_module.py(), name)?;
+    func(&child_module)?;
+
+    // https://github.com/PyO3/pyo3/issues/1517#issuecomment-808664021
+    // https://github.com/PyO3/pyo3/issues/759
+    let _ = Python::with_gil(|py| {
+        py.import("sys")?
+            .getattr("modules")?
+            .set_item(&format!("deebot_client.rs.{}", name), &child_module)
+    });
+
+    parent_module.add_submodule(&child_module)
+}

src/map.rs

@@ -0,0 +1,64 @@
+use std::error::Error;
+
+use super::util::decompress_7z_base64_data;
+use byteorder::{LittleEndian, ReadBytesExt};
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use std::io::Cursor;
+
+/// Trace point
+#[pyclass]
+struct TracePoint {
+    #[pyo3(get)]
+    x: i16,
+
+    #[pyo3(get)]
+    y: i16,
+
+    #[pyo3(get)]
+    connected: bool,
+}
+
+#[pymethods]
+impl TracePoint {
+    #[new]
+    fn new(x: i16, y: i16, connected: bool) -> Self {
+        TracePoint { x, y, connected }
+    }
+}
+
+fn process_trace_points(trace_points: &[u8]) -> Result<Vec<TracePoint>, Box<dyn Error>> {
+    let mut trace_values = Vec::new();
+    for i in (0..trace_points.len()).step_by(5) {
+        if i + 4 >= trace_points.len() {
+            return Err("Invalid trace points length".into());
+        }
+
+        let mut cursor = Cursor::new(&trace_points[i..i + 4]);
+        let x = cursor.read_i16::<LittleEndian>()?;
+        let y = cursor.read_i16::<LittleEndian>()?;
+
+        // Determine connection status
+        let connected = (trace_points[i + 4] >> 7 & 1) == 0;
+
+        trace_values.push(TracePoint { x, y, connected });
+    }
+    Ok(trace_values)
+}
+
+fn extract_trace_points(value: String) -> Result<Vec<TracePoint>, Box<dyn Error>> {
+    let decompressed_data = decompress_7z_base64_data(value)?;
+    Ok(process_trace_points(&decompressed_data)?)
+}
+
+#[pyfunction(name = "extract_trace_points")]
+/// Extract trace points from 7z compressed data string.
+fn python_extract_trace_points(value: String) -> Result<Vec<TracePoint>, PyErr> {
+    Ok(extract_trace_points(value).map_err(|err| PyValueError::new_err(err.to_string()))?)
+}
+
+pub fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(python_extract_trace_points, m)?)?;
+    m.add_class::<TracePoint>()?;
+    Ok(())
+}

src/util.rs

@@ -0,0 +1,30 @@
+use std::error::Error;
+
+use base64::{engine::general_purpose, Engine as _};
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+
+pub fn decompress_7z_base64_data(value: String) -> Result<Vec<u8>, Box<dyn Error>> {
+    let mut bytes = general_purpose::STANDARD.decode(value)?;
+
+    if bytes.len() < 8 {
+        return Err("Invalid 7z compressed data".into());
+    }
+
+    for _ in 0..=3 {
+        bytes.insert(8, 0);
+    }
+
+    Ok(lzma::decompress(&bytes)?)
+}
+
+/// Decompress base64 decoded 7z compressed string.
+#[pyfunction(name = "decompress_7z_base64_data")]
+fn python_decompress_7z_base64_data(value: String) -> Result<Vec<u8>, PyErr> {
+    Ok(decompress_7z_base64_data(value).map_err(|err| PyValueError::new_err(err.to_string()))?)
+}
+
+pub fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(python_decompress_7z_base64_data, m)?)?;
+    Ok(())
+}

tests/rs/__init__.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from deebot_client.rs import decompress_7z_base64_data
+from deebot_client.rs.util import decompress_7z_base64_data
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager

tests/test_rs.py renamed to tests/rs/test_util.py

@@ -48,7 +48,7 @@
     _points_to_svg_path,
 )
 from deebot_client.models import Room
-from deebot_client.rs import TracePoint
+from deebot_client.rs.map import TracePoint
 
 from .common import block_till_done