return borrows where possible, use real iterators

Author: circuitsacul

regex_rs.pyi

@@ -1,3 +1,29 @@
+class Captures:
+    def get(self, i: int) -> Match | None: ...
+    def name(self, name: str) -> Match | None: ...
+    def expand(self, replacement: str, dst: str) -> str: ...
+    def __getitem__(self, i: int) -> Match: ...
+    def __len__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+
+class CapturesIter:
+    def __iter__(slf) -> CapturesIter: ...
+    def __next__(self) -> Captures: ...
+    def __repr__(self) -> str: ...
+
+class Match:
+    matched_text: str
+    start: int
+    end: int
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...
+
+class Matches:
+    def __iter__(self) -> Matches: ...
+    def __next__(self) -> Match: ...
+    def __repr__(self) -> str: ...
+
 class Regex:
     def __init__(
         self,
@@ -16,26 +42,15 @@ class Regex:
     ) -> None: ...
     def is_match(self, text: str, start: int | None = None) -> bool: ...
     def find(self, text: str, start: int | None) -> Match | None: ...
-    def find_iter(self, text: str) -> list[Match]: ...
+    def find_iter(self, text: str) -> Matches: ...
     def captures(self, text: str) -> Captures | None: ...
-    def split(self, text: str, limit: int | None = None) -> list[str]: ...
+    def captures_iter(self, text: str) -> CapturesIter: ...
+    def split(self, text: str, limit: int | None = None) -> Split: ...
     def replace(self, text: str, rep: str, limit: int | None = None) -> str: ...
-    def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
-
-
-class Captures:
-    def get(self, i: int) -> Match | None: ...
-    def name(self, name: str) -> Match | None: ...
-    def expand(self, replacement: str, dst: str) -> str: ...
-    def __getitem__(self, i: int) -> Match | None: ...
-    def __len__(self) -> int: ...
     def __repr__(self) -> str: ...
 
-
-class Match:
-    matched_text: str
-    start: int
-    end: int
-    def __str__(self) -> str: ...
+class Split:
+    def __iter__(slf) -> Split: ...
+    def __next__(self) -> str: ...
     def __repr__(self) -> str: ...

src/captures.rs

@@ -1,8 +1,7 @@
 use std::sync::Arc;
 
 use ouroboros::self_referencing;
-use pyo3::prelude::*;
-use regex::Captures as ReCaptures;
+use pyo3::{exceptions::PyIndexError, prelude::*};
 
 use crate::match_struct::Match;
 
@@ -13,7 +12,7 @@ pub struct Captures {
 
     #[borrows(text)]
     #[covariant]
-    pub captures: ReCaptures<'this>,
+    pub captures: regex::Captures<'this>,
 }
 
 #[pymethods]
@@ -31,16 +30,45 @@ impl Captures {
         dst
     }
 
-    pub fn __getitem__(&self, i: usize) -> Option<Match> {
-        self.get(i)
+    pub fn __getitem__(&self, i: usize) -> PyResult<Match> {
+        self.get(i).ok_or(PyIndexError::new_err(i))
     }
 
     pub fn __len__(&self) -> usize {
         self.borrow_captures().len()
     }
 
     pub fn __repr__(&self) -> String {
-        let dep = self.borrow_captures();
-        format!("{dep:#?}")
+        format!("{:#?}", self.borrow_captures())
+    }
+}
+
+#[pyclass]
+#[self_referencing(pub_extras)]
+pub struct CapturesIter {
+    pub text: Arc<String>,
+    pub re: Arc<regex::Regex>,
+
+    #[borrows(text, re)]
+    #[not_covariant]
+    pub capture_matches: regex::CaptureMatches<'this, 'this>,
+}
+
+#[pymethods]
+impl CapturesIter {
+    pub fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    pub fn __next__(&mut self) -> Option<Captures> {
+        let text = self.borrow_text().clone();
+        self.with_capture_matches_mut(|iter| {
+            iter.next()
+                .map(|caps| Captures::new(text, |text| caps.adopt(text)))
+        })
+    }
+
+    pub fn __repr__(&self) -> String {
+        self.with_capture_matches(|caps| format!("{caps:#?}"))
     }
 }

src/lib.rs

@@ -2,14 +2,18 @@ mod captures;
 mod error;
 mod match_struct;
 mod regex;
