instant-labs
/
instant-distance
mirror of https://github.com/instant-labs/instant-distance.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

py: make Neighbor.value a Python object

This commit is contained in:
Dirkjan Ochtman 2021-05-21 17:14:54 +02:00
parent 1660ad1521
commit 0dd6a4ece8
1 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
instant-distance-py/src/lib.rs
View File

@ -5,10 +5,13 @@ use std::io::{BufReader, BufWriter};
use std::iter::FromIterator; use std::iter::FromIterator;
use instant_distance::Point; use instant_distance::Point;
use pyo3::conversion::IntoPy;
use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::proc_macro::{pyclass, pymethods, pymodule, pyproto}; use pyo3::proc_macro::{pyclass, pymethods, pymodule, pyproto};
use pyo3::types::{PyList, PyModule}; use pyo3::types::{PyList, PyModule};
use pyo3::{Py, PyAny, PyErr, PyIterProtocol, PyObjectProtocol, PyRef, PyRefMut, PyResult, Python}; use pyo3::{
Py, PyAny, PyErr, PyIterProtocol, PyObject, PyObjectProtocol, PyRef, PyRefMut, PyResult, Python,
};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_big_array::big_array; use serde_big_array::big_array;
@ -173,7 +176,7 @@ impl PyIterProtocol for Search {
item.map(|item| Neighbor { item.map(|item| Neighbor {
distance: item.distance, distance: item.distance,
pid: item.pid.into_inner(), pid: item.pid.into_inner(),
value: None, value: py.None(),
}) })
} }
HnswType::Map(map) => { HnswType::Map(map) => {
@ -182,7 +185,7 @@ impl PyIterProtocol for Search {
item.map(|item| Neighbor { item.map(|item| Neighbor {
distance: item.distance, distance: item.distance,
pid: item.pid.into_inner(), pid: item.pid.into_inner(),
value: Some(item.value.to_owned()), value: item.value.into_py(py),
}) })
} }
}; };
@ -319,18 +322,20 @@ struct Neighbor {
pid: u32, pid: u32,
/// Value for the neighboring point (only set for `HnswMap` results) /// Value for the neighboring point (only set for `HnswMap` results)
#[pyo3(get)] #[pyo3(get)]
value: Option<String>, value: PyObject,
} }
#[pyproto] #[pyproto]
impl PyObjectProtocol for Neighbor { impl PyObjectProtocol for Neighbor {
fn __repr__(&self) -> PyResult<String> { fn __repr__(&self) -> PyResult<String> {
match &self.value { match Python::with_gil(|py| self.value.is_none(py)) {
Some(s) => Ok(format!( false => Ok(format!(
"instant_distance.Neighbor(distance={}, pid={}, value={})", "instant_distance.Neighbor(distance={}, pid={}, value={})",
self.distance, self.pid, s, self.distance,
self.pid,
Python::with_gil(|py| self.value.as_ref(py).repr().map(|s| s.to_string()))?,
)), )),
None => Ok(format!( true => Ok(format!(
"instant_distance.Item(distance={}, pid={})", "instant_distance.Item(distance={}, pid={})",
self.distance, self.pid, self.distance, self.pid,
)), )),