minor refactoring

pyinfra/callback.py

@@ -35,7 +35,7 @@ class Callback:
         """
         TODO: Since data and metadata are passed as singletons, there is no buffering and hence no batching happening
          within the pipeline. However, the queue acknowledgment logic needs to be changed in order to facilitate
-         passing   non-singletons, to only ack a message, once a response is pulled from the output queue of the
+         passing non-singletons, to only ack a message, once a response is pulled from the output queue of the
          pipeline. Probably the pipeline return value needs to contains the queue message frame (or so), in order for
          the queue manager to tell which message to ack.
 

pyinfra/server/buffering/stream.py

@@ -1,7 +1,7 @@
 from itertools import chain, takewhile
 from typing import Iterable
 
-from funcy import first, repeatedly, flatten
+from funcy import first, repeatedly, mapcat
 
 from pyinfra.server.buffering.bufferize import bufferize
 from pyinfra.server.dispatcher.dispatcher import Nothing, is_not_nothing
@@ -16,7 +16,7 @@ class FlatStreamBuffer:
 
     def __call__(self, items):
         items = chain(items, [Nothing])
-        yield from flatten(map(self.stream_buffer, items))
+        yield from mapcat(self.stream_buffer, items)
 
 
 class StreamBuffer:

test/unit_tests/server/pipeline_test.py

@@ -1,5 +1,5 @@
 import pytest
-from funcy import rcompose, compose, project, second, merge, lpluck
+from funcy import rcompose, compose, project, merge
 
 from pyinfra.server.buffering.stream import FlatStreamBuffer
 from pyinfra.server.client_pipeline import ClientPipeline
@@ -40,9 +40,9 @@ def test_pipeline(
         for storage_mdt, queue_mdt in zip(metadata, queue_message_metadata)
     ]
 
-    output = compose(llift(unpack), client_pipeline)(input_data_items, metadata)
-    assert n_items == 0 or len(output) > 0
-    assert output == targets
+    outputs = compose(llift(unpack), client_pipeline)(input_data_items, metadata)
+    assert n_items == 0 or len(outputs) > 0
+    assert outputs == targets
 
 
 @pytest.mark.parametrize("item_type", ["string"])