feature: added convert strategy

shenek · shenek · commit a892bd43f982 · 2020-10-13T23:45:52.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -32,4 +32,4 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = { version = "~0.11.1", features = ["extension-module"] }
-streamson-lib = "~5.0.0"
+streamson-lib = "~6.0.0"
diff --git a/src/convert.rs b/src/convert.rs
@@ -0,0 +1,100 @@
+use super::{RustMatcher, StreamsonError};
+use pyo3::{prelude::*, types::PyBytes};
+use std::sync::{Arc, Mutex};
+use streamson_lib::{error, handler, path::Path, strategy};
+
+/// TODO
+#[pyclass]
+#[derive(Clone)]
+pub struct PythonConverter {
+    callable: PyObject,
+    use_path: bool,
+}
+
+#[pymethods]
+impl PythonConverter {
+    /// Create instance of PythonConverter
+    ///
+    /// # Arguments
+    /// * `callable` - python callable (3 arguments)
+    #[new]
+    pub fn new(callable: PyObject, use_path: bool) -> Self {
+        Self { callable, use_path }
+    }
+}
+
+impl handler::Handler for PythonConverter {
+    fn handle(
+        &mut self,
+        path: &Path,
+        matcher_idx: usize,
+        data: Option<&[u8]>,
+    ) -> Result<Option<Vec<u8>>, error::Handler> {
+        let gil = Python::acquire_gil();
+        let res = self
+            .callable
+            .call1(
+                gil.python(),
+                (
+                    if self.use_path {
+                        Some(path.to_string())
+                    } else {
+                        None
+                    },
+                    matcher_idx,
+                    data.unwrap().to_vec(),
+                ),
+            )
+            .map_err(|_| error::Handler::new("Failed to call handler function"))?;
+        let bytes = res.cast_as::<PyBytes>(gil.python()).unwrap();
+        Ok(FromPyObject::extract(bytes)
+            .map_err(|_| error::Handler::new("Function does not return bytes."))?)
+    }
+}
+
+/// Low level Python wrapper for Convert strategy
+#[pyclass]
+pub struct Convert {
+    convert: strategy::Convert,
+}
+
+#[pymethods]
+impl Convert {
+    /// Create a new instance of Convert
+    ///
+    /// # Arguments
+    /// * `export_path` - indicator whether path is required in further processing
+    #[new]
+    pub fn new() -> PyResult<Self> {
+        let convert = strategy::Convert::new();
+        Ok(Self { convert })
+    }
+
+    /// Adds matcher for Convert
+    ///
+    /// # Arguments
+    /// * `matcher` - matcher to be added (`Simple` or `Depth`)
+    /// * `handlers` - list of handlers to process
+    pub fn add_matcher(&mut self, matcher: &RustMatcher, mut handlers: Vec<PythonConverter>) {
+        let mut res: Vec<Arc<Mutex<dyn handler::Handler>>> = vec![];
+        handlers
+            .drain(..)
+            .for_each(|e| res.push(Arc::new(Mutex::new(e))));
+
+        self.convert
+            .add_matcher(Box::new(matcher.inner.clone()), res)
+    }
+
+    pub fn process(&mut self, input_data: &[u8]) -> PyResult<String> {
+        match self.convert.process(input_data) {
+            Err(err) => Err(StreamsonError::from(err).into()),
+            Ok(data_list) => {
+                let mut res = String::new();
+                for out_data in data_list {
+                    res += &String::from_utf8(out_data.to_vec())?;
+                }
+                Ok(res)
+            }
+        }
+    }
+}
diff --git a/src/filter.rs b/src/filter.rs
@@ -31,7 +31,7 @@ impl Filter {
     /// * `data` - input data to be processed
     ///
     /// # Returns
-    /// * `vector of tuples` - (path_or_none, data)
+    /// * `filtered output`
     pub fn process(&mut self, data: &[u8]) -> PyResult<String> {
         match self.filter.process(data) {
             Err(err) => Err(StreamsonError::from(err).into()),
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,7 +1,9 @@
+pub mod convert;
 pub mod extract;
 pub mod filter;
 pub mod trigger;
 
+pub use convert::{Convert, PythonConverter};
 pub use extract::Extract;
 pub use filter::Filter;
 pub use trigger::{PythonHandler, Trigger};
@@ -88,6 +90,8 @@ impl PyNumberProtocol for RustMatcher {
 /// This module is a python module implemented in Rust.
 #[pymodule]
 fn streamson(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<Convert>()?;
+    m.add_class::<PythonConverter>()?;
     m.add_class::<Extract>()?;
     m.add_class::<Filter>()?;
     m.add_class::<RustMatcher>()?;
diff --git a/src/trigger.rs b/src/trigger.rs
@@ -39,7 +39,7 @@ impl handler::Handler for PythonHandler {
         path: &Path,
         matcher_idx: usize,
         data: Option<&[u8]>,
-    ) -> Result<(), error::Handler> {
+    ) -> Result<Option<Vec<u8>>, error::Handler> {
         let gil = Python::acquire_gil();
         self.callable
             .call1(
@@ -55,7 +55,7 @@ impl handler::Handler for PythonHandler {
                 ),
             )
             .map_err(|_| error::Handler::new("Calling python failed"))?;
-        Ok(())
+        Ok(None)
     }
 }
 
diff --git a/streamson/__init__.py b/streamson/__init__.py
@@ -1,3 +1,4 @@
+from .convert import convert_async, convert_fd, convert_iter  # noqa
 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
diff --git a/streamson/convert.py b/streamson/convert.py
@@ -0,0 +1,75 @@
+import typing
+
+from streamson.streamson import Convert, PythonConverter
+
+from .matcher import Matcher
+
+
+def convert_iter(
+    input_gen: typing.Generator[bytes, None, None],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
+    matcher: Matcher,
+    require_path: bool = True,
+) -> typing.Generator[str, None, None]:
+    """Converts handlers on matched data from a file description
+    :param input_gen: input generator
+    :param handlers: list of handlers which should be called on match
+    :param matcher: used matcher
+    :param: require_path: is path required for handlers
+
+    :yields: converted data
+    """
+    convert = Convert()
+    convert.add_matcher(matcher.inner, [PythonConverter(handler, require_path) for handler in handlers])
+    for item in input_gen:
+        for output in convert.process(item):
+            yield output
+
+
+def convert_fd(
+    input_fd: typing.IO[bytes],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
+    matcher: Matcher,
+    buffer_size: int = 1024 * 1024,
+    require_path: bool = True,
+) -> typing.Generator[str, None, None]:
+    """Converts handlers on matched data from a file description
+    :param input_fd: input generator
+    :param handlers: list of handlers which should be called on match
+    :param matcher: used matcher
+    :param: buffer_size: how many bytes can be read from a file at once
+    :param: require_path: is path required for handlers
+
+    :yields: converted data
+    """
+    convert = Convert()
+    convert.add_matcher(matcher.inner, [PythonConverter(handler, require_path) for handler in handlers])
+
+    input_data = input_fd.read(buffer_size)
+
+    while input_data:
+        for item in convert.process(input_data):
+            yield item
+        input_data = input_fd.read(buffer_size)
+
+
+async def convert_async(
+    input_gen: typing.AsyncGenerator[bytes, None],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
+    matcher: Matcher,
+    require_path: bool = True,
+) -> typing.AsyncGenerator[str, None]:
+    """Convert handlers on matched data from async generator
+    :param: input_gen: input generator
+    :param handlers: list of handlers which should be called on match
+    :param: matcher: used matcher
+    :param: require_path: is path required for handlers
+
+    :yields: input data
+    """
+    convert = Convert()
+    convert.add_matcher(matcher.inner, [PythonConverter(handler, require_path) for handler in handlers])
+
+    async for input_data in input_gen:
+        for item in convert.process(input_data):
+            yield item
diff --git a/streamson/trigger.py b/streamson/trigger.py
@@ -7,7 +7,7 @@
 
 def trigger_iter(
     input_gen: typing.Generator[bytes, None, None],
-    handlers: typing.List[typing.Callable[[typing.Optional[str], int, bytes], None]],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
     matcher: Matcher,
     require_path: bool = True,
 ) -> typing.Generator[bytes, None, None]:
@@ -28,7 +28,7 @@ def trigger_iter(
 
 def trigger_fd(
     input_fd: typing.IO[bytes],
-    handlers: typing.List[typing.Callable[[typing.Optional[str], int, bytes], None]],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
     matcher: Matcher,
     buffer_size: int = 1024 * 1024,
     require_path: bool = True,
@@ -55,7 +55,7 @@ def trigger_fd(
 
 async def trigger_async(
     input_gen: typing.AsyncGenerator[bytes, None],
-    handlers: typing.List[typing.Callable[[typing.Optional[str], int, bytes], None]],
+    handlers: typing.List[typing.Callable[[typing.Optional[str], int, typing.Optional[bytes]], typing.Optional[bytes]]],
     matcher: Matcher,
     require_path: bool = True,
 ):
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -36,9 +36,17 @@ def io_reader() -> io.BytesIO:
 
 @pytest.fixture(scope="function")
 def handler():
-    output: typing.List[typing.Tuple[typing.Optional[str], int, bytes]] = []
+    output: typing.List[typing.Tuple[typing.Optional[str], int, typing.Optional[bytes]]] = []
 
     def store_to_output(path: typing.Optional[str], matcher_idx: int, data: bytes):
         output.append((path, matcher_idx, data))
 
     return store_to_output, output
+
+
+@pytest.fixture(scope="function")
+def converter():
+    def convert_to_stars(path: typing.Optional[str], matcher_idx: int, data: typing.Optional[bytes]):
+        return br'"***"'
+
+    return convert_to_stars
diff --git a/tests/test_convert.py b/tests/test_convert.py
@@ -0,0 +1,38 @@
+from enum import Enum, auto
+
+import pytest
+
+import streamson
+
+
+class Kind(Enum):
+    FD = auto()
+    ITER = auto()
+
+
+@pytest.mark.parametrize(
+    "kind,extract_path",
+    [
+        (Kind.FD, True),
+        (Kind.FD, False),
+        (Kind.ITER, True),
+        (Kind.ITER, False),
+    ],
+    ids=[
+        "fd-path",
+        "fd-nopath",
+        "iter-path",
+        "iter-nopath",
+    ],
+)
+def test_simple(io_reader, data, converter, kind, extract_path):
+
+    matcher = streamson.SimpleMatcher('{"users"}[1]')
+    output_data = ""
+    if kind == Kind.ITER:
+        for e in streamson.convert_iter((e for e in data), [converter], matcher, extract_path):
+            output_data += e
+    elif kind == Kind.FD:
+        for e in streamson.convert_fd(io_reader, [converter], matcher, 5, extract_path):
+            output_data += e
+    assert output_data == '{"users": ["john", "***", "bob"], "groups": ["admins", "users"]}'
diff --git a/tests/test_convert_async.py b/tests/test_convert_async.py
@@ -0,0 +1,23 @@
+import pytest
+
+import streamson
+
+
+@pytest.mark.parametrize(
+    "extract_path",
+    [True, False],
+    ids=[
+        "path",
+        "nopath",
+    ],
+)
+@pytest.mark.asyncio
+async def test_simple(make_async_gen, converter, extract_path):
+
+    matcher = streamson.SimpleMatcher('{"users"}[1]')
+
+    output_data = ""
+    async for e in streamson.convert_async(make_async_gen()(), [converter], matcher, extract_path):
+        output_data += e
+
+    assert output_data == '{"users": ["john", "***", "bob"]}'

Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ impl handler::Handler for PythonHandler {`
`39`	`39`	`path: &Path,`
`40`	`40`	`matcher_idx: usize,`
`41`	`41`	`data: Option<&[u8]>,`
`42`		`- ) -> Result<(), error::Handler> {`
	`42`	`+ ) -> Result<Option<Vec<u8>>, error::Handler> {`
`43`	`43`	`let gil = Python::acquire_gil();`
`44`	`44`	`self.callable`
`45`	`45`	`.call1(`
`@@ -55,7 +55,7 @@ impl handler::Handler for PythonHandler {`
`55`	`55`	`),`
`56`	`56`	`)`
`57`	`57`	`.map_err(\|_\| error::Handler::new("Calling python failed"))?;`
`58`		`- Ok(())`
	`58`	`+ Ok(None)`
`59`	`59`	`}`
`60`	`60`	`}`
`61`	`61`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+from .convert import convert_async, convert_fd, convert_iter # noqa`
`1`	`2`	`from .extract import extract_async, extract_fd, extract_iter # noqa`
`2`	`3`	`from .filter import filter_async, filter_fd, filter_iter # noqa`
`3`	`4`	`from .matcher import DepthMatcher, Matcher, SimpleMatcher # noqa`