use crate::parser::raw_match::RawMatchData;
use pyo3::exceptions::{PyIndexError, PyTypeError};
use pyo3::prelude::*;
use pyo3::types::PyList;
use pyo3::IntoPyObjectExt;
#[pyclass]
pub struct Results {
raw_data: Vec<RawMatchData>,
cached_results: Option<PyObject>,
}
impl Results {
pub fn new(raw_data: Vec<RawMatchData>) -> Self {
Self {
raw_data,
cached_results: None,
}
}
fn convert_all(&mut self, py: Python) -> PyResult<PyObject> {
if let Some(ref cached) = self.cached_results {
return Ok(cached.clone_ref(py));
}
let mut py_results: Vec<PyObject> = Vec::with_capacity(self.raw_data.len());
for raw_data in &self.raw_data {
let parse_result = raw_data.to_parse_result(py)?;
py_results.push(parse_result.into_py_any(py)?);
}
let items: Vec<_> = py_results.iter().map(|obj| obj.bind(py)).collect();
let results_list = PyList::new(py, items)?;
let list_obj = results_list.into_py_any(py)?;
self.cached_results = Some(list_obj.clone_ref(py));
Ok(list_obj)
}
pub fn convert_item(&self, index: usize, py: Python) -> PyResult<PyObject> {
if index >= self.raw_data.len() {
return Err(PyIndexError::new_err("list index out of range"));
}
let raw_data = &self.raw_data[index];
let parse_result = raw_data.to_parse_result(py)?;
parse_result.into_py_any(py)
}
}
#[pymethods]
impl Results {
fn __len__(&self) -> usize {
self.raw_data.len()
}
fn __getitem__(&self, key: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
if let Ok(index) = key.extract::<usize>() {
self.convert_item(index, py)
} else if let Ok(index) = key.extract::<isize>() {
let len = self.raw_data.len();
let actual_index = if index < 0 {
match (len as i128).checked_add(index as i128) {
Some(sum) if sum >= 0 && (sum as usize) < len => sum as usize,
_ => {
return Err(PyIndexError::new_err("list index out of range"));
}
}
} else {
let u = index as usize;
if u >= len {
return Err(PyIndexError::new_err("list index out of range"));
}
u
};
self.convert_item(actual_index, py)
} else if key.is_instance_of::<pyo3::types::PySlice>() {
let mut results = Results::new(self.raw_data.clone());
let list = results.convert_all(py)?;
let list_bound = list.bind(py);
let slice_result = list_bound.get_item(key)?;
Ok(slice_result.into_py_any(py)?)
} else {
Err(PyTypeError::new_err(
"list indices must be integers or slices",
))
}
}
fn __iter__(slf: PyRef<'_, Self>) -> PyResult<ResultsIterator> {
Ok(ResultsIterator {
results: slf.into(),
cached_list: None,
index: 0,
})
}
#[allow(clippy::wrong_self_convention)] fn to_list(&mut self, py: Python) -> PyResult<PyObject> {
self.convert_all(py)
}
fn __repr__(&self) -> String {
format!("<Results {} matches>", self.raw_data.len())
}
fn __str__(&self) -> String {
self.__repr__()
}
}
#[pyclass]
pub struct ResultsIterator {
results: Py<Results>,
cached_list: Option<PyObject>, index: usize,
}
#[pymethods]
impl ResultsIterator {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
fn __next__(&mut self, py: Python) -> PyResult<Option<PyObject>> {
if self.cached_list.is_none() {
let results = self.results.bind(py);
let list = results.call_method0("to_list")?;
self.cached_list = Some(list.into_py_any(py)?);
}
let list_bound = self
.cached_list
.as_ref()
.expect("cached_list set immediately above when None")
.bind(py)
.downcast::<pyo3::types::PyList>()?;
let len = list_bound.len();
if self.index >= len {
return Ok(None);
}
let item = list_bound.get_item(self.index)?;
self.index += 1;
Ok(Some(item.into_py_any(py)?))
}
}