Merge pull request #225 from jessekrubin/pickling

Support picking of several structs

crates/ryo3-bytes/src/bytes_ext.rs

@@ -2,12 +2,17 @@
 
 use crate::bytes::PyBytes;
 use pyo3::prelude::*;
-use pyo3::types::{PyString, PyType};
+use pyo3::types::{PyString, PyTuple, PyType};
 use pyo3::IntoPyObjectExt;
 use std::hash::Hash;
 
 #[pymethods]
 impl PyBytes {
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        let pybytes = pyo3::types::PyBytes::new(py, self.as_ref()).into_bound_py_any(py)?;
+        PyTuple::new(py, vec![pybytes])
+    }
+
     /// Hash bytes
     fn __hash__(&self) -> u64 {
         // STD-HASHER VERSION

crates/ryo3-bytes/src/lib.rs

@@ -27,7 +27,7 @@ pub fn pymod_add(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyBytes>()?;
 
     // rename bytes module to `ry`
-    m.getattr("Bytes")?.setattr("__module__", "ryo3")?;
+    m.getattr("Bytes")?.setattr("__module__", "ry.ryo3")?;
 
     Ok(())
 }

crates/ryo3-fspath/src/fspath.rs

@@ -19,10 +19,9 @@ const MAIN_SEPARATOR: char = std::path::MAIN_SEPARATOR;
 
 type ArcPathBuf = std::sync::Arc<PathBuf>;
 
-#[pyclass(name = "FsPath", module = "ryo3", frozen)]
+#[pyclass(name = "FsPath", module = "ry", frozen)]
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct PyFsPath {
-    // pth: PathBuf,
     pth: ArcPathBuf,
 }
 
@@ -68,6 +67,11 @@ impl PyFsPath {
         }
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        let os_str = self.pth.as_os_str();
+        PyTuple::new(py, vec![os_str])
+    }
+
     fn string(&self) -> String {
         path2str(self.path())
     }

crates/ryo3-globset/src/lib.rs

@@ -294,7 +294,7 @@ fn glob(
     name = "globster",
     signature = (patterns, /, *, case_insensitive=None, literal_separator=None, backslash_escape=None)
 )]
-fn globster_fn(
+fn py_globster(
     patterns: StringOrStrings,
     case_insensitive: Option<bool>,
     literal_separator: Option<bool>,
@@ -311,7 +311,7 @@ fn globster_fn(
 
 pub fn pymod_add(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(glob, m)?)?;
-    m.add_function(wrap_pyfunction!(globster_fn, m)?)?;
+    m.add_function(wrap_pyfunction!(py_globster, m)?)?;
     m.add_class::<PyGlob>()?;
     m.add_class::<PyGlobSet>()?;
     m.add_class::<PyGlobster>()?;

crates/ryo3-http/Cargo.toml

@@ -14,6 +14,7 @@ categories.workspace = true
 [dependencies]
 http.workspace = true
 pyo3 = { workspace = true, features = ["experimental-inspect"] }
+ryo3-core.workspace = true
 ryo3-macros.workspace = true
 
 [build-dependencies]

crates/ryo3-http/src/headers.rs

@@ -4,10 +4,9 @@ use crate::PyHeadersLike;
 use http::header::HeaderMap;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
-use pyo3::types::PyString;
-use std::collections::HashMap;
+use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyTuple};
 
-#[pyclass(name = "Headers", module = "ry.ryo3.http")]
+#[pyclass(name = "Headers", module = "ry")]
 #[derive(Clone, Debug)]
 pub struct PyHeaders(pub HeaderMap);
 
@@ -21,25 +20,18 @@ impl From<HeaderMap> for PyHeaders {
 impl PyHeaders {
     #[new]
     #[pyo3(signature = (d = None))]
-    fn py_new(_py: Python<'_>, d: Option<HashMap<String, String>>) -> PyResult<Self> {
-        let mut headers = HeaderMap::new();
+    fn py_new(_py: Python<'_>, d: Option<PyHeadersLike>) -> PyResult<Self> {
         if let Some(d) = d {
-            for (k, v) in d {
-                let header_name =
-                    http::header::HeaderName::from_bytes(k.as_bytes()).map_err(|e| {
-                        PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                            "header-name-error: {e} (k={k}, v={v})"
-                        ))
-                    })?;
-                let header_value = http::header::HeaderValue::from_str(&v).map_err(|e| {
-                    PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                        "header-value-error: {e} (k={k}, v={v})"
-                    ))
-                })?;
-                headers.insert(header_name, header_value);
-            }
+            let headers_map = HeaderMap::try_from(d)?;
+            Ok(Self(headers_map))
+        } else {
+            Ok(PyHeaders(HeaderMap::new()))
         }
-        Ok(Self(headers))
+    }
+
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        let dict = self.asdict(py)?;
+        PyTuple::new(py, vec![dict])
     }
 
     /// Return struct Debug-string