+mod split;
 
 use pyo3::prelude::*;
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn regex_rs(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_class::<regex::Regex>()?;
-    m.add_class::<match_struct::Match>()?;
     m.add_class::<captures::Captures>()?;
+    m.add_class::<captures::CapturesIter>()?;
+    m.add_class::<match_struct::Match>()?;
+    m.add_class::<match_struct::Matches>()?;
+    m.add_class::<regex::Regex>()?;
+    m.add_class::<split::Split>()?;
     Ok(())
 }

src/match_struct.rs

@@ -1,3 +1,6 @@
+use std::sync::Arc;
+
+use ouroboros::self_referencing;
 use pyo3::prelude::*;
 
 #[pyclass]
@@ -31,3 +34,29 @@ impl Match {
         format!("{self:#?}")
     }
 }
+
+#[pyclass]
+#[self_referencing(pub_extras)]
+pub struct Matches {
+    text: String,
+    re: Arc<regex::Regex>,
+
+    #[borrows(text, re)]
+    #[not_covariant]
+    matches: regex::Matches<'this, 'this>,
+}
+
+#[pymethods]
+impl Matches {
+    pub fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    pub fn __next__(&mut self) -> Option<Match> {
+        self.with_matches_mut(|iter| iter.next().map(|m| m.into()))
+    }
+
+    pub fn __repr__(&self) -> String {
+        self.with_matches(|matches| format!("{matches:#?}"))
+    }
+}

src/regex.rs

@@ -2,7 +2,12 @@ use std::sync::Arc;
 
 use pyo3::prelude::*;
 
-use crate::{captures::Captures, error::RegexResult, match_struct::Match};
+use crate::{
+    captures::{Captures, CapturesIter},
+    error::RegexResult,
+    match_struct::{Match, Matches},
+    split::Split,
+};
 
 #[pyclass]
 #[derive(Debug)]
@@ -96,28 +101,22 @@ impl Regex {
         mat.map(|m| m.into())
     }
 
-    pub fn find_iter(&self, text: &str) -> Vec<Match> {
-        self.0.find_iter(text).map(|m| m.into()).collect()
+    pub fn find_iter(&self, text: String) -> Matches {
+        Matches::new(text, self.0.clone(), |text, re| re.find_iter(text))
     }
 
     pub fn captures(&self, text: String) -> Option<Captures> {
         Captures::try_new(Arc::new(text), |text| self.0.captures(text).ok_or(())).ok()
     }
 
-    pub fn captures_iter(&self, text: String) -> Vec<Captures> {
+    pub fn captures_iter(&self, text: String) -> CapturesIter {
         let text = Arc::new(text);
-        self.0.captures_iter(&text)
-        .map(|caps| Captures::new(text.clone(), |text| caps.adopt(text)))
-        .collect()
+        CapturesIter::new(text, self.0.clone(), |text, re| re.captures_iter(text))
     }
 
     #[pyo3(signature = (text, limit=None))]
-    pub fn split(&self, text: &str, limit: Option<usize>) -> Vec<String> {
-        if let Some(limit) = limit {
-            self.0.splitn(text, limit).map(|v| v.to_owned()).collect()
-        } else {
-            self.0.split(text).map(|v| v.to_owned()).collect()
-        }
+    pub fn split(&self, text: String, limit: Option<usize>) -> Split {
+        Split::new(text, self.0.clone(), limit, |text, re| re.split(text))
     }
 
     #[pyo3(signature = (text, rep, limit=None))]

src/split.rs

@@ -0,0 +1,47 @@
+use std::sync::Arc;
+
+use ouroboros::self_referencing;
+use pyo3::prelude::*;
+
+#[pyclass]
+#[self_referencing(pub_extras)]
+#[derive(Debug)]
+pub struct Split {
+    pub text: String,
+    pub re: Arc<regex::Regex>,
+    pub limit: Option<usize>,
+
+    #[borrows(text, re)]
+    #[not_covariant]
+    pub split: regex::Split<'this, 'this>,
+}
+
+#[pymethods]
+impl Split {
+    pub fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    pub fn __next__(&mut self) -> Option<&str> {
+        let ret = self.with_limit_mut(|limit| {
+            if let Some(limit) = limit {
+                if *limit == 0 {
+                    return false;
+                }
+
+                *limit -= 1;
+            }
+
+            true
+        });
+        if !ret {
+            return None;
+        }
+
+        self.with_split_mut(|split| split.next())
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!("{self:#?}")
+    }
+}