quarb_pyvalue/
foreign.rs

1//! A Python object as a library value (spec: Command Libraries, "A
2//! library value"): what a Python-wrapping library keeps in its own
3//! form between two commands — a frame, an array, any object that is
4//! not plain data. The engine sees its kind, its text and its two
5//! exports; the libraries take the object back as it is.
6
7use pyo3::prelude::*;
8use pyo3::types::PyDict;
9use quarb::{Foreign, Value};
10use std::ffi::CString;
11use std::sync::OnceLock;
12
13/// The Python side of the two readings and the kind. Numpy and
14/// pandas are recognised by the module their types come from, so
15/// neither is imported for an object that is not theirs.
16const GLUE: &str = r#"
17import math
18
19def _mod(o):
20    return type(o).__module__.split(".")[0]
21
22def _plain(x):
23    """A value out of numpy/pandas as plain Python: scalars by item(),
24    timestamps as datetimes, NaT as None; containers recursively."""
25    m = _mod(x)
26    if m == "numpy":
27        if hasattr(x, "tolist"):
28            return _plain(x.tolist()) if hasattr(x, "shape") and x.shape != () else x.item()
29        return x
30    if m == "pandas":
31        t = type(x).__name__
32        if t == "Timestamp":
33            return x.to_pydatetime()
34        if t == "Timedelta":
35            return x.to_pytimedelta()
36        if t == "NaTType":
37            return None
38        if t in ("Period", "Interval"):
39            return str(x)
40    if isinstance(x, dict):
41        return {str(k): _plain(v) for k, v in x.items()}
42    if isinstance(x, (list, tuple)):
43        return [_plain(v) for v in x]
44    return x
45
46def _kind(o):
47    m, t = _mod(o), type(o).__name__
48    if m == "pandas" and t == "DataFrame":
49        r, c = o.shape
50        return f"DataFrame {r}×{c}"
51    if m == "pandas" and t == "Series":
52        return f"Series {len(o)}"
53    if m == "numpy" and hasattr(o, "shape"):
54        return f"ndarray {o.shape} {o.dtype}"
55    if m == "scipy" and hasattr(o, "toarray"):
56        return f"sparse {t} {o.shape}"
57    return type(o).__qualname__
58
59def _records(df):
60    import pandas as pd
61    if isinstance(df.index, pd.RangeIndex):
62        return _plain(df.to_dict("records"))
63    name = df.index.name
64    cols = [str(c) for c in df.columns]
65    out = []
66    for idx, row in zip(df.index.tolist(), df.itertuples(index=False, name=None)):
67        rec = {name if name is not None else "index": idx}
68        for c, v in zip(cols, row):
69            rec[c] = v
70        out.append(rec)
71    return _plain(out)
72
73def _rows(o):
74    m, t = _mod(o), type(o).__name__
75    if m == "pandas" and t == "DataFrame":
76        return _records(o)
77    if m == "pandas" and t in ("Series", "Index"):
78        return _plain(o.tolist())
79    if m == "numpy" and hasattr(o, "tolist"):
80        return _plain(o.tolist())
81    if m == "scipy" and hasattr(o, "toarray"):
82        return _plain(o.toarray().tolist())
83    if hasattr(o, "_fields"):
84        names = list(o._fields) + list(getattr(o, "_extra_fields", ()))
85        return {n: _plain(getattr(o, n)) for n in names}
86    if hasattr(o, "__iter__") and not hasattr(o, "read"):
87        return [_plain(v) for v in o]
88    raise TypeError(f"{_kind(o)} has no row reading")
89
90def _table(o):
91    """A DataFrame or a 2-d ndarray as a table (header and rows, a
92    null None), a 1-d ndarray or a Series on a RangeIndex as a sample
93    (one column, no header), None for anything else. A RangeIndex is left out; another index
94    is the first column, by its name or `index`."""
95    m, t = _mod(o), type(o).__name__
96    def cell(v):
97        # null is the one missing value: None, NaT and pd.NA come back
98        # as null; a nan is a float and stays one.
99        try:
100            import pandas as pd
101            if v is pd.NA or v is pd.NaT:
102                return None
103        except ImportError:
104            pass
105        v = _plain(v)
106        if v is None:
107            return None
108        try:
109            import pandas as pd
110            if v is pd.NA or v is pd.NaT:
111                return None
112        except ImportError:
113            pass
114        return v
115    if m == "pandas" and t == "DataFrame":
116        import pandas as pd
117        cols = [str(c) for c in o.columns]
118        rows = [[cell(v) for v in r] for r in o.itertuples(index=False, name=None)]
119        if not isinstance(o.index, pd.RangeIndex):
120            name = str(o.index.name) if o.index.name is not None else "index"
121            cols = [name] + cols
122            rows = [[cell(i)] + r for i, r in zip(o.index.tolist(), rows)]
123        return {"header": cols, "rows": rows}
124    if m == "pandas" and t == "Series":
125        import pandas as pd
126        if isinstance(o.index, pd.RangeIndex):
127            return {"header": None, "rows": [[cell(v)] for v in o.tolist()]}
128        name = str(o.name) if o.name is not None else "value"
129        iname = str(o.index.name) if o.index.name is not None else "index"
130        return {"header": [iname, name],
131                "rows": [[cell(i), cell(v)] for i, v in zip(o.index.tolist(), o.tolist())]}
132    if m == "numpy" and hasattr(o, "ndim"):
133        if o.ndim == 2:
134            return {"header": None, "rows": [[cell(v) for v in r] for r in o.tolist()]}
135        if o.ndim == 1:
136            return {"header": None, "rows": [[cell(v)] for v in o.tolist()]}
137    return None
138
139def _columns(o):
140    m, t = _mod(o), type(o).__name__
141    if m == "pandas" and t == "DataFrame":
142        import pandas as pd
143        if isinstance(o.index, pd.RangeIndex):
144            return _plain({str(c): o[c].tolist() for c in o.columns})
145        return _plain({str(c): {str(k): v for k, v in o[c].to_dict().items()} for c in o.columns})
146    if m == "pandas" and t == "Series":
147        return _plain({str(k): v for k, v in o.to_dict().items()})
148    if m == "numpy":
149        raise TypeError(f"{_kind(o)} has no names; '| @' reads its rows, pd: keeps names")
150    if hasattr(o, "_fields"):
151        return _rows(o)
152    if hasattr(o, "items"):
153        return {str(k): _plain(v) for k, v in o.items()}
154    raise TypeError(f"{_kind(o)} has no column reading")
155"#;
156
157/// The site-packages the embedded interpreter would miss: an
158/// active virtualenv's (`VIRTUAL_ENV`), and the prefix the running
159/// binary was installed under (a wheel's `bin/qua` beside its
160/// `lib/pythonX.Y/site-packages`), added once, after the
161/// interpreter's own, through `site.addsitedir` so `.pth` files
162/// count too.
163pub(crate) fn extend_sys_path(py: Python<'_>) -> PyResult<()> {
164    static DONE: OnceLock<()> = OnceLock::new();
165    if DONE.get().is_some() {
166        return Ok(());
167    }
168    DONE.set(()).ok();
169    let exe_prefix = std::env::current_exe()
170        .ok()
171        .and_then(|p| p.parent().and_then(|b| b.parent()).map(|p| p.display().to_string()))
172        .unwrap_or_default();
173    let code = c"
174import os, site, sys
175def _extend(prefix):
176    if not prefix:
177        return
178    v = f'python{sys.version_info[0]}.{sys.version_info[1]}'
179    for d in (os.path.join(prefix, 'lib', v, 'site-packages'),
180              os.path.join(prefix, 'lib64', v, 'site-packages'),
181              os.path.join(prefix, 'Lib', 'site-packages')):
182        if os.path.isdir(d) and d not in sys.path:
183            site.addsitedir(d)
184_extend(os.environ.get('VIRTUAL_ENV'))
185_extend(_EXE_PREFIX)
186";
187    // One namespace for the code: the function body resolves its
188    // imports in the globals it was defined under.
189    let ns = PyDict::new(py);
190    ns.set_item("_EXE_PREFIX", exe_prefix)?;
191    py.run(code, Some(&ns), Some(&ns))?;
192    Ok(())
193}
194
195static MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
196
197fn module<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyModule>> {
198    if let Some(m) = MODULE.get() {
199        return Ok(m.bind(py).clone());
200    }
201    extend_sys_path(py)?;
202    let m = PyModule::from_code(
203        py,
204        &CString::new(GLUE).expect("no NUL"),
205        c"quarb_foreign.py",
206        c"quarb_foreign",
207    )?;
208    // Importable as `quarb_foreign` from every other glue page.
209    py.import("sys")?
210        .getattr("modules")?
211        .set_item("quarb_foreign", &m)?;
212    let _ = MODULE.set(m.clone().unbind());
213    Ok(m)
214}
215
216/// Compile the helpers (once) so a glue page can import them.
217pub(crate) fn register(py: Python<'_>) -> PyResult<()> {
218    module(py).map(|_| ())
219}
220
221/// Whether a Python object is plain data for the engine (converted
222/// as it crosses) or a library value (kept as it is): the builtins
223/// and the standard containers convert, numpy scalars convert by
224/// their item, everything else is foreign.
225pub(crate) fn is_plain(obj: &Bound<'_, PyAny>) -> bool {
226    let module = obj
227        .get_type()
228        .getattr("__module__")
229        .ok()
230        .and_then(|m| m.extract::<String>().ok())
231        .unwrap_or_default();
232    matches!(
233        module.split('.').next().unwrap_or(""),
234        "builtins" | "collections" | "itertools" | "datetime" | "decimal" | "fractions" | "uuid" | "pathlib" | "quarb"
235    )
236}
237
238/// A Python object held as a library value.
239pub struct PyForeign {
240    pub obj: Py<PyAny>,
241}
242
243impl PyForeign {
244    pub fn new(obj: Py<PyAny>) -> PyForeign {
245        PyForeign { obj }
246    }
247
248    fn glue<T>(&self, name: &str, f: impl Fn(Bound<'_, PyAny>) -> PyResult<T>) -> Result<T, String> {
249        Python::attach(|py| {
250            let m = module(py).map_err(|e| e.to_string())?;
251            let out = m
252                .getattr(name)
253                .and_then(|g| g.call1((self.obj.bind(py),)))
254                .map_err(|e| {
255                    // The message without the exception's class.
256                    e.value(py).to_string()
257                })?;
258            f(out).map_err(|e| e.to_string())
259        })
260    }
261}
262
263impl Foreign for PyForeign {
264    fn kind(&self) -> String {
265        self.glue("_kind", |o| o.extract::<String>())
266            .unwrap_or_else(|_| "a Python object".to_string())
267    }
268
269    fn display(&self) -> String {
270        Python::attach(|py| {
271            self.obj
272                .bind(py)
273                .str()
274                .map(|s| s.to_string_lossy().into_owned())
275                .unwrap_or_else(|e| e.to_string())
276        })
277    }
278
279    fn rows(&self) -> Result<Value, String> {
280        self.glue("_rows", |o| crate::py_to_value(&o))
281    }
282
283    fn columns(&self) -> Result<Value, String> {
284        self.glue("_columns", |o| crate::py_to_value(&o))
285    }
286
287    fn as_any(&self) -> &dyn std::any::Any {
288        self
289    }
290}
291
292/// A frame or an array out of the library as the engine's own
293/// table, a one-dimensional one as a list; `None` for an object of
294/// another kind (kept as a library value).
295pub(crate) fn as_table(obj: &Bound<'_, PyAny>) -> PyResult<Option<Value>> {
296    let py = obj.py();
297    let m = module(py)?;
298    let d = m.getattr("_table")?.call1((obj,))?;
299    if d.is_none() {
300        return Ok(None);
301    }
302    let d = d.cast::<PyDict>().map_err(|e| PyErr::from(e))?;
303    if let Some(list) = d.get_item("list")? {
304        return Ok(Some(crate::py_to_value(&list)?));
305    }
306    let header: Option<Vec<String>> = match d.get_item("header")? {
307        Some(h) if !h.is_none() => Some(h.extract()?),
308        _ => None,
309    };
310    let mut rows: Vec<Vec<Value>> = Vec::new();
311    if let Some(rs) = d.get_item("rows")? {
312        for r in rs.try_iter()? {
313            let mut row = Vec::new();
314            for c in r?.try_iter()? {
315                row.push(crate::py_to_value(&c?)?);
316            }
317            rows.push(row);
318        }
319    }
320    Ok(Some(Value::table(quarb::Table { header, rows })))
321}
quarb_pyvalue/foreign.rs

quarb_pyvalue/
foreign.rs