@@ -278,16 +270,35 @@ impl PyHeaders {
         Ok(PyHeaders(headers))
     }
 
-    // pub fn __ior__<'py>(
-    //     mut slf: PyRefMut<'py, Self>,
-    //     other: &PyHeaders,
-    // ) -> PyResult<PyRefMut<'py, Self>> {
-    //     let other_headers = other.0.clone();
-    //     for (k, v) in other_headers {
-    //         if let Some(k) = k {
-    //             slf.0.insert(k, v);
-    //         }
-    //     }
-    //     Ok(slf)
-    // }
+    fn asdict<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        let d = PyDict::new(py);
+
+        for key in self.0.keys() {
+            let key_str = key.as_str();
+            let values: Vec<_> = self.0.get_all(key).iter().collect();
+
+            if values.len() == 1 {
+                let v = values[0];
+                if let Ok(vstr) = v.to_str() {
+                    d.set_item(key_str, vstr)?;
+                } else {
+                    let pybytes = PyBytes::new(py, v.as_bytes());
+                    d.set_item(key_str, pybytes)?;
+                }
+            } else {
+                let py_list = PyList::empty(py);
+                for v in values {
+                    if let Ok(vstr) = v.to_str() {
+                        py_list.append(vstr)?;
+                    } else {
+                        let pybytes = PyBytes::new(py, v.as_bytes());
+                        py_list.append(pybytes)?;
+                    }
+                }
+                d.set_item(key_str, py_list)?;
+            }
+        }
+
+        Ok(d)
+    }
 }

crates/ryo3-http/src/headers_like.rs

@@ -2,23 +2,38 @@ use crate::PyHeaders;
 use http::HeaderMap;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use ryo3_core::types::StringOrStrings;
 use std::collections::HashMap;
 
 #[derive(Debug, Clone, FromPyObject)]
 pub enum PyHeadersLike {
     Headers(PyHeaders),
-    Map(HashMap<String, String>),
+    Map(HashMap<String, StringOrStrings>),
 }
 
 impl PyHeadersLike {
-    pub fn map2headers(d: &HashMap<String, String>) -> PyResult<HeaderMap> {
+    pub fn map2headers(d: &HashMap<String, StringOrStrings>) -> PyResult<HeaderMap> {
         let mut headers = HeaderMap::new();
         for (k, v) in d {
-            let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
-                .map_err(|e| PyValueError::new_err(format!("header-name-error: {e}")))?;
-            let header_value = http::header::HeaderValue::from_str(v)
-                .map_err(|e| PyValueError::new_err(format!("header-value-error: {e}")))?;
-            headers.insert(header_name, header_value);
+            match v {
+                StringOrStrings::String(s) => {
+                    let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
+                        .map_err(|e| PyValueError::new_err(format!("header-name-error: {e}")))?;
+                    let header_value = http::header::HeaderValue::from_str(s)
+                        .map_err(|e| PyValueError::new_err(format!("header-value-error: {e}")))?;
+                    headers.insert(header_name, header_value);
+                }
+                StringOrStrings::Strings(v) => {
+                    let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
+                        .map_err(|e| PyValueError::new_err(format!("header-name-error: {e}")))?;
+                    for s in v {
+                        let header_value = http::header::HeaderValue::from_str(s).map_err(|e| {
+                            PyValueError::new_err(format!("header-value-error: {e}"))
+                        })?;
+                        headers.append(&header_name, header_value);
+                    }
+                }
+            }
         }
         Ok(headers)
     }
@@ -33,12 +48,32 @@ impl TryFrom<PyHeadersLike> for HeaderMap {
                 let mut default_headers = HeaderMap::new();
                 for (k, v) in d {
                     let k = k.to_string();
-                    let v = v.to_string();
-                    let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
-                        .map_err(|e| PyValueError::new_err(format!("header-name-error: {e}")))?;
-                    let header_value = http::header::HeaderValue::from_str(&v)
-                        .map_err(|e| PyValueError::new_err(format!("header-value-error: {e}")))?;
-                    default_headers.insert(header_name, header_value);
+                    match v {
+                        StringOrStrings::String(s) => {
+                            let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
+                                .map_err(|e| {
+                                    PyValueError::new_err(format!("header-name-error: {e}"))
+                                })?;
+                            let header_value =
+                                http::header::HeaderValue::from_str(&s).map_err(|e| {
+                                    PyValueError::new_err(format!("header-value-error: {e}"))
+                                })?;
+                            default_headers.insert(header_name, header_value);
+                        }
+                        StringOrStrings::Strings(v) => {
+                            let header_name = http::header::HeaderName::from_bytes(k.as_bytes())
+                                .map_err(|e| {
+                                    PyValueError::new_err(format!("header-name-error: {e}"))
+                                })?;
+                            for s in v {
+                                let header_value = http::header::HeaderValue::from_str(&s)
+                                    .map_err(|e| {
+                                        PyValueError::new_err(format!("header-value-error: {e}"))
+                                    })?;
+                                default_headers.append(&header_name, header_value);
+                            }
+                        }
+                    }
                 }
                 Ok(default_headers)
             }

crates/ryo3-http/src/status_code.rs

@@ -2,6 +2,7 @@
 
 use pyo3::prelude::*;
 use pyo3::pyclass::CompareOp;
+use pyo3::types::PyTuple;
 
 #[pyclass(name = "HttpStatus", module = "ry.ryo3.http", frozen)]
 #[derive(Clone, Debug)]
@@ -17,6 +18,10 @@ impl PyHttpStatus {
         })?))
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(py, [self.0.as_u16()])
+    }
+
     #[must_use]
     pub fn __str__(&self) -> String {
         format!("{:?}", self.0)

crates/ryo3-jiff/src/ry_date.rs

@@ -23,7 +23,7 @@ use std::borrow::BorrowMut;
 use std::fmt::Display;
 use std::hash::{DefaultHasher, Hash, Hasher};
 #[derive(Debug, Clone)]
-#[pyclass(name = "Date", module = "ryo3", frozen)]
+#[pyclass(name = "Date", module = "ry", frozen)]
 pub struct RyDate(pub(crate) Date);
 
 #[pymethods]
@@ -123,6 +123,17 @@ impl RyDate {
         )
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(
+            py,
+            vec![
+                self.year().into_pyobject(py)?,
+                self.month().into_pyobject(py)?,
+                self.day().into_pyobject(py)?,
+            ],
+        )
+    }
+
     fn sub_date(&self, other: &RyDate) -> RySpan {
         RySpan::from(self.0 - other.0)
     }

