feature: add filter strategy

Author: shenek

src/filter.rs

@@ -0,0 +1,41 @@
+use super::{RustMatcher, StreamsonError};
+use pyo3::prelude::*;
+use streamson_lib::strategy;
+
+/// Low level Python wrapper for Filter strategy
+#[pyclass]
+pub struct Filter {
+    filter: strategy::Filter,
+}
+
+#[pymethods]
+impl Filter {
+    /// Create a new instance of Filter
+    #[new]
+    pub fn new() -> PyResult<Self> {
+        let filter = strategy::Filter::new();
+        Ok(Self { filter })
+    }
+
+    /// Adds matcher for Filter
+    ///
+    /// # Arguments
+    /// * `matcher` - matcher to be added (`Simple` or `Depth`)
+    pub fn add_matcher(&mut self, matcher: &RustMatcher) {
+        self.filter.add_matcher(Box::new(matcher.inner.clone()));
+    }
+
+    /// Process data for Filter strategy
+    ///
+    /// # Arguments
+    /// * `data` - input data to be processed
+    ///
+    /// # Returns
+    /// * `vector of tuples` - (path_or_none, data)
+    pub fn process(&mut self, data: &[u8]) -> PyResult<String> {
+        match self.filter.process(data) {
+            Err(err) => Err(StreamsonError::from(err).into()),
+            Ok(data) => Ok(String::from_utf8(data)?),
+        }
+    }
+}

src/lib.rs

@@ -1,7 +1,9 @@
 pub mod extract;
+pub mod filter;
 pub mod trigger;
 
 pub use extract::Extract;
+pub use filter::Filter;
 pub use trigger::{PythonHandler, Trigger};
 
 use pyo3::{class::PyNumberProtocol, create_exception, exceptions, prelude::*};
@@ -87,6 +89,7 @@ impl PyNumberProtocol for RustMatcher {
 #[pymodule]
 fn streamson(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<Extract>()?;
+    m.add_class::<Filter>()?;
     m.add_class::<RustMatcher>()?;
     m.add_class::<Trigger>()?;
     m.add_class::<PythonHandler>()?;

streamson/__init__.py

@@ -1,3 +1,4 @@
 from .extract import extract_async, extract_fd, extract_iter  # noqa
+from .filter import filter_async, filter_fd, filter_iter  # noqa
 from .matcher import DepthMatcher, Matcher, SimpleMatcher  # noqa
 from .trigger import trigger_async, trigger_fd, trigger_iter  # noqa

streamson/filter.py

@@ -0,0 +1,60 @@
+import typing
+
+from streamson.streamson import Filter
+
+from .matcher import Matcher
+
+
+def filter_iter(
+    input_gen: typing.Generator[bytes, None, None],
+    matcher: Matcher,
+) -> typing.Generator[typing.Tuple[str, typing.Any], None, None]:
+    """Filters json parts from generator specified by given matcher
+    :param: input_gen: input generator
+    :param: matcher: used matcher
+
+    :yields: filtered data
+    """
+    filter_strategy = Filter()
+    filter_strategy.add_matcher(matcher.inner)
+    for item in input_gen:
+        yield filter_strategy.process(item)
+
+
+def filter_fd(
+    input_fd: typing.IO[bytes],
+    matcher: Matcher,
+    buffer_size: int = 1024 * 1024,
+) -> typing.Generator[typing.Tuple[str, typing.Any], None, None]:
+    """Filters json parts from input file specified by given matcher
+    :param: input_fd: input fd
+    :param: matcher: used matcher
+    :param: buffer_size: how many bytes can be read from a file at once
+
+    :yields: filtered data
+    """
+    filter_strategy = Filter()
+    filter_strategy.add_matcher(matcher.inner)
+
+    input_data = input_fd.read(buffer_size)
+
+    while input_data:
+        yield filter_strategy.process(input_data)
+        input_data = input_fd.read(buffer_size)
+
+
+async def filter_async(
+    input_gen: typing.AsyncGenerator[bytes, None],
+    matcher: Matcher,
+):
+    """Filters json parts from given async generator specified by given matcher
+    :param: input_gen: input generator
+    :param: matcher: used matcher
+
+    :yields: filtered data
+    """
+    filter_strategy = Filter()
+    filter_strategy.add_matcher(matcher.inner)
+
+    async for input_data in input_gen:
+        yield filter_strategy.process(input_data)

tests/test_filter.py

@@ -0,0 +1,31 @@
+from enum import Enum, auto
+
+import pytest
+
+import streamson
+
+
+class Kind(Enum):
+    FD = auto()
+    ITER = auto()
+
+
+@pytest.mark.parametrize(
+    "kind",
+    [
+        (Kind.FD),
+        (Kind.ITER),
+    ],
+)
+def test_simple(io_reader, data, kind):
+    matcher = streamson.SimpleMatcher('{"users"}[]')
+
+    output_data = ""
+    if kind == Kind.ITER:
+        for e in streamson.filter_iter((e for e in data), matcher):
+            output_data += e
+    elif kind == Kind.FD:
+        for e in streamson.filter_fd(io_reader, matcher, 5):
+            output_data += e
+
+    assert output_data == '{"users": [], "groups": ["admins", "users"]}'

tests/test_filter_async.py

@@ -0,0 +1,15 @@
+import pytest
+
+import streamson
+
+
+@pytest.mark.asyncio
+async def test_simple(make_async_gen, handler):
+    matcher = streamson.SimpleMatcher('{"users"}[]')
+    async_out = streamson.filter_async(make_async_gen()(), matcher)
+
+    output_data = ""
+    async for rec in async_out:
+        output_data += rec
+
+    assert output_data == '{"users": []}'