fix: address regex filter review findings

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

plugins/rust/python-package/regex_filter/cpex_regex_filter/regex_filter_rust/__init__.pyi

@@ -1,9 +1,11 @@
 # This file is automatically generated by pyo3_stub_gen
 # ruff: noqa: E501, F401, F403, F405
 
+import builtins
 import typing
 __all__ = [
     "RegexFilterPluginCore",
+    "SearchReplacePluginRust",
 ]
 
 @typing.final
@@ -13,3 +15,9 @@ class RegexFilterPluginCore:
     def prompt_post_fetch(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
     def tool_pre_invoke(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
     def tool_post_invoke(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
+
+@typing.final
+class SearchReplacePluginRust:
+    def __new__(cls, config_dict: dict) -> SearchReplacePluginRust: ...
+    def apply_patterns(self, text: builtins.str) -> builtins.str: ...
+    def process_nested(self, data: typing.Any) -> tuple[builtins.bool, typing.Any]: ...

plugins/rust/python-package/regex_filter/src/bin/stub_gen.rs

@@ -10,13 +10,21 @@ use regex_filter_rust::stub_info;
 
 const EXTENSION_STUB_PATH: &str = "cpex_regex_filter/regex_filter_rust/__init__.pyi";
 const ORPHAN_EXTENSION_STUB_PATH: &str = "python/regex_filter_rust/__init__.pyi";
-const GENERATED_ALL_MARKER: &str = "\"SearchReplacePluginRust\",\n]";
-const CURATED_ALL_ENTRY: &str = "\"RegexFilterPluginCore\",\n    \"SearchReplacePluginRust\",\n]";
+const GENERATED_ALL_MARKER: &str = "__all__ = [\n";
 const PLUGIN_CORE_CLASS_MARKER: &str = "class RegexFilterPluginCore:";
-const PLUGIN_CORE_CLASS_DEF: &str = "\n\n@typing.final\nclass RegexFilterPluginCore:\n    def __new__(cls, config: typing.Any) -> RegexFilterPluginCore: ...\n    def prompt_pre_fetch(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def prompt_post_fetch(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def tool_pre_invoke(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def tool_post_invoke(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n";
+const ENGINE_CLASS_MARKER: &str = "class SearchReplacePluginRust:";
+const CURATED_ALL_BLOCK: &str =
+    "__all__ = [\n    \"RegexFilterPluginCore\",\n    \"SearchReplacePluginRust\",\n]\n";
+const PLUGIN_CORE_CLASS_DEF: &str = "\n\n@typing.final\nclass RegexFilterPluginCore:\n    def __new__(cls, config: dict) -> RegexFilterPluginCore: ...\n    def prompt_pre_fetch(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def prompt_post_fetch(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def tool_pre_invoke(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n    def tool_post_invoke(self, payload: typing.Any, context: typing.Any) -> typing.Any: ...\n";
 
 fn curate_extension_stub_content(content: &str) -> String {
-    let mut curated = content.replace(GENERATED_ALL_MARKER, CURATED_ALL_ENTRY);
+    let mut curated = content.to_string();
+    if let Some(all_start) = curated.find(GENERATED_ALL_MARKER)
+        && let Some(relative_end) = curated[all_start..].find("]\n")
+    {
+        let all_end = all_start + relative_end + 2;
+        curated.replace_range(all_start..all_end, CURATED_ALL_BLOCK);
+    }
     if !curated.contains(PLUGIN_CORE_CLASS_MARKER) {
         curated.push_str(PLUGIN_CORE_CLASS_DEF);
     }
@@ -25,10 +33,18 @@ fn curate_extension_stub_content(content: &str) -> String {
         curated.contains("\"RegexFilterPluginCore\""),
         "curated extension stub is missing RegexFilterPluginCore in __all__",
     );
+    assert!(
+        curated.contains("\"SearchReplacePluginRust\""),
+        "curated extension stub is missing SearchReplacePluginRust in __all__",
+    );
     assert!(
         curated.contains(PLUGIN_CORE_CLASS_MARKER),
         "curated extension stub is missing RegexFilterPluginCore class definition",
     );
+    assert!(
+        curated.contains(ENGINE_CLASS_MARKER),
+        "curated extension stub is missing SearchReplacePluginRust class definition",
+    );
 
     curated
 }
@@ -68,9 +84,10 @@ mod tests {
 
     #[test]
     fn test_curate_extension_stub_adds_required_exports_and_class() {
-        let generated = "# This file is automatically generated by pyo3_stub_gen\n# ruff: noqa: E501, F401, F403, F405\n\nimport typing\n\n__all__ = [\n    \"SearchReplacePluginRust\",\n]\n";
+        let generated = "# This file is automatically generated by pyo3_stub_gen\n# ruff: noqa: E501, F401, F403, F405\n\nimport builtins\nimport typing\n__all__ = [\n    \"SearchReplacePluginRust\",\n]\n\n@typing.final\nclass SearchReplacePluginRust:\n    def __new__(cls, config: typing.Any) -> SearchReplacePluginRust: ...\n";
         let curated = curate_extension_stub_content(generated);
         assert!(curated.contains("\"RegexFilterPluginCore\""));
+        assert!(curated.contains("\"SearchReplacePluginRust\""));
         assert!(curated.contains(PLUGIN_CORE_CLASS_MARKER));
     }
 }

plugins/rust/python-package/regex_filter/src/lib.rs

@@ -4,16 +4,25 @@
 // Regex Filter Plugin - Rust Implementation
 
 use std::borrow::Cow;
+use std::collections::HashSet;
 use std::sync::Once;
 
 use log::debug;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyModule, PyTuple};
 use pyo3_stub_gen::define_stub_info_gatherer;
+use pyo3_stub_gen::derive::*;
 use regex::{Regex, RegexSet};
 
 pub mod plugin;
 
+const MAX_NESTED_DEPTH: usize = 64;
+
+enum TraversalResult {
+    Unchanged(Py<PyAny>),
+    Modified(Py<PyAny>),
+}
+
 #[derive(Debug, Clone)]
 pub struct SearchReplace {
     pub search: String,
@@ -85,118 +94,185 @@ impl SearchReplaceConfig {
     }
 }
 
+#[gen_stub_pyclass]
 #[derive(Debug)]
 #[pyclass]
 pub struct SearchReplacePluginRust {
     pub config: SearchReplaceConfig,
 }
 
-#[pymethods]
-impl SearchReplacePluginRust {
-    #[new]
-    pub fn new(config_dict: &Bound<'_, PyDict>) -> PyResult<Self> {
-        let config = SearchReplaceConfig::from_py_dict(config_dict).map_err(|error| {
-            PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid config: {}", error))
-        })?;
-        Ok(Self { config })
+fn apply_patterns_impl<'a>(config: &'a SearchReplaceConfig, text: &'a str) -> Cow<'a, str> {
+    if let Some(ref pattern_set) = config.pattern_set
+        && !pattern_set.is_match(text)
+    {
+        return Cow::Borrowed(text);
     }
 
-    pub fn apply_patterns(&self, text: &str) -> String {
-        if let Some(ref pattern_set) = self.config.pattern_set
-            && !pattern_set.is_match(text)
-        {
-            return text.to_string();
+    let mut result = Cow::Borrowed(text);
+    let mut modified = false;
+
+    for pattern in &config.words {
+        if pattern.compiled.is_match(&result) {
+            let replaced = pattern.compiled.replace_all(&result, &pattern.replace);
+            if let Cow::Owned(new_text) = replaced {
+                result = Cow::Owned(new_text);
+                modified = true;
+            } else if modified {
+                result = Cow::Owned(replaced.into_owned());
+            }
+        }
+    }
+
+    result
+}
+
+fn process_nested_impl(
+    plugin: &SearchReplacePluginRust,
+    py: Python<'_>,
+    data: &Bound<'_, PyAny>,
+    depth: usize,
+    seen: &mut HashSet<usize>,
+) -> PyResult<TraversalResult> {
+    if depth >= MAX_NESTED_DEPTH {
+        return Err(pyo3::exceptions::PyValueError::new_err(format!(
+            "Maximum nested depth of {} exceeded",
+            MAX_NESTED_DEPTH
+        )));
+    }
+
+    if let Ok(text) = data.extract::<String>() {
+        let modified_text = apply_patterns_impl(&plugin.config, &text);
+        return match modified_text {
+            Cow::Borrowed(_) => Ok(TraversalResult::Unchanged(data.clone().unbind())),
+            Cow::Owned(value) => Ok(TraversalResult::Modified(
+                value.into_pyobject(py)?.into_any().unbind(),
+            )),
+        };
+    }
+
+    if let Ok(dict) = data.cast::<PyDict>() {
+        let identity = dict.as_ptr() as usize;
+        if !seen.insert(identity) {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "Cyclic containers are not supported",
+            ));
         }
 
-        let mut result = Cow::Borrowed(text);
-        let mut modified = false;
-
-        for pattern in &self.config.words {
-            if pattern.compiled.is_match(&result) {
-                let replaced = pattern.compiled.replace_all(&result, &pattern.replace);
-                if let Cow::Owned(new_text) = replaced {
-                    result = Cow::Owned(new_text);
-                    modified = true;
-                } else if modified {
-                    result = Cow::Owned(replaced.into_owned());
+        let mut any_modified = false;
+        let mut processed_items = Vec::with_capacity(dict.len());
+        for (key, value) in dict.iter() {
+            match process_nested_impl(plugin, py, &value, depth + 1, seen)? {
+                TraversalResult::Unchanged(new_value) => {
+                    processed_items.push((key.clone().unbind(), new_value));
+                }
+                TraversalResult::Modified(new_value) => {
+                    any_modified = true;
+                    processed_items.push((key.clone().unbind(), new_value));
                 }
             }
         }
+        seen.remove(&identity);
 
-        result.into_owned()
+        if !any_modified {
+            return Ok(TraversalResult::Unchanged(data.clone().unbind()));
+        }
+
+        let new_dict = PyDict::new(py);
+        for (key, value) in processed_items {
+            new_dict.set_item(key.bind(py), value.bind(py))?;
+        }
+        return Ok(TraversalResult::Modified(new_dict.into_any().unbind()));
     }
 
-    pub fn process_nested(
-        &self,
-        py: Python<'_>,
-        data: &Bound<'_, PyAny>,
-    ) -> PyResult<(bool, Py<PyAny>)> {
-        if let Ok(text) = data.extract::<String>() {
-            let modified_text = self.apply_patterns(&text);
-            if modified_text == text {
-                return Ok((false, data.clone().unbind()));
-            }
-            return Ok((true, modified_text.into_pyobject(py)?.into_any().unbind()));
+    if let Ok(list) = data.cast::<PyList>() {
+        let identity = list.as_ptr() as usize;
+        if !seen.insert(identity) {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "Cyclic containers are not supported",
+            ));
         }
 
-        if let Ok(dict) = data.cast::<PyDict>() {
-            let mut any_modified = false;
-            let mut processed_items = Vec::with_capacity(dict.len());
-            for (key, value) in dict.iter() {
-                let (item_modified, new_value) = self.process_nested(py, &value)?;
-                any_modified |= item_modified;
-                processed_items.push((key.clone().unbind(), new_value));
+        let mut any_modified = false;
+        let mut new_items = Vec::with_capacity(list.len());
+        for item in list.iter() {
+            match process_nested_impl(plugin, py, &item, depth + 1, seen)? {
+                TraversalResult::Unchanged(new_item) => new_items.push(new_item),
+                TraversalResult::Modified(new_item) => {
+                    any_modified = true;
+                    new_items.push(new_item);
+                }
             }
+        }
+        seen.remove(&identity);
 
-            if !any_modified {
-                return Ok((false, data.clone().unbind()));
-            }
+        if !any_modified {
+            return Ok(TraversalResult::Unchanged(data.clone().unbind()));
+        }
 
-            let new_dict = PyDict::new(py);
-            for (key, value) in processed_items {
-                new_dict.set_item(key.bind(py), value.bind(py))?;
-            }
-            return Ok((true, new_dict.into_any().unbind()));
+        let new_list = PyList::empty(py);
+        for item in new_items {
+            new_list.append(item.bind(py))?;
         }
+        return Ok(TraversalResult::Modified(new_list.into_any().unbind()));
+    }
 
-        if let Ok(list) = data.cast::<PyList>() {
-            let mut any_modified = false;
-            let mut new_items = Vec::with_capacity(list.len());
-            for item in list.iter() {
-                let (item_modified, new_item) = self.process_nested(py, &item)?;
-                any_modified |= item_modified;
-                new_items.push(new_item);
-            }
+    if let Ok(tuple) = data.cast::<PyTuple>() {
+        let identity = tuple.as_ptr() as usize;
+        if !seen.insert(identity) {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "Cyclic containers are not supported",
+            ));
+        }
 
-            if !any_modified {
-                return Ok((false, data.clone().unbind()));
+        let mut any_modified = false;
+        let mut new_items = Vec::with_capacity(tuple.len());
+        for item in tuple.iter() {
+            match process_nested_impl(plugin, py, &item, depth + 1, seen)? {
+                TraversalResult::Unchanged(new_item) => new_items.push(new_item),
+                TraversalResult::Modified(new_item) => {
+                    any_modified = true;
+                    new_items.push(new_item);
+                }
             }
+        }
+        seen.remove(&identity);
 
-            let new_list = PyList::empty(py);
-            for item in new_items {
-                new_list.append(item.bind(py))?;
-            }
-            return Ok((true, new_list.into_any().unbind()));
+        if !any_modified {
+            return Ok(TraversalResult::Unchanged(data.clone().unbind()));
         }
 
-        if let Ok(tuple) = data.cast::<PyTuple>() {
-            let mut any_modified = false;
-            let mut new_items = Vec::with_capacity(tuple.len());
-            for item in tuple.iter() {
-                let (item_modified, new_item) = self.process_nested(py, &item)?;
-                any_modified |= item_modified;
-                new_items.push(new_item);
-            }
+        let new_tuple = PyTuple::new(py, new_items.iter().map(|item| item.bind(py)))?;
+        return Ok(TraversalResult::Modified(new_tuple.into_any().unbind()));
+    }
 
-            if !any_modified {
-                return Ok((false, data.clone().unbind()));
-            }
+    Ok(TraversalResult::Unchanged(data.clone().unbind()))
+}
 
-            let new_tuple = PyTuple::new(py, new_items.iter().map(|item| item.bind(py)))?;
-            return Ok((true, new_tuple.into_any().unbind()));
-        }
+#[gen_stub_pymethods]
+#[pymethods]
+impl SearchReplacePluginRust {
+    #[new]
+    pub fn new(config_dict: &Bound<'_, PyDict>) -> PyResult<Self> {
+        let config = SearchReplaceConfig::from_py_dict(config_dict).map_err(|error| {
+            PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid config: {}", error))
+        })?;
+        Ok(Self { config })
+    }
 
-        Ok((false, data.clone().unbind()))
+    pub fn apply_patterns(&self, text: &str) -> String {
+        apply_patterns_impl(&self.config, text).into_owned()
+    }
+
+    pub fn process_nested(
+        &self,
+        py: Python<'_>,
+        data: &Bound<'_, PyAny>,
+    ) -> PyResult<(bool, Py<PyAny>)> {
+        let mut seen = HashSet::new();
+        Ok(match process_nested_impl(self, py, data, 0, &mut seen)? {
+            TraversalResult::Unchanged(value) => (false, value),
+            TraversalResult::Modified(value) => (true, value),
+        })
     }
 }