crates/ryo3-jiff/src/ry_datetime.rs

@@ -13,15 +13,15 @@ use crate::{JiffEraYear, JiffRoundMode, JiffUnit, JiffWeekday, RyDate, RyDateTim
 use jiff::civil::{DateTime, DateTimeRound, Weekday};
 use jiff::Zoned;
 use pyo3::basic::CompareOp;
-use pyo3::intern;
 use pyo3::prelude::*;
-use pyo3::types::{PyDateTime, PyDict, PyType};
+use pyo3::types::{PyDateTime, PyDict, PyTuple, PyType};
+use pyo3::{intern, IntoPyObjectExt};
 use std::borrow::BorrowMut;
 use std::fmt::Display;
 use std::hash::{DefaultHasher, Hash, Hasher};
 use std::str::FromStr;
 #[derive(Debug, Clone)]
-#[pyclass(name = "DateTime", module = "ryo3", frozen)]
+#[pyclass(name = "DateTime", module = "ry", frozen)]
 pub struct RyDateTime(pub(crate) DateTime);
 
 impl From<DateTime> for RyDateTime {
@@ -56,6 +56,21 @@ impl RyDateTime {
         .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{e}")))
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(
+            py,
+            vec![
+                self.year().into_bound_py_any(py)?,
+                self.month().into_bound_py_any(py)?,
+                self.day().into_bound_py_any(py)?,
+                self.hour().into_bound_py_any(py)?,
+                self.minute().into_bound_py_any(py)?,
+                self.second().into_bound_py_any(py)?,
+                self.subsec_nanosecond().into_bound_py_any(py)?,
+            ],
+        )
+    }
+
     #[expect(non_snake_case)]
     #[classattr]
     fn MIN() -> Self {

crates/ryo3-jiff/src/ry_iso_week_date.rs

@@ -3,11 +3,11 @@ use crate::{JiffWeekday, RyDate};
 use jiff::civil::ISOWeekDate;
 use jiff::Zoned;
 use pyo3::prelude::*;
-use pyo3::types::PyType;
+use pyo3::types::{PyTuple, PyType};
 use std::hash::{DefaultHasher, Hash, Hasher};
 
 #[derive(Debug, Clone)]
-#[pyclass(name = "ISOWeekDate", module = "ryo3", frozen)]
+#[pyclass(name = "ISOWeekDate", module = "ry", frozen)]
 pub struct RyISOWeekDate(pub(crate) ISOWeekDate);
 
 #[pymethods]
@@ -19,6 +19,17 @@ impl RyISOWeekDate {
             .map_err(map_py_value_err)
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(
+            py,
+            vec![
+                self.year().into_pyobject(py)?,
+                self.week().into_pyobject(py)?,
+                self.weekday().into_pyobject(py)?,
+            ],
+        )
+    }
+
     // ========================================================================
     // CLASSATTR
     // ========================================================================