From 16ec6dcce8a6dfa61bab24d560ec7789649cd3ce Mon Sep 17 00:00:00 2001
From: Dimitri Merejkowsky <dimitri.merejkowsky@tanker.io>
Date: Thu, 27 Jun 2019 13:55:02 +0200
Subject: [PATCH] proofreading

---
 sources/stegano/stegano.py      | 18 ++++++++++--------
 sources/stegano/test_stegano.py |  6 ++++++
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/sources/stegano/stegano.py b/sources/stegano/stegano.py
index e3539a1..baa4274 100644
--- a/sources/stegano/stegano.py
+++ b/sources/stegano/stegano.py
@@ -1,5 +1,4 @@
 import PIL.Image
-import itertools
 
 import sys
 
@@ -22,9 +21,12 @@ def parse_bitstream(stream):
 
 def by_chunk(iterable, size, fillvalue=0):
     "Collect data into fixed-length chunks or blocks"
-    args = [iter(iterable)] * size
-    return itertools.zip_longest(*args, fillvalue=fillvalue)
-
+    chunk = []
+    for x in iterable:
+        chunk.append(x)
+        if len(chunk) >= size:
+            yield chunk
+            chunk = []
 
 def set_bit(old_byte, new_bit):
     b = list(bin(old_byte))
@@ -33,8 +35,8 @@ def set_bit(old_byte, new_bit):
 
 
 def main_encrypt():
-    base_name = sys.argv[2]
-    message = (sys.argv[3] + "\n").encode()
+    message = (sys.argv[2] + "\n").encode()
+    base_name = sys.argv[3]
     carrier_name = base_name.replace(".png", ".carrier.png")
     base_image = PIL.Image.open(base_name)
     width, height = base_image.size
@@ -42,7 +44,7 @@ def main_encrypt():
     bitstream = get_bitstream(message)
     for row in range(height):
         for col in range(width):
-            r, g, b = base_image.getpixel((col, row))
+            r, g, b, = base_image.getpixel((col, row))
             value = None
             try:
                 value = next(bitstream)
@@ -60,7 +62,7 @@ def yield_bits(carrier_image):
     width, height = carrier_image.size
     for row in range(height):
         for col in range(width):
-            r, g, b = carrier_image.getpixel((col, row))
+            r, g, b, = carrier_image.getpixel((col, row))
             last_bit = int(bin(r)[-1])
             yield last_bit
 
diff --git a/sources/stegano/test_stegano.py b/sources/stegano/test_stegano.py
index a23bf44..7ea3ad3 100644
--- a/sources/stegano/test_stegano.py
+++ b/sources/stegano/test_stegano.py
@@ -1,6 +1,12 @@
 import stegano
 
 
+def test_by_chunks():
+    x = [1, 2, 3, 4, 5, 6]
+    chunks = stegano.by_chunk(x, 3)
+    assert list(chunks) == [[1, 2, 3], [4, 5, 6]]
+
+
 def test_message_conversion():
     message = b"I love you\ngarbage"
     encoded = list(stegano.get_bitstream(message))