From db5de637555a9686514558a0364c3700a9f91fc2 Mon Sep 17 00:00:00 2001
From: jesko <rattle@nullteilerfrei.de>
Date: Tue, 21 Nov 2023 12:24:26 +0100
Subject: [PATCH] makes stego easier to use outside frames

---
 refinery/units/formats/stego.py  | 13 +++++++++++--
 test/units/formats/test_stego.py |  5 ++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/refinery/units/formats/stego.py b/refinery/units/formats/stego.py
index 6806261b9a..c7e5febd10 100644
--- a/refinery/units/formats/stego.py
+++ b/refinery/units/formats/stego.py
@@ -16,11 +16,12 @@ class PIXEL_PART(IntEnum):
 class stego(Unit):
     """
     Decodes the RGBA (red/green/blue/alpha) values of the pixels of a given image file and outputs
-    these values as bytes. Each row of the image is transformed and output as an individual chunk.
+    these values as bytes. By default, the pixels are converted left to right, top to bottom.
     """
     def __init__(
         self,
         transpose: Arg.Switch('-t', help='Return the columns of the image rather than the rows.'),
+        split: Arg.Switch('-m', help='Emit the individual rows or columns as separate outputs.') = False,
         parts: Arg('parts', nargs='?', type=str, help=(
             'A string containing any ordering of the letters R, G, B, and A (case-insensitive). '
             'These pixel components will be extracted from every pixel in the given order. The '
@@ -29,6 +30,7 @@ def __init__(
     ):
         super().__init__(
             transpose=transpose,
+            split=split,
             parts=tuple(Arg.AsOption(p, PIXEL_PART) for p in parts)
         )
 
@@ -38,12 +40,14 @@ def _image():
         return Image
 
     def process(self, data):
+        split = self.args.split
         parts = self.args.parts
         image = self._image.open(MemoryFile(data))
         if self.args.transpose:
             image = image.transpose(self._image.Transpose.ROTATE_90)
         width, height = image.size
         chunk_size = len(parts)
+        output = MemoryFile()
         buffer = bytearray(chunk_size * width)
         for y in range(height):
             offset = 0
@@ -52,4 +56,9 @@ def process(self, data):
                 next_offset = offset + chunk_size
                 buffer[offset:next_offset] = (pixel[p] for p in parts)
                 offset = next_offset
-            yield buffer
+            if split:
+                yield buffer
+            else:
+                output.write(buffer)
+        if not split:
+            yield output.getvalue()
diff --git a/test/units/formats/test_stego.py b/test/units/formats/test_stego.py
index 79892e2cb4..988ee57f6b 100644
--- a/test/units/formats/test_stego.py
+++ b/test/units/formats/test_stego.py
@@ -32,5 +32,8 @@ def test_simple(self):
             'E002000000DCECECB4E2E2F5222299A7A7D9D9D9020202EAFDFDB4DFDFD8FEFE8A567CCFC3DE9AB90000000049454E'
             '44AE426082'
         )
-        stego = self.load('RGB')
+        stego = self.load('RGB', split=True)
         self.assertListEqual(expected_rows, list(image | stego))
+
+        stego = self.load('RGB', split=False)
+        self.assertEqual(b''.join(expected_rows), next(image | stego))