Refined arguments to functions. Made sure system argument 0 wasn't i…

README.md

@@ -33,3 +33,24 @@ https://www.sultanik.com/<br />
 
 Initial Python3 port by [zenarcher007](https://github.com/zenarcher007), along with implementation of the `-v` and `-s`
 options.
+
+
+-----------  
+Fork by zenarcher007:  
+Latest commit:  
+• Changed --no-progress to -p --progress in order to eliminate confusion in the code and
+when being used as a module.  
+• You can now use Encoder.encode(infile, outfile=(path), square=(bool), width=(number), height=(number), progress=(bool), verbose=(bool))
+when importing bin2png as a module in a script.  
+
+
+Implemented decode function  
+To be used as a module in a script, the encode and decode functions  
+are split into a separate file. Here is an example of usage:  
+```
+from bin2png import Encoder
+f = Encoder.encode("path/test")    # <<(returns a png image)
+print(f)   # <<<  <PIL.Image.Image image mode=RGB size=72x1987 at 0x10752DD90>
+
+g = Encoder.decode(f)    # <<(returns a bytes object)
+```

bin2png.py

@@ -68,18 +68,18 @@ def read(self, n):
             return map(ord, self.file.read(n))
 
 
-def choose_file_dimensions(infile, input_dimensions=None):
+def choose_file_dimensions(infile, input_dimensions=None, square=False, verbose=False):
     if input_dimensions is not None and len(input_dimensions) >= 2 and input_dimensions[0] is not None \
             and input_dimensions[1] is not None:
         # the dimensions were already fully specified
         return input_dimensions
-    infile = FileReader.new(infile)
+    #infile = FileReader.new(infile)
     num_bytes = len(infile)
     num_pixels = int(math.ceil(float(num_bytes) / 3.0))
     sqrt = math.sqrt(num_pixels)
     sqrt_max = int(math.ceil(sqrt))
 
-    if args.square is True:
+    if square is True:
         return sqrt_max, sqrt_max
 
     if input_dimensions is not None and len(input_dimensions) >= 1:
@@ -111,17 +111,15 @@ def choose_file_dimensions(infile, input_dimensions=None):
         if is_perfect:
             break
     if best_extra_bytes > 0:
-        # TODO: If verbose mode is on...
-        if args.verbose is True:
+        if verbose is True:
             sys.stderr.write("Could not find PNG dimensions that perfectly encode "
                              "%s bytes; the encoding will be tail-padded with %s zeros.\n"
                              % (num_bytes, int(best_extra_bytes)))
     return best_dimensions
 
 
-def file_to_png(infile, outfile, dimensions=None):
-    reader = FileReader.new(infile)
-    dimensions = choose_file_dimensions(reader, dimensions)
+def file_to_png(reader, outfile, dimensions, progress=False):
+
     dim = (int(dimensions[0]), int(dimensions[1]))
     img = Image.new('RGB', dim)
     pixels = img.load()
@@ -136,7 +134,7 @@ def file_to_png(infile, outfile, dimensions=None):
         if column >= img.size[0]:
             column = 0
             row += 1
-            if args.no_progress is False:
+            if progress is True:
                 percent = float(((row + 1) // dimensions[1]) * 100)
                 sys.stderr.write("\r%s%s" % (round(percent, 2), "%"))
 
@@ -151,21 +149,46 @@ def file_to_png(infile, outfile, dimensions=None):
 
         if not row >= img.size[1]:
             pixels[column, row] = tuple(color)
-    if args.no_progress is False:
+    if progress is True:
         sys.stderr.write("\n")
-    if sys.version_info.major >= 3 and outfile.name == '<stdout>' and hasattr(outfile, 'buffer'):
+    if outfile and sys.version_info.major >= 3 and outfile.name == '<stdout>' and hasattr(outfile, 'buffer'):
         outfile = outfile.buffer
-    img.save(outfile, format="PNG")
-
-
-def png_to_file(infile, outfile, no_progress=False, verbose=False):
-    with FileReader.new(infile, file_backed=True) as reader:
-        img = Image.open(reader.name)
+    if outfile:
+        img.save(outfile, format="PNG")
+    else:
+        return img
+
+
+def png_to_file(infile, outfile, progress=True, verbose=False):
+    #with FileReader.new(infile, file_backed=True) as reader:
+    if isinstance(infile, Image.Image):
+        #In this mode, a full PNG image was already provided through input
+        isInputImage = True
+        returnArray = bytearray()
+
+    with infile as reader:
+        if isInputImage:
+            img = infile
+        else:
+            img = Image.open(reader.name)
         rgb_im = img.convert('RGB')
+        #TODO: Only store the image that is needed to save RAM?
 
+        def writeOutput(data): #Nested function to improve readability below
+            if sys.version_info.major >= 3:
+                if outfile:
+                    outfile.write(bytes([data]))
+                else:
+                    returnArray.append(data)
+            else:
+                if outfile:
+                    outfile.write(chr(data))
+                else:
+                    returnArray.append(data)
+
         pix_buffer = 0
         for row in range(img.size[1]):
-            if not no_progress:
+            if progress is True:
                 percent = float(((row + 1) // img.size[1]) * 100)
                 sys.stderr.write("\r%s%s" % (round(percent, 2), "%"))
             for col in range(img.size[0]):
@@ -182,30 +205,80 @@ def png_to_file(infile, outfile, no_progress=False, verbose=False):
                         if pix_buffer != 0:
                             for color in range(pix_buffer):
                                 # flush the cache to the file if a non-null byte was detected
-                                if sys.version_info.major >= 3:
-                                    outfile.write(bytes([0]))
-                                else:
-                                    outfile.write(chr(0))
+                                writeOutput(0)
+                                #returnArray.append(0)
                             pix_buffer = 0
-                        if sys.version_info.major >= 3:
-                            outfile.write(bytes([segment]))
-                        else:
-                            outfile.write(chr(segment))
+                        writeOutput(segment)
+                        #returnArray.append(segment)
 
-        if not no_progress:
+        if progress is True:
             sys.stderr.write("\n")
         if pix_buffer != 0 and verbose:
             length = pix_buffer
             if length == 1:
                 sys.stderr.write("Omitting %s zero from end of file\n" % pix_buffer)
             else:  # Why not...
                 sys.stderr.write("Omitting %s zeroes from end of file\n" % pix_buffer)
+
+        if not outfile:
+            #return bytearray(str(returnArray,'char'), 'char')
+            return bytes(returnArray) #.decode(encoding='UTF-8')
+
+
+
+def encode(infile, outfile=None, square=False, width=None, height=None, progress=False, verbose=False):
+    if sys.version_info.major >= 3:
+        file_mode = 'wb'
+        read_mode = 'rb'
+    else:
+        file_mode = 'w'
+        read_mode = 'r'
+
+    if isinstance(infile, str):
+        infile = open(infile, read_mode)
+    if isinstance(outfile, str):
+        outfile = open(outfile, file_mode)
+
+    dims = None
+    if width is not None or height is not None:
+        dims = (width, height)
+    #Both functions need to use the same file reader for best results.
+    reader = FileReader.new(infile)
+    calc_dimensions = choose_file_dimensions(reader, dims, square=square,verbose=verbose)
+    return file_to_png(reader, outfile, calc_dimensions, progress=progress)
+
+
+
+
+def decode(infile, outfile=None, progress=False, verbose=False):
+    if sys.version_info.major >= 3:
+        file_mode = 'wb'
+        read_mode = 'rb'
+    else:
+        file_mode = 'w'
+        read_mode = 'r'
+
+    if isinstance(infile, str):
+        infile = open(infile, read_mode)
+    if isinstance(outfile, str):
+        outfile = open(outfile, file_mode)
+
+    #if not isinstance(infile, Image.Image) and not hasattr(infile, "name"):
+    #    print("noarggggg")
+    reader = None
+    if not isinstance(infile, Image.Image): #Don't need a reader because
+        #a full image is already passed in
+        reader = FileReader.new(infile, file_backed=True)
+    else:
+        reader = infile
+
+    return png_to_file(reader, outfile, progress=progress, verbose=verbose)
 
 
 def main(argv=None):
     parser = argparse.ArgumentParser(description="A simple cross-platform script for encoding any binary file into a "
-                                                 "lossless PNG.", prog="bin2png")
-
+        "lossless PNG.", prog="bin2png")
+    
     if sys.version_info.major >= 3:
         read_mode = 'rb'
         write_mode = 'wb'
@@ -214,7 +287,8 @@ def main(argv=None):
         read_mode = 'r'
         write_mode = 'w'
         out_default = sys.stdout
-    parser.add_argument('file', type=argparse.FileType(read_mode), default=sys.stdin,
+
+    parser.add_argument('file', type=argparse.FileType(read_mode), default=sys.stdin, nargs='?',
                         help="the file to encode as a PNG (defaults to '-', which is stdin)")
     parser.add_argument("-o", "--outfile", type=argparse.FileType(write_mode), default=out_default,
                         help="the output file (defaults to '-', which is stdout)")
@@ -226,21 +300,19 @@ def main(argv=None):
                         help="constrain the output PNG to a specific height")
     parser.add_argument("-s", "--square", action="store_true", default=False, help="generate only square images")
     parser.add_argument("-v", "--verbose", action="store_true", default=False, help="enable debugging messages")
-    parser.add_argument("--no-progress", action="store_true", default=False, help="don't display percent progress")
+    parser.add_argument("-p", "--progress", action="store_true", default=False, help="display percent progress")
 
     if argv is None:
         argv = sys.argv
 
+    #Remove the first argument (the name of this file itsself)
+    argv.pop(0)
     args = parser.parse_args(argv)
 
     if args.decode:
-        png_to_file(args.file, args.outfile, no_progress=args.no_progress, verbose=args.verbose)
+        decode(reader, args.outfile, progress=args.progress, verbose=args.verbose)
     else:
-        dims = None
-        if args.height is not None or args.width is not None:
-            dims = (args.width, args.height)
-
-        file_to_png(args.file, args.outfile, dimensions=dims)
+        encode(args.file, outfile=args.outfile, square=args.square, width=args.width, height=args.height, progress=args.progress, verbose=args.verbose)
 
 
 if __name__ == "__main__":