1use crate::parser::raw_match::RawMatchData;
2use pyo3::exceptions::{PyIndexError, PyTypeError};
3use pyo3::prelude::*;
4use pyo3::types::PyList;
5use pyo3::IntoPyObjectExt;
6
7#[pyclass]
11pub struct Results {
12 raw_data: Vec<RawMatchData>,
13 cached_results: Option<Py<PyAny>>,
15}
16
17impl Results {
18 pub fn new(raw_data: Vec<RawMatchData>) -> Self {
19 Self {
20 raw_data,
21 cached_results: None,
22 }
23 }
24
25 fn convert_all(&mut self, py: Python) -> PyResult<Py<PyAny>> {
27 if let Some(ref cached) = self.cached_results {
28 return Ok(cached.clone_ref(py));
29 }
30
31 let mut py_results: Vec<Py<PyAny>> = Vec::with_capacity(self.raw_data.len());
32 for raw_data in &self.raw_data {
33 let parse_result = raw_data.to_parse_result(py)?;
34 py_results.push(parse_result.into_py_any(py)?);
35 }
36
37 let items: Vec<_> = py_results.iter().map(|obj| obj.bind(py)).collect();
38 let results_list = PyList::new(py, items)?;
39 let list_obj = results_list.into_py_any(py)?;
40
41 self.cached_results = Some(list_obj.clone_ref(py));
43 Ok(list_obj)
44 }
45
46 pub fn convert_item(&self, index: usize, py: Python) -> PyResult<Py<PyAny>> {
48 if index >= self.raw_data.len() {
49 return Err(PyIndexError::new_err("list index out of range"));
50 }
51
52 let raw_data = &self.raw_data[index];
53 let parse_result = raw_data.to_parse_result(py)?;
54 parse_result.into_py_any(py)
55 }
56}
57
58#[pymethods]
59impl Results {
60 fn __len__(&self) -> usize {
62 self.raw_data.len()
63 }
64
65 fn __getitem__(&self, key: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
67 if let Ok(index) = key.extract::<usize>() {
69 self.convert_item(index, py)
71 } else if let Ok(index) = key.extract::<isize>() {
72 let len = self.raw_data.len();
73 let actual_index = if index < 0 {
74 match (len as i128).checked_add(index as i128) {
75 Some(sum) if sum >= 0 && (sum as usize) < len => sum as usize,
76 _ => {
77 return Err(PyIndexError::new_err("list index out of range"));
78 }
79 }
80 } else {
81 let u = index as usize;
82 if u >= len {
83 return Err(PyIndexError::new_err("list index out of range"));
84 }
85 u
86 };
87 self.convert_item(actual_index, py)
88 } else if let Ok(slice) = key.cast::<pyo3::types::PySlice>() {
89 let len = self.raw_data.len() as isize;
90 let indices = slice.indices(len)?;
91 let mut items: Vec<Py<PyAny>> = Vec::new();
92 let mut i = indices.start;
93 while (indices.step > 0 && i < indices.stop) || (indices.step < 0 && i > indices.stop) {
94 if i >= 0 && (i as usize) < self.raw_data.len() {
95 items.push(self.convert_item(i as usize, py)?);
96 }
97 i += indices.step;
98 }
99 let py_items: Vec<_> = items.iter().map(|obj| obj.bind(py)).collect();
100 Ok(PyList::new(py, py_items)?.into_py_any(py)?)
101 } else {
102 Err(PyTypeError::new_err(
103 "list indices must be integers or slices",
104 ))
105 }
106 }
107
108 fn __iter__(slf: PyRef<'_, Self>) -> PyResult<ResultsIterator> {
110 Ok(ResultsIterator {
111 results: slf.into(),
112 cached_list: None,
113 index: 0,
114 })
115 }
116
117 #[allow(clippy::wrong_self_convention)] fn to_list(&mut self, py: Python) -> PyResult<Py<PyAny>> {
121 self.convert_all(py)
122 }
123
124 fn __repr__(&self) -> String {
126 format!("<Results {} matches>", self.raw_data.len())
127 }
128
129 fn __str__(&self) -> String {
130 self.__repr__()
131 }
132}
133
134#[pyclass]
137pub struct ResultsIterator {
138 results: Py<Results>,
139 cached_list: Option<Py<PyAny>>, index: usize,
141}
142
143#[pymethods]
144impl ResultsIterator {
145 fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
146 slf
147 }
148
149 fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyAny>>> {
150 if self.cached_list.is_none() {
152 let results = self.results.bind(py);
153 let list = results.call_method0("to_list")?;
155 self.cached_list = Some(list.into_py_any(py)?);
156 }
157
158 let list_bound = self
160 .cached_list
161 .as_ref()
162 .expect("cached_list set immediately above when None")
163 .bind(py)
164 .cast::<pyo3::types::PyList>()?;
165 let len = list_bound.len();
166
167 if self.index >= len {
168 return Ok(None);
169 }
170
171 let item = list_bound.get_item(self.index)?;
172 self.index += 1;
173 Ok(Some(item.into_py_any(py)?))
174 }
175}