Advertisement
TerusTheBird

makeimg_standalone.py

Apr 26th, 2020
463
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 114.83 KB | None | 0 0
  1. #!/usr/bin/python3
  2. # -*- coding: utf-8 -*-
  3. from __future__ import print_function
  4.  
  5. '''
  6. STANDALONE SPRITE CONVERTER
  7. HOW TO RUN:
  8. python makeimg_standalone.py <input_sprite>      ->  DAT/CBB/2bpp > PNG
  9. python makeimg_standalone.py -H <input_sprite>   ->  Text sprite > PNG
  10. python makeimg_standalone.py -C <input_sprite>   ->  Compressed binary sprite > PNG
  11. python makeimg_standalone.py -HC <input_sprite>  ->  Compressed text sprite > PNG
  12.  
  13. python makeimg_standalone.py -HC -d 2 -l grayhir -q <input_sprite> <output_png>
  14. '''
  15.  
  16. import os, sys, argparse, textwrap, re, math, subprocess, itertools, struct, warnings, zlib, collections
  17. sys.dont_write_bytecode = True
  18. from array import array
  19. from functools import reduce
  20. sqrt = math.sqrt
  21.  
  22. if sys.version_info > (3,0):
  23.     global xrange
  24.     xrange = range
  25.  
  26. # hex_char_encoding = 'utf-8'
  27. hex_char_encoding = 'shift_jis'
  28.  
  29. # -------------------------------------------------------------------------------------------
  30. # -------------------------------------------------------------------------------------------
  31. # -------------------------------------------------------------------------------------------
  32. # -------------------------------------------------------------------------------------------
  33. # -------------------------------------------------------------------------------------------
  34. # -------------------------------------------------------------------------------------------
  35. # -------------------------------------------------------------------------------------------
  36. # -------------------------------------------------------------------------------------------
  37. # -------------------------------------------------------------------------------------------
  38.  
  39. # png.py - PNG encoder/decoder in pure Python
  40. #
  41. # Copyright (C) 2006 Johann C. Rocholl <johann@browsershots.org>
  42. # Portions Copyright (C) 2009 David Jones <drj@pobox.com>
  43. # And probably portions Copyright (C) 2006 Nicko van Someren <nicko@nicko.org>
  44. #
  45. # Original concept by Johann C. Rocholl.
  46. #
  47. # LICENCE (MIT)
  48. #
  49. # Permission is hereby granted, free of charge, to any person
  50. # obtaining a copy of this software and associated documentation files
  51. # (the "Software"), to deal in the Software without restriction,
  52. # including without limitation the rights to use, copy, modify, merge,
  53. # publish, distribute, sublicense, and/or sell copies of the Software,
  54. # and to permit persons to whom the Software is furnished to do so,
  55. # subject to the following conditions:
  56. #
  57. # The above copyright notice and this permission notice shall be
  58. # included in all copies or substantial portions of the Software.
  59. #
  60. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  61. # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  62. # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  63. # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
  64. # BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
  65. # ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
  66. # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  67. # SOFTWARE.
  68.  
  69. """
  70. The ``png`` module can read and write PNG files.
  71.  
  72. Installation and Overview
  73. -------------------------
  74.  
  75. ``pip install pypng``
  76.  
  77. For help, type ``import png; help(png)`` in your python interpreter.
  78.  
  79. A good place to start is the :class:`Reader` and :class:`Writer` classes.
  80.  
  81. Coverage of PNG formats is fairly complete;
  82. all allowable bit depths (1/2/4/8/16/24/32/48/64 bits per pixel) and
  83. colour combinations are supported:
  84.  
  85. - greyscale (1/2/4/8/16 bit);
  86. - RGB, RGBA, LA (greyscale with alpha) with 8/16 bits per channel;
  87. - colour mapped images (1/2/4/8 bit).
  88.  
  89. Interlaced images,
  90. which support a progressive display when downloading,
  91. are supported for both reading and writing.
  92.  
  93. A number of optional chunks can be specified (when writing)
  94. and understood (when reading): ``tRNS``, ``bKGD``, ``gAMA``.
  95.  
  96. The ``sBIT`` chunk can be used to specify precision for
  97. non-native bit depths.
  98.  
  99. Requires Python 3.4 or higher (or Python 2.7).
  100. Installation is trivial,
  101. but see the ``README.txt`` file (with the source distribution) for details.
  102.  
  103. Full use of all features will need some reading of the PNG specification
  104. http://www.w3.org/TR/2003/REC-PNG-20031110/.
  105.  
  106. The package also comes with command line utilities.
  107.  
  108. - ``pripamtopng`` converts
  109.  `Netpbm <http://netpbm.sourceforge.net/>`_ PAM/PNM files to PNG;
  110. - ``pripngtopam`` converts PNG to file PAM/PNM.
  111.  
  112. There are a few more for simple PNG manipulations.
  113.  
  114. Spelling and Terminology
  115. ------------------------
  116.  
  117. Generally British English spelling is used in the documentation.
  118. So that's "greyscale" and "colour".
  119. This not only matches the author's native language,
  120. it's also used by the PNG specification.
  121.  
  122. Colour Models
  123. -------------
  124.  
  125. The major colour models supported by PNG (and hence by PyPNG) are:
  126.  
  127. - greyscale;
  128. - greyscale--alpha;
  129. - RGB;
  130. - RGB--alpha.
  131.  
  132. Also referred to using the abbreviations: L, LA, RGB, RGBA.
  133. Each letter codes a single channel:
  134. *L* is for Luminance or Luma or Lightness (greyscale images);
  135. *A* stands for Alpha, the opacity channel
  136. (used for transparency effects, but higher values are more opaque,
  137. so it makes sense to call it opacity);
  138. *R*, *G*, *B* stand for Red, Green, Blue (colour image).
  139.  
  140. Lists, arrays, sequences, and so on
  141. -----------------------------------
  142.  
  143. When getting pixel data out of this module (reading) and
  144. presenting data to this module (writing) there are
  145. a number of ways the data could be represented as a Python value.
  146.  
  147. The preferred format is a sequence of *rows*,
  148. which each row being a sequence of *values*.
  149. In this format, the values are in pixel order,
  150. with all the values from all the pixels in a row
  151. being concatenated into a single sequence for that row.
  152.  
  153. Consider an image that is 3 pixels wide by 2 pixels high, and each pixel
  154. has RGB components:
  155.  
  156. Sequence of rows::
  157.  
  158.  list([R,G,B, R,G,B, R,G,B],
  159.       [R,G,B, R,G,B, R,G,B])
  160.  
  161. Each row appears as its own list,
  162. but the pixels are flattened so that three values for one pixel
  163. simply follow the three values for the previous pixel.
  164.  
  165. This is the preferred because
  166. it provides a good compromise between space and convenience.
  167. PyPNG regards itself as at liberty to replace any sequence type with
  168. any sufficiently compatible other sequence type;
  169. in practice each row is an array (``bytearray`` or ``array.array``).
  170.  
  171. To allow streaming the outer list is sometimes
  172. an iterator rather than an explicit list.
  173.  
  174. An alternative format is a single array holding all the values.
  175.  
  176. Array of values::
  177.  
  178.  [R,G,B, R,G,B, R,G,B,
  179.   R,G,B, R,G,B, R,G,B]
  180.  
  181. The entire image is one single giant sequence of colour values.
  182. Generally an array will be used (to save space), not a list.
  183.  
  184. The top row comes first,
  185. and within each row the pixels are ordered from left-to-right.
  186. Within a pixel the values appear in the order R-G-B-A
  187. (or L-A for greyscale--alpha).
  188.  
  189. There is another format, which should only be used with caution.
  190. It is mentioned because it is used internally,
  191. is close to what lies inside a PNG file itself,
  192. and has some support from the public API.
  193. This format is called *packed*.
  194. When packed, each row is a sequence of bytes (integers from 0 to 255),
  195. just as it is before PNG scanline filtering is applied.
  196. When the bit depth is 8 this is the same as a sequence of rows;
  197. when the bit depth is less than 8 (1, 2 and 4),
  198. several pixels are packed into each byte;
  199. when the bit depth is 16 each pixel value is decomposed into 2 bytes
  200. (and `packed` is a misnomer).
  201. This format is used by the :meth:`Writer.write_packed` method.
  202. It isn't usually a convenient format,
  203. but may be just right if the source data for
  204. the PNG image comes from something that uses a similar format
  205. (for example, 1-bit BMPs, or another PNG file).
  206. """
  207.  
  208. __version__ = "0.0.20"
  209.  
  210. __all__ = ['Image', 'Reader', 'Writer', 'write_chunks', 'from_array']
  211.  
  212.  
  213. # The PNG signature.
  214. # http://www.w3.org/TR/PNG/#5PNG-file-signature
  215. signature = struct.pack('8B', 137, 80, 78, 71, 13, 10, 26, 10)
  216.  
  217. # The xstart, ystart, xstep, ystep for the Adam7 interlace passes.
  218. adam7 = ((0, 0, 8, 8),
  219.          (4, 0, 8, 8),
  220.          (0, 4, 4, 8),
  221.          (2, 0, 4, 4),
  222.          (0, 2, 2, 4),
  223.          (1, 0, 2, 2),
  224.          (0, 1, 1, 2))
  225.  
  226.  
  227. def adam7_generate(width, height):
  228.     """
  229.    Generate the coordinates for the reduced scanlines
  230.    of an Adam7 interlaced image
  231.    of size `width` by `height` pixels.
  232.  
  233.    Yields a generator for each pass,
  234.    and each pass generator yields a series of (x, y, xstep) triples,
  235.    each one identifying a reduced scanline consisting of
  236.    pixels starting at (x, y) and taking every xstep pixel to the right.
  237.    """
  238.  
  239.     for xstart, ystart, xstep, ystep in adam7:
  240.         if xstart >= width:
  241.             continue
  242.         yield ((xstart, y, xstep) for y in range(ystart, height, ystep))
  243.  
  244.  
  245. # Models the 'pHYs' chunk (used by the Reader)
  246. Resolution = collections.namedtuple('_Resolution', 'x y unit_is_meter')
  247.  
  248.  
  249. def group(s, n):
  250.     return list(zip(* [iter(s)] * n))
  251.  
  252.  
  253. def isarray(x):
  254.     return isinstance(x, array)
  255.  
  256.  
  257. def check_palette(palette):
  258.     """
  259.    Check a palette argument (to the :class:`Writer` class) for validity.
  260.    Returns the palette as a list if okay;
  261.    raises an exception otherwise.
  262.    """
  263.  
  264.     # None is the default and is allowed.
  265.     if palette is None:
  266.         return None
  267.  
  268.     p = list(palette)
  269.     if not (0 < len(p) <= 256):
  270.         raise ProtocolError(
  271.             "a palette must have between 1 and 256 entries,"
  272.             " see https://www.w3.org/TR/PNG/#11PLTE")
  273.     seen_triple = False
  274.     for i, t in enumerate(p):
  275.         if len(t) not in (3, 4):
  276.             raise ProtocolError(
  277.                 "palette entry %d: entries must be 3- or 4-tuples." % i)
  278.         if len(t) == 3:
  279.             seen_triple = True
  280.         if seen_triple and len(t) == 4:
  281.             raise ProtocolError(
  282.                 "palette entry %d: all 4-tuples must precede all 3-tuples" % i)
  283.         for x in t:
  284.             if int(x) != x or not(0 <= x <= 255):
  285.                 raise ProtocolError(
  286.                     "palette entry %d: "
  287.                     "values must be integer: 0 <= x <= 255" % i)
  288.     return p
  289.  
  290.  
  291. def check_sizes(size, width, height):
  292.     """
  293.    Check that these arguments, if supplied, are consistent.
  294.    Return a (width, height) pair.
  295.    """
  296.  
  297.     if not size:
  298.         return width, height
  299.  
  300.     if len(size) != 2:
  301.         raise ProtocolError(
  302.             "size argument should be a pair (width, height)")
  303.     if width is not None and width != size[0]:
  304.         raise ProtocolError(
  305.             "size[0] (%r) and width (%r) should match when both are used."
  306.             % (size[0], width))
  307.     if height is not None and height != size[1]:
  308.         raise ProtocolError(
  309.             "size[1] (%r) and height (%r) should match when both are used."
  310.             % (size[1], height))
  311.     return size
  312.  
  313.  
  314. def check_color(c, greyscale, which):
  315.     """
  316.    Checks that a colour argument for transparent or background options
  317.    is the right form.
  318.    Returns the colour
  319.    (which, if it's a bare integer, is "corrected" to a 1-tuple).
  320.    """
  321.  
  322.     if c is None:
  323.         return c
  324.     if greyscale:
  325.         try:
  326.             len(c)
  327.         except TypeError:
  328.             c = (c,)
  329.         if len(c) != 1:
  330.             raise ProtocolError("%s for greyscale must be 1-tuple" % which)
  331.         if not is_natural(c[0]):
  332.             raise ProtocolError(
  333.                 "%s colour for greyscale must be integer" % which)
  334.     else:
  335.         if not (len(c) == 3 and
  336.                 is_natural(c[0]) and
  337.                 is_natural(c[1]) and
  338.                 is_natural(c[2])):
  339.             raise ProtocolError(
  340.                 "%s colour must be a triple of integers" % which)
  341.     return c
  342.  
  343.  
  344. class Error(Exception):
  345.     def __str__(self):
  346.         return self.__class__.__name__ + ': ' + ' '.join(self.args)
  347.  
  348.  
  349. class FormatError(Error):
  350.     """
  351.    Problem with input file format.
  352.    In other words, PNG file does not conform to
  353.    the specification in some way and is invalid.
  354.    """
  355.  
  356.  
  357. class ProtocolError(Error):
  358.     """
  359.    Problem with the way the programming interface has been used,
  360.    or the data presented to it.
  361.    """
  362.  
  363.  
  364. class ChunkError(FormatError):
  365.     pass
  366.  
  367.  
  368. class Default:
  369.     """The default for the greyscale paramter."""
  370.  
  371.  
  372. class Writer:
  373.     """
  374.    PNG encoder in pure Python.
  375.    """
  376.  
  377.     def __init__(self, width=None, height=None,
  378.                  size=None,
  379.                  greyscale=Default,
  380.                  alpha=False,
  381.                  bitdepth=8,
  382.                  palette=None,
  383.                  transparent=None,
  384.                  background=None,
  385.                  gamma=None,
  386.                  compression=None,
  387.                  interlace=False,
  388.                  planes=None,
  389.                  colormap=None,
  390.                  maxval=None,
  391.                  chunk_limit=2**20,
  392.                  x_pixels_per_unit=None,
  393.                  y_pixels_per_unit=None,
  394.                  unit_is_meter=False):
  395.         """
  396.        Create a PNG encoder object.
  397.  
  398.        Arguments:
  399.  
  400.        width, height
  401.          Image size in pixels, as two separate arguments.
  402.        size
  403.          Image size (w,h) in pixels, as single argument.
  404.        greyscale
  405.          Pixels are greyscale, not RGB.
  406.        alpha
  407.          Input data has alpha channel (RGBA or LA).
  408.        bitdepth
  409.          Bit depth: from 1 to 16 (for each channel).
  410.        palette
  411.          Create a palette for a colour mapped image (colour type 3).
  412.        transparent
  413.          Specify a transparent colour (create a ``tRNS`` chunk).
  414.        background
  415.          Specify a default background colour (create a ``bKGD`` chunk).
  416.        gamma
  417.          Specify a gamma value (create a ``gAMA`` chunk).
  418.        compression
  419.          zlib compression level: 0 (none) to 9 (more compressed);
  420.          default: -1 or None.
  421.        interlace
  422.          Create an interlaced image.
  423.        chunk_limit
  424.          Write multiple ``IDAT`` chunks to save memory.
  425.        x_pixels_per_unit
  426.          Number of pixels a unit along the x axis (write a
  427.          `pHYs` chunk).
  428.        y_pixels_per_unit
  429.          Number of pixels a unit along the y axis (write a
  430.          `pHYs` chunk). Along with `x_pixel_unit`, this gives
  431.          the pixel size ratio.
  432.        unit_is_meter
  433.          `True` to indicate that the unit (for the `pHYs`
  434.          chunk) is metre.
  435.  
  436.        The image size (in pixels) can be specified either by using the
  437.        `width` and `height` arguments, or with the single `size`
  438.        argument.
  439.        If `size` is used it should be a pair (*width*, *height*).
  440.  
  441.        The `greyscale` argument indicates whether input pixels
  442.        are greyscale (when true), or colour (when false).
  443.        The default is true unless `palette=` is used.
  444.  
  445.        The `alpha` argument (a boolean) specifies
  446.        whether input pixels have an alpha channel (or not).
  447.  
  448.        `bitdepth` specifies the bit depth of the source pixel values.
  449.        Each channel may have a different bit depth.
  450.        Each source pixel must have values that are
  451.        an integer between 0 and ``2**bitdepth-1``, where
  452.        `bitdepth` is the bit depth for the corresponding channel.
  453.        For example, 8-bit images have values between 0 and 255.
  454.        PNG only stores images with bit depths of
  455.        1,2,4,8, or 16 (the same for all channels).
  456.        When `bitdepth` is not one of these values or where
  457.        channels have different bit depths,
  458.        the next highest valid bit depth is selected,
  459.        and an ``sBIT`` (significant bits) chunk is generated
  460.        that specifies the original precision of the source image.
  461.        In this case the supplied pixel values will be rescaled to
  462.        fit the range of the selected bit depth.
  463.  
  464.        The PNG file format supports many bit depth / colour model
  465.        combinations, but not all.
  466.        The details are somewhat arcane
  467.        (refer to the PNG specification for full details).
  468.        Briefly:
  469.        Bit depths < 8 (1,2,4) are only allowed with greyscale and
  470.        colour mapped images;
  471.        colour mapped images cannot have bit depth 16.
  472.  
  473.        For colour mapped images
  474.        (in other words, when the `palette` argument is specified)
  475.        the `bitdepth` argument must match one of
  476.        the valid PNG bit depths: 1, 2, 4, or 8.
  477.        (It is valid to have a PNG image with a palette and
  478.        an ``sBIT`` chunk, but the meaning is slightly different;
  479.        it would be awkward to use the `bitdepth` argument for this.)
  480.  
  481.        The `palette` option, when specified,
  482.        causes a colour mapped image to be created:
  483.        the PNG colour type is set to 3;
  484.        `greyscale` must not be true; `alpha` must not be true;
  485.        `transparent` must not be set.
  486.        The bit depth must be 1,2,4, or 8.
  487.        When a colour mapped image is created,
  488.        the pixel values are palette indexes and
  489.        the `bitdepth` argument specifies the size of these indexes
  490.        (not the size of the colour values in the palette).
  491.  
  492.        The palette argument value should be a sequence of 3- or
  493.        4-tuples.
  494.        3-tuples specify RGB palette entries;
  495.        4-tuples specify RGBA palette entries.
  496.        All the 4-tuples (if present) must come before all the 3-tuples.
  497.        A ``PLTE`` chunk is created;
  498.        if there are 4-tuples then a ``tRNS`` chunk is created as well.
  499.        The ``PLTE`` chunk will contain all the RGB triples in the same
  500.        sequence;
  501.        the ``tRNS`` chunk will contain the alpha channel for
  502.        all the 4-tuples, in the same sequence.
  503.        Palette entries are always 8-bit.
  504.  
  505.        If specified, the `transparent` and `background` parameters must be
  506.        a tuple with one element for each channel in the image.
  507.        Either a 3-tuple of integer (RGB) values for a colour image, or
  508.        a 1-tuple of a single integer for a greyscale image.
  509.  
  510.        If specified, the `gamma` parameter must be a positive number
  511.        (generally, a `float`).
  512.        A ``gAMA`` chunk will be created.
  513.        Note that this will not change the values of the pixels as
  514.        they appear in the PNG file,
  515.        they are assumed to have already
  516.        been converted appropriately for the gamma specified.
  517.  
  518.        The `compression` argument specifies the compression level to
  519.        be used by the ``zlib`` module.
  520.        Values from 1 to 9 (highest) specify compression.
  521.        0 means no compression.
  522.        -1 and ``None`` both mean that the ``zlib`` module uses
  523.        the default level of compession (which is generally acceptable).
  524.  
  525.        If `interlace` is true then an interlaced image is created
  526.        (using PNG's so far only interace method, *Adam7*).
  527.        This does not affect how the pixels should be passed in,
  528.        rather it changes how they are arranged into the PNG file.
  529.        On slow connexions interlaced images can be
  530.        partially decoded by the browser to give
  531.        a rough view of the image that is
  532.        successively refined as more image data appears.
  533.  
  534.        .. note ::
  535.  
  536.          Enabling the `interlace` option requires the entire image
  537.          to be processed in working memory.
  538.  
  539.        `chunk_limit` is used to limit the amount of memory used whilst
  540.        compressing the image.
  541.        In order to avoid using large amounts of memory,
  542.        multiple ``IDAT`` chunks may be created.
  543.        """
  544.  
  545.         # At the moment the `planes` argument is ignored;
  546.         # its purpose is to act as a dummy so that
  547.         # ``Writer(x, y, **info)`` works, where `info` is a dictionary
  548.         # returned by Reader.read and friends.
  549.         # Ditto for `colormap`.
  550.  
  551.         width, height = check_sizes(size, width, height)
  552.         del size
  553.  
  554.         if not is_natural(width) or not is_natural(height):
  555.             raise ProtocolError("width and height must be integers")
  556.         if width <= 0 or height <= 0:
  557.             raise ProtocolError("width and height must be greater than zero")
  558.         # http://www.w3.org/TR/PNG/#7Integers-and-byte-order
  559.         if width > 2 ** 31 - 1 or height > 2 ** 31 - 1:
  560.             raise ProtocolError("width and height cannot exceed 2**31-1")
  561.  
  562.         if alpha and transparent is not None:
  563.             raise ProtocolError(
  564.                 "transparent colour not allowed with alpha channel")
  565.  
  566.         # bitdepth is either single integer, or tuple of integers.
  567.         # Convert to tuple.
  568.         try:
  569.             len(bitdepth)
  570.         except TypeError:
  571.             bitdepth = (bitdepth, )
  572.         for b in bitdepth:
  573.             valid = is_natural(b) and 1 <= b <= 16
  574.             if not valid:
  575.                 raise ProtocolError(
  576.                     "each bitdepth %r must be a positive integer <= 16" %
  577.                     (bitdepth,))
  578.  
  579.         # Calculate channels, and
  580.         # expand bitdepth to be one element per channel.
  581.         palette = check_palette(palette)
  582.         alpha = bool(alpha)
  583.         colormap = bool(palette)
  584.         if greyscale is Default and palette:
  585.             greyscale = False
  586.         greyscale = bool(greyscale)
  587.         if colormap:
  588.             color_planes = 1
  589.             planes = 1
  590.         else:
  591.             color_planes = (3, 1)[greyscale]
  592.             planes = color_planes + alpha
  593.         if len(bitdepth) == 1:
  594.             bitdepth *= planes
  595.  
  596.         bitdepth, self.rescale = check_bitdepth_rescale(
  597.                 palette,
  598.                 bitdepth,
  599.                 transparent, alpha, greyscale)
  600.  
  601.         # These are assertions, because above logic should have
  602.         # corrected or raised all problematic cases.
  603.         if bitdepth < 8:
  604.             assert greyscale or palette
  605.             assert not alpha
  606.         if bitdepth > 8:
  607.             assert not palette
  608.  
  609.         transparent = check_color(transparent, greyscale, 'transparent')
  610.         background = check_color(background, greyscale, 'background')
  611.  
  612.         # It's important that the true boolean values
  613.         # (greyscale, alpha, colormap, interlace) are converted
  614.         # to bool because Iverson's convention is relied upon later on.
  615.         self.width = width
  616.         self.height = height
  617.         self.transparent = transparent
  618.         self.background = background
  619.         self.gamma = gamma
  620.         self.greyscale = greyscale
  621.         self.alpha = alpha
  622.         self.colormap = colormap
  623.         self.bitdepth = int(bitdepth)
  624.         self.compression = compression
  625.         self.chunk_limit = chunk_limit
  626.         self.interlace = bool(interlace)
  627.         self.palette = palette
  628.         self.x_pixels_per_unit = x_pixels_per_unit
  629.         self.y_pixels_per_unit = y_pixels_per_unit
  630.         self.unit_is_meter = bool(unit_is_meter)
  631.  
  632.         self.color_type = (4 * self.alpha +
  633.                            2 * (not greyscale) +
  634.                            1 * self.colormap)
  635.         assert self.color_type in (0, 2, 3, 4, 6)
  636.  
  637.         self.color_planes = color_planes
  638.         self.planes = planes
  639.         # :todo: fix for bitdepth < 8
  640.         self.psize = (self.bitdepth / 8) * self.planes
  641.  
  642.     def write(self, outfile, rows):
  643.         """
  644.        Write a PNG image to the output file.
  645.        `rows` should be an iterable that yields each row
  646.        (each row is a sequence of values).
  647.        The rows should be the rows of the original image,
  648.        so there should be ``self.height`` rows of
  649.        ``self.width * self.planes`` values.
  650.        If `interlace` is specified (when creating the instance),
  651.        then an interlaced PNG file will be written.
  652.        Supply the rows in the normal image order;
  653.        the interlacing is carried out internally.
  654.  
  655.        .. note ::
  656.  
  657.          Interlacing requires the entire image to be in working memory.
  658.        """
  659.  
  660.         # Values per row
  661.         vpr = self.width * self.planes
  662.  
  663.         def check_rows(rows):
  664.             """
  665.            Yield each row in rows,
  666.            but check each row first (for correct width).
  667.            """
  668.             for i, row in enumerate(rows):
  669.                 try:
  670.                     wrong_length = len(row) != vpr
  671.                 except TypeError:
  672.                     # When using an itertools.ichain object or
  673.                     # other generator not supporting __len__,
  674.                     # we set this to False to skip the check.
  675.                     wrong_length = False
  676.                 if wrong_length:
  677.                     # Note: row numbers start at 0.
  678.                     raise ProtocolError(
  679.                         "Expected %d values but got %d value, in row %d" %
  680.                         (vpr, len(row), i))
  681.                 yield row
  682.  
  683.         if self.interlace:
  684.             fmt = 'BH'[self.bitdepth > 8]
  685.             a = array(fmt, itertools.chain(*check_rows(rows)))
  686.             return self.write_array(outfile, a)
  687.  
  688.         nrows = self.write_passes(outfile, check_rows(rows))
  689.         if nrows != self.height:
  690.             raise ProtocolError(
  691.                 "rows supplied (%d) does not match height (%d)" %
  692.                 (nrows, self.height))
  693.  
  694.     def write_passes(self, outfile, rows):
  695.         """
  696.        Write a PNG image to the output file.
  697.  
  698.        Most users are expected to find the :meth:`write` or
  699.        :meth:`write_array` method more convenient.
  700.  
  701.        The rows should be given to this method in the order that
  702.        they appear in the output file.
  703.        For straightlaced images, this is the usual top to bottom ordering.
  704.        For interlaced images the rows should have been interlaced before
  705.        passing them to this function.
  706.  
  707.        `rows` should be an iterable that yields each row
  708.        (each row being a sequence of values).
  709.        """
  710.  
  711.         # Ensure rows are scaled (to 4-/8-/16-bit),
  712.         # and packed into bytes.
  713.  
  714.         if self.rescale:
  715.             rows = rescale_rows(rows, self.rescale)
  716.  
  717.         if self.bitdepth < 8:
  718.             rows = pack_rows(rows, self.bitdepth)
  719.         elif self.bitdepth == 16:
  720.             rows = unpack_rows(rows)
  721.  
  722.         return self.write_packed(outfile, rows)
  723.  
  724.     def write_packed(self, outfile, rows):
  725.         """
  726.        Write PNG file to `outfile`.
  727.        `rows` should be an iterator that yields each packed row;
  728.        a packed row being a sequence of packed bytes.
  729.  
  730.        The rows have a filter byte prefixed and
  731.        are then compressed into one or more IDAT chunks.
  732.        They are not processed any further,
  733.        so if bitdepth is other than 1, 2, 4, 8, 16,
  734.        the pixel values should have been scaled
  735.        before passing them to this method.
  736.  
  737.        This method does work for interlaced images but it is best avoided.
  738.        For interlaced images, the rows should be
  739.        presented in the order that they appear in the file.
  740.        """
  741.  
  742.         self.write_preamble(outfile)
  743.  
  744.         # http://www.w3.org/TR/PNG/#11IDAT
  745.         if self.compression is not None:
  746.             compressor = zlib.compressobj(self.compression)
  747.         else:
  748.             compressor = zlib.compressobj()
  749.  
  750.         # data accumulates bytes to be compressed for the IDAT chunk;
  751.         # it's compressed when sufficiently large.
  752.         data = bytearray()
  753.  
  754.         for i, row in enumerate(rows):
  755.             # Add "None" filter type.
  756.             # Currently, it's essential that this filter type be used
  757.             # for every scanline as
  758.             # we do not mark the first row of a reduced pass image;
  759.             # that means we could accidentally compute
  760.             # the wrong filtered scanline if we used
  761.             # "up", "average", or "paeth" on such a line.
  762.             data.append(0)
  763.             data.extend(row)
  764.             if len(data) > self.chunk_limit:
  765.                 # :todo: bytes() only necessary in Python 2
  766.                 compressed = compressor.compress(bytes(data))
  767.                 if len(compressed):
  768.                     write_chunk(outfile, b'IDAT', compressed)
  769.                 data = bytearray()
  770.  
  771.         compressed = compressor.compress(bytes(data))
  772.         flushed = compressor.flush()
  773.         if len(compressed) or len(flushed):
  774.             write_chunk(outfile, b'IDAT', compressed + flushed)
  775.         # http://www.w3.org/TR/PNG/#11IEND
  776.         write_chunk(outfile, b'IEND')
  777.         return i + 1
  778.  
  779.     def write_preamble(self, outfile):
  780.         # http://www.w3.org/TR/PNG/#5PNG-file-signature
  781.         outfile.write(signature)
  782.  
  783.         # http://www.w3.org/TR/PNG/#11IHDR
  784.         write_chunk(outfile, b'IHDR',
  785.                     struct.pack("!2I5B", self.width, self.height,
  786.                                 self.bitdepth, self.color_type,
  787.                                 0, 0, self.interlace))
  788.  
  789.         # See :chunk:order
  790.         # http://www.w3.org/TR/PNG/#11gAMA
  791.         if self.gamma is not None:
  792.             write_chunk(outfile, b'gAMA',
  793.                         struct.pack("!L", int(round(self.gamma * 1e5))))
  794.  
  795.         # See :chunk:order
  796.         # http://www.w3.org/TR/PNG/#11sBIT
  797.         if self.rescale:
  798.             write_chunk(
  799.                 outfile, b'sBIT',
  800.                 struct.pack('%dB' % self.planes,
  801.                             * [s[0] for s in self.rescale]))
  802.  
  803.         # :chunk:order: Without a palette (PLTE chunk),
  804.         # ordering is relatively relaxed.
  805.         # With one, gAMA chunk must precede PLTE chunk
  806.         # which must precede tRNS and bKGD.
  807.         # See http://www.w3.org/TR/PNG/#5ChunkOrdering
  808.         if self.palette:
  809.             p, t = make_palette_chunks(self.palette)
  810.             write_chunk(outfile, b'PLTE', p)
  811.             if t:
  812.                 # tRNS chunk is optional;
  813.                 # Only needed if palette entries have alpha.
  814.                 write_chunk(outfile, b'tRNS', t)
  815.  
  816.         # http://www.w3.org/TR/PNG/#11tRNS
  817.         if self.transparent is not None:
  818.             if self.greyscale:
  819.                 fmt = "!1H"
  820.             else:
  821.                 fmt = "!3H"
  822.             write_chunk(outfile, b'tRNS',
  823.                         struct.pack(fmt, *self.transparent))
  824.  
  825.         # http://www.w3.org/TR/PNG/#11bKGD
  826.         if self.background is not None:
  827.             if self.greyscale:
  828.                 fmt = "!1H"
  829.             else:
  830.                 fmt = "!3H"
  831.             write_chunk(outfile, b'bKGD',
  832.                         struct.pack(fmt, *self.background))
  833.  
  834.         # http://www.w3.org/TR/PNG/#11pHYs
  835.         if (self.x_pixels_per_unit is not None and
  836.                 self.y_pixels_per_unit is not None):
  837.             tup = (self.x_pixels_per_unit,
  838.                    self.y_pixels_per_unit,
  839.                    int(self.unit_is_meter))
  840.             write_chunk(outfile, b'pHYs', struct.pack("!LLB", *tup))
  841.  
  842.     def write_array(self, outfile, pixels):
  843.         """
  844.        Write an array that holds all the image values
  845.        as a PNG file on the output file.
  846.        See also :meth:`write` method.
  847.        """
  848.  
  849.         if self.interlace:
  850.             if type(pixels) != array:
  851.                 # Coerce to array type
  852.                 fmt = 'BH'[self.bitdepth > 8]
  853.                 pixels = array(fmt, pixels)
  854.             self.write_passes(outfile, self.array_scanlines_interlace(pixels))
  855.         else:
  856.             self.write_passes(outfile, self.array_scanlines(pixels))
  857.  
  858.     def array_scanlines(self, pixels):
  859.         """
  860.        Generates rows (each a sequence of values) from
  861.        a single array of values.
  862.        """
  863.  
  864.         # Values per row
  865.         vpr = self.width * self.planes
  866.         stop = 0
  867.         for y in range(self.height):
  868.             start = stop
  869.             stop = start + vpr
  870.             yield pixels[start:stop]
  871.  
  872.     def array_scanlines_interlace(self, pixels):
  873.         """
  874.        Generator for interlaced scanlines from an array.
  875.        `pixels` is the full source image as a single array of values.
  876.        The generator yields each scanline of the reduced passes in turn,
  877.        each scanline being a sequence of values.
  878.        """
  879.  
  880.         # http://www.w3.org/TR/PNG/#8InterlaceMethods
  881.         # Array type.
  882.         fmt = 'BH'[self.bitdepth > 8]
  883.         # Value per row
  884.         vpr = self.width * self.planes
  885.  
  886.         # Each iteration generates a scanline starting at (x, y)
  887.         # and consisting of every xstep pixels.
  888.         for lines in adam7_generate(self.width, self.height):
  889.             for x, y, xstep in lines:
  890.                 # Pixels per row (of reduced image)
  891.                 ppr = int(math.ceil((self.width - x) / float(xstep)))
  892.                 # Values per row (of reduced image)
  893.                 reduced_row_len = ppr * self.planes
  894.                 if xstep == 1:
  895.                     # Easy case: line is a simple slice.
  896.                     offset = y * vpr
  897.                     yield pixels[offset: offset + vpr]
  898.                     continue
  899.                 # We have to step by xstep,
  900.                 # which we can do one plane at a time
  901.                 # using the step in Python slices.
  902.                 row = array(fmt)
  903.                 # There's no easier way to set the length of an array
  904.                 row.extend(pixels[0:reduced_row_len])
  905.                 offset = y * vpr + x * self.planes
  906.                 end_offset = (y + 1) * vpr
  907.                 skip = self.planes * xstep
  908.                 for i in range(self.planes):
  909.                     row[i::self.planes] = \
  910.                         pixels[offset + i: end_offset: skip]
  911.                 yield row
  912.  
  913.  
  914. def write_chunk(outfile, tag, data=b''):
  915.     """
  916.    Write a PNG chunk to the output file, including length and
  917.    checksum.
  918.    """
  919.  
  920.     data = bytes(data)
  921.     # http://www.w3.org/TR/PNG/#5Chunk-layout
  922.     outfile.write(struct.pack("!I", len(data)))
  923.     outfile.write(tag)
  924.     outfile.write(data)
  925.     checksum = zlib.crc32(tag)
  926.     checksum = zlib.crc32(data, checksum)
  927.     checksum &= 2 ** 32 - 1
  928.     outfile.write(struct.pack("!I", checksum))
  929.  
  930.  
  931. def write_chunks(out, chunks):
  932.     """Create a PNG file by writing out the chunks."""
  933.  
  934.     out.write(signature)
  935.     for chunk in chunks:
  936.         write_chunk(out, *chunk)
  937.  
  938.  
  939. def rescale_rows(rows, rescale):
  940.     """
  941.    Take each row in rows (an iterator) and yield
  942.    a fresh row with the pixels scaled according to
  943.    the rescale parameters in the list `rescale`.
  944.    Each element of `rescale` is a tuple of
  945.    (source_bitdepth, target_bitdepth),
  946.    with one element per channel.
  947.    """
  948.  
  949.     # One factor for each channel
  950.     fs = [float(2 ** s[1] - 1)/float(2 ** s[0] - 1)
  951.           for s in rescale]
  952.  
  953.     # Assume all target_bitdepths are the same
  954.     target_bitdepths = set(s[1] for s in rescale)
  955.     assert len(target_bitdepths) == 1
  956.     (target_bitdepth, ) = target_bitdepths
  957.     typecode = 'BH'[target_bitdepth > 8]
  958.  
  959.     # Number of channels
  960.     n_chans = len(rescale)
  961.  
  962.     for row in rows:
  963.         rescaled_row = array(typecode, iter(row))
  964.         for i in range(n_chans):
  965.             channel = array(
  966.                 typecode,
  967.                 (int(round(fs[i] * x)) for x in row[i::n_chans]))
  968.             rescaled_row[i::n_chans] = channel
  969.         yield rescaled_row
  970.  
  971.  
  972. def pack_rows(rows, bitdepth):
  973.     """Yield packed rows that are a byte array.
  974.    Each byte is packed with the values from several pixels.
  975.    """
  976.  
  977.     assert bitdepth < 8
  978.     assert 8 % bitdepth == 0
  979.  
  980.     # samples per byte
  981.     spb = int(8 / bitdepth)
  982.  
  983.     def make_byte(block):
  984.         """Take a block of (2, 4, or 8) values,
  985.        and pack them into a single byte.
  986.        """
  987.  
  988.         res = 0
  989.         for v in block:
  990.             res = (res << bitdepth) + v
  991.         return res
  992.  
  993.     for row in rows:
  994.         a = bytearray(row)
  995.         # Adding padding bytes so we can group into a whole
  996.         # number of spb-tuples.
  997.         n = float(len(a))
  998.         extra = math.ceil(n / spb) * spb - n
  999.         a.extend([0] * int(extra))
  1000.         # Pack into bytes.
  1001.         # Each block is the samples for one byte.
  1002.         blocks = group(a, spb)
  1003.         yield bytearray(make_byte(block) for block in blocks)
  1004.  
  1005.  
  1006. def unpack_rows(rows):
  1007.     """Unpack each row from being 16-bits per value,
  1008.    to being a sequence of bytes.
  1009.    """
  1010.     for row in rows:
  1011.         fmt = '!%dH' % len(row)
  1012.         yield bytearray(struct.pack(fmt, *row))
  1013.  
  1014.  
  1015. def make_palette_chunks(palette):
  1016.     """
  1017.    Create the byte sequences for a ``PLTE`` and
  1018.    if necessary a ``tRNS`` chunk.
  1019.    Returned as a pair (*p*, *t*).
  1020.    *t* will be ``None`` if no ``tRNS`` chunk is necessary.
  1021.    """
  1022.  
  1023.     p = bytearray()
  1024.     t = bytearray()
  1025.  
  1026.     for x in palette:
  1027.         p.extend(x[0:3])
  1028.         if len(x) > 3:
  1029.             t.append(x[3])
  1030.     if t:
  1031.         return p, t
  1032.     return p, None
  1033.  
  1034.  
  1035. def check_bitdepth_rescale(
  1036.         palette, bitdepth, transparent, alpha, greyscale):
  1037.     """
  1038.    Returns (bitdepth, rescale) pair.
  1039.    """
  1040.  
  1041.     if palette:
  1042.         if len(bitdepth) != 1:
  1043.             raise ProtocolError(
  1044.                 "with palette, only a single bitdepth may be used")
  1045.         (bitdepth, ) = bitdepth
  1046.         if bitdepth not in (1, 2, 4, 8):
  1047.             raise ProtocolError(
  1048.                 "with palette, bitdepth must be 1, 2, 4, or 8")
  1049.         if transparent is not None:
  1050.             raise ProtocolError("transparent and palette not compatible")
  1051.         if alpha:
  1052.             raise ProtocolError("alpha and palette not compatible")
  1053.         if greyscale:
  1054.             raise ProtocolError("greyscale and palette not compatible")
  1055.         return bitdepth, None
  1056.  
  1057.     # No palette, check for sBIT chunk generation.
  1058.  
  1059.     if greyscale and not alpha:
  1060.         # Single channel, L.
  1061.         (bitdepth,) = bitdepth
  1062.         if bitdepth in (1, 2, 4, 8, 16):
  1063.             return bitdepth, None
  1064.         if bitdepth > 8:
  1065.             targetbitdepth = 16
  1066.         elif bitdepth == 3:
  1067.             targetbitdepth = 4
  1068.         else:
  1069.             assert bitdepth in (5, 6, 7)
  1070.             targetbitdepth = 8
  1071.         return targetbitdepth, [(bitdepth, targetbitdepth)]
  1072.  
  1073.     assert alpha or not greyscale
  1074.  
  1075.     depth_set = tuple(set(bitdepth))
  1076.     if depth_set in [(8,), (16,)]:
  1077.         # No sBIT required.
  1078.         (bitdepth, ) = depth_set
  1079.         return bitdepth, None
  1080.  
  1081.     targetbitdepth = (8, 16)[max(bitdepth) > 8]
  1082.     return targetbitdepth, [(b, targetbitdepth) for b in bitdepth]
  1083.  
  1084.  
  1085. # Regex for decoding mode string
  1086. RegexModeDecode = re.compile("(LA?|RGBA?);?([0-9]*)", flags=re.IGNORECASE)
  1087.  
  1088.  
  1089. def from_array(a, mode=None, info={}):
  1090.     """
  1091.    Create a PNG :class:`Image` object from a 2-dimensional array.
  1092.    One application of this function is easy PIL-style saving:
  1093.    ``png.from_array(pixels, 'L').save('foo.png')``.
  1094.  
  1095.    Unless they are specified using the *info* parameter,
  1096.    the PNG's height and width are taken from the array size.
  1097.    The first axis is the height; the second axis is the
  1098.    ravelled width and channel index.
  1099.    The array is treated is a sequence of rows,
  1100.    each row being a sequence of values (``width*channels`` in number).
  1101.    So an RGB image that is 16 pixels high and 8 wide will
  1102.    occupy a 2-dimensional array that is 16x24
  1103.    (each row will be 8*3 = 24 sample values).
  1104.  
  1105.    *mode* is a string that specifies the image colour format in a
  1106.    PIL-style mode.  It can be:
  1107.  
  1108.    ``'L'``
  1109.      greyscale (1 channel)
  1110.    ``'LA'``
  1111.      greyscale with alpha (2 channel)
  1112.    ``'RGB'``
  1113.      colour image (3 channel)
  1114.    ``'RGBA'``
  1115.      colour image with alpha (4 channel)
  1116.  
  1117.    The mode string can also specify the bit depth
  1118.    (overriding how this function normally derives the bit depth,
  1119.    see below).
  1120.    Appending ``';16'`` to the mode will cause the PNG to be
  1121.    16 bits per channel;
  1122.    any decimal from 1 to 16 can be used to specify the bit depth.
  1123.  
  1124.    When a 2-dimensional array is used *mode* determines how many
  1125.    channels the image has, and so allows the width to be derived from
  1126.    the second array dimension.
  1127.  
  1128.    The array is expected to be a ``numpy`` array,
  1129.    but it can be any suitable Python sequence.
  1130.    For example, a list of lists can be used:
  1131.    ``png.from_array([[0, 255, 0], [255, 0, 255]], 'L')``.
  1132.    The exact rules are: ``len(a)`` gives the first dimension, height;
  1133.    ``len(a[0])`` gives the second dimension.
  1134.    It's slightly more complicated than that because
  1135.    an iterator of rows can be used, and it all still works.
  1136.    Using an iterator allows data to be streamed efficiently.
  1137.  
  1138.    The bit depth of the PNG is normally taken from
  1139.    the array element's datatype
  1140.    (but if *mode* specifies a bitdepth then that is used instead).
  1141.    The array element's datatype is determined in a way which
  1142.    is supposed to work both for ``numpy`` arrays and for Python
  1143.    ``array.array`` objects.
  1144.    A 1 byte datatype will give a bit depth of 8,
  1145.    a 2 byte datatype will give a bit depth of 16.
  1146.    If the datatype does not have an implicit size,
  1147.    like the above example where it is a plain Python list of lists,
  1148.    then a default of 8 is used.
  1149.  
  1150.    The *info* parameter is a dictionary that can
  1151.    be used to specify metadata (in the same style as
  1152.    the arguments to the :class:`png.Writer` class).
  1153.    For this function the keys that are useful are:
  1154.  
  1155.    height
  1156.      overrides the height derived from the array dimensions and
  1157.      allows *a* to be an iterable.
  1158.    width
  1159.      overrides the width derived from the array dimensions.
  1160.    bitdepth
  1161.      overrides the bit depth derived from the element datatype
  1162.      (but must match *mode* if that also specifies a bit depth).
  1163.  
  1164.    Generally anything specified in the *info* dictionary will
  1165.    override any implicit choices that this function would otherwise make,
  1166.    but must match any explicit ones.
  1167.    For example, if the *info* dictionary has a ``greyscale`` key then
  1168.    this must be true when mode is ``'L'`` or ``'LA'`` and
  1169.    false when mode is ``'RGB'`` or ``'RGBA'``.
  1170.    """
  1171.  
  1172.     # We abuse the *info* parameter by modifying it.  Take a copy here.
  1173.     # (Also typechecks *info* to some extent).
  1174.     info = dict(info)
  1175.  
  1176.     # Syntax check mode string.
  1177.     match = RegexModeDecode.match(mode)
  1178.     if not match:
  1179.         raise Error("mode string should be 'RGB' or 'L;16' or similar.")
  1180.  
  1181.     mode, bitdepth = match.groups()
  1182.     if bitdepth:
  1183.         bitdepth = int(bitdepth)
  1184.  
  1185.     # Colour format.
  1186.     if 'greyscale' in info:
  1187.         if bool(info['greyscale']) != ('L' in mode):
  1188.             raise ProtocolError("info['greyscale'] should match mode.")
  1189.     info['greyscale'] = 'L' in mode
  1190.  
  1191.     alpha = 'A' in mode
  1192.     if 'alpha' in info:
  1193.         if bool(info['alpha']) != alpha:
  1194.             raise ProtocolError("info['alpha'] should match mode.")
  1195.     info['alpha'] = alpha
  1196.  
  1197.     # Get bitdepth from *mode* if possible.
  1198.     if bitdepth:
  1199.         if info.get("bitdepth") and bitdepth != info['bitdepth']:
  1200.             raise ProtocolError(
  1201.                 "bitdepth (%d) should match bitdepth of info (%d)." %
  1202.                 (bitdepth, info['bitdepth']))
  1203.         info['bitdepth'] = bitdepth
  1204.  
  1205.     # Fill in and/or check entries in *info*.
  1206.     # Dimensions.
  1207.     width, height = check_sizes(
  1208.         info.get("size"),
  1209.         info.get("width"),
  1210.         info.get("height"))
  1211.     if width:
  1212.         info["width"] = width
  1213.     if height:
  1214.         info["height"] = height
  1215.  
  1216.     if "height" not in info:
  1217.         try:
  1218.             info['height'] = len(a)
  1219.         except TypeError:
  1220.             raise ProtocolError(
  1221.                 "len(a) does not work, supply info['height'] instead.")
  1222.  
  1223.     planes = len(mode)
  1224.     if 'planes' in info:
  1225.         if info['planes'] != planes:
  1226.             raise Error("info['planes'] should match mode.")
  1227.  
  1228.     # In order to work out whether we the array is 2D or 3D we need its
  1229.     # first row, which requires that we take a copy of its iterator.
  1230.     # We may also need the first row to derive width and bitdepth.
  1231.     a, t = itertools.tee(a)
  1232.     row = next(t)
  1233.     del t
  1234.  
  1235.     testelement = row
  1236.     if 'width' not in info:
  1237.         width = len(row) // planes
  1238.         info['width'] = width
  1239.  
  1240.     if 'bitdepth' not in info:
  1241.         try:
  1242.             dtype = testelement.dtype
  1243.             # goto the "else:" clause.  Sorry.
  1244.         except AttributeError:
  1245.             try:
  1246.                 # Try a Python array.array.
  1247.                 bitdepth = 8 * testelement.itemsize
  1248.             except AttributeError:
  1249.                 # We can't determine it from the array element's datatype,
  1250.                 # use a default of 8.
  1251.                 bitdepth = 8
  1252.         else:
  1253.             # If we got here without exception,
  1254.             # we now assume that the array is a numpy array.
  1255.             if dtype.kind == 'b':
  1256.                 bitdepth = 1
  1257.             else:
  1258.                 bitdepth = 8 * dtype.itemsize
  1259.         info['bitdepth'] = bitdepth
  1260.  
  1261.     for thing in ["width", "height", "bitdepth", "greyscale", "alpha"]:
  1262.         assert thing in info
  1263.  
  1264.     return Image(a, info)
  1265.  
  1266.  
  1267. # So that refugee's from PIL feel more at home.  Not documented.
  1268. fromarray = from_array
  1269.  
  1270.  
  1271. class Image:
  1272.     """A PNG image.  You can create an :class:`Image` object from
  1273.    an array of pixels by calling :meth:`png.from_array`.  It can be
  1274.    saved to disk with the :meth:`save` method.
  1275.    """
  1276.  
  1277.     def __init__(self, rows, info):
  1278.         """
  1279.        .. note ::
  1280.  
  1281.          The constructor is not public.  Please do not call it.
  1282.        """
  1283.  
  1284.         self.rows = rows
  1285.         self.info = info
  1286.  
  1287.     def save(self, file):
  1288.         """Save the image to the named *file*.
  1289.  
  1290.        See `.write()` if you already have an open file object.
  1291.  
  1292.        In general, you can only call this method once;
  1293.        after it has been called the first time the PNG image is written,
  1294.        the source data will have been streamed, and
  1295.        cannot be streamed again.
  1296.        """
  1297.  
  1298.         w = Writer(**self.info)
  1299.  
  1300.         with open(file, 'wb') as fd:
  1301.             w.write(fd, self.rows)
  1302.  
  1303.     def write(self, file):
  1304.         """Write the image to the open file object.
  1305.  
  1306.        See `.save()` if you have a filename.
  1307.  
  1308.        In general, you can only call this method once;
  1309.        after it has been called the first time the PNG image is written,
  1310.        the source data will have been streamed, and
  1311.        cannot be streamed again.
  1312.        """
  1313.  
  1314.         w = Writer(**self.info)
  1315.         w.write(file, self.rows)
  1316.  
  1317.  
  1318. class Reader:
  1319.     """
  1320.    Pure Python PNG decoder in pure Python.
  1321.    """
  1322.  
  1323.     def __init__(self, _guess=None, filename=None, file=None, bytes=None):
  1324.         """
  1325.        The constructor expects exactly one keyword argument.
  1326.        If you supply a positional argument instead,
  1327.        it will guess the input type.
  1328.        Choose from the following keyword arguments:
  1329.  
  1330.        filename
  1331.          Name of input file (a PNG file).
  1332.        file
  1333.          A file-like object (object with a read() method).
  1334.        bytes
  1335.          ``bytes`` or ``bytearray`` with PNG data.
  1336.  
  1337.        """
  1338.         keywords_supplied = (
  1339.             (_guess is not None) +
  1340.             (filename is not None) +
  1341.             (file is not None) +
  1342.             (bytes is not None))
  1343.         if keywords_supplied != 1:
  1344.             raise TypeError("Reader() takes exactly 1 argument")
  1345.  
  1346.         # Will be the first 8 bytes, later on.  See validate_signature.
  1347.         self.signature = None
  1348.         self.transparent = None
  1349.         # A pair of (len,type) if a chunk has been read but its data and
  1350.         # checksum have not (in other words the file position is just
  1351.         # past the 4 bytes that specify the chunk type).
  1352.         # See preamble method for how this is used.
  1353.         self.atchunk = None
  1354.  
  1355.         if _guess is not None:
  1356.             if isarray(_guess):
  1357.                 bytes = _guess
  1358.             elif isinstance(_guess, str):
  1359.                 filename = _guess
  1360.             elif hasattr(_guess, 'read'):
  1361.                 file = _guess
  1362.  
  1363.         if bytes is not None:
  1364.             self.file = io.BytesIO(bytes)
  1365.         elif filename is not None:
  1366.             self.file = open(filename, "rb")
  1367.         elif file is not None:
  1368.             self.file = file
  1369.         else:
  1370.             raise ProtocolError("expecting filename, file or bytes array")
  1371.  
  1372.     def chunk(self, lenient=False):
  1373.         """
  1374.        Read the next PNG chunk from the input file;
  1375.        returns a (*type*, *data*) tuple.
  1376.        *type* is the chunk's type as a byte string
  1377.        (all PNG chunk types are 4 bytes long).
  1378.        *data* is the chunk's data content, as a byte string.
  1379.  
  1380.        If the optional `lenient` argument evaluates to `True`,
  1381.        checksum failures will raise warnings rather than exceptions.
  1382.        """
  1383.  
  1384.         self.validate_signature()
  1385.  
  1386.         # http://www.w3.org/TR/PNG/#5Chunk-layout
  1387.         if not self.atchunk:
  1388.             self.atchunk = self._chunk_len_type()
  1389.         if not self.atchunk:
  1390.             raise ChunkError("No more chunks.")
  1391.         length, type = self.atchunk
  1392.         self.atchunk = None
  1393.  
  1394.         data = self.file.read(length)
  1395.         if len(data) != length:
  1396.             raise ChunkError(
  1397.                 'Chunk %s too short for required %i octets.'
  1398.                 % (type, length))
  1399.         checksum = self.file.read(4)
  1400.         if len(checksum) != 4:
  1401.             raise ChunkError('Chunk %s too short for checksum.' % type)
  1402.         verify = zlib.crc32(type)
  1403.         verify = zlib.crc32(data, verify)
  1404.         # Whether the output from zlib.crc32 is signed or not varies
  1405.         # according to hideous implementation details, see
  1406.         # http://bugs.python.org/issue1202 .
  1407.         # We coerce it to be positive here (in a way which works on
  1408.         # Python 2.3 and older).
  1409.         verify &= 2**32 - 1
  1410.         verify = struct.pack('!I', verify)
  1411.         if checksum != verify:
  1412.             (a, ) = struct.unpack('!I', checksum)
  1413.             (b, ) = struct.unpack('!I', verify)
  1414.             message = ("Checksum error in %s chunk: 0x%08X != 0x%08X."
  1415.                        % (type.decode('ascii'), a, b))
  1416.             if lenient:
  1417.                 warnings.warn(message, RuntimeWarning)
  1418.             else:
  1419.                 raise ChunkError(message)
  1420.         return type, data
  1421.  
  1422.     def chunks(self):
  1423.         """Return an iterator that will yield each chunk as a
  1424.        (*chunktype*, *content*) pair.
  1425.        """
  1426.  
  1427.         while True:
  1428.             t, v = self.chunk()
  1429.             yield t, v
  1430.             if t == b'IEND':
  1431.                 break
  1432.  
  1433.     def undo_filter(self, filter_type, scanline, previous):
  1434.         """
  1435.        Undo the filter for a scanline.
  1436.        `scanline` is a sequence of bytes that
  1437.        does not include the initial filter type byte.
  1438.        `previous` is decoded previous scanline
  1439.        (for straightlaced images this is the previous pixel row,
  1440.        but for interlaced images, it is
  1441.        the previous scanline in the reduced image,
  1442.        which in general is not the previous pixel row in the final image).
  1443.        When there is no previous scanline
  1444.        (the first row of a straightlaced image,
  1445.        or the first row in one of the passes in an interlaced image),
  1446.        then this argument should be ``None``.
  1447.  
  1448.        The scanline will have the effects of filtering removed;
  1449.        the result will be returned as a fresh sequence of bytes.
  1450.        """
  1451.  
  1452.         # :todo: Would it be better to update scanline in place?
  1453.         result = scanline
  1454.  
  1455.         if filter_type == 0:
  1456.             return result
  1457.  
  1458.         if filter_type not in (1, 2, 3, 4):
  1459.             raise FormatError(
  1460.                 'Invalid PNG Filter Type.  '
  1461.                 'See http://www.w3.org/TR/2003/REC-PNG-20031110/#9Filters .')
  1462.  
  1463.         # Filter unit.  The stride from one pixel to the corresponding
  1464.         # byte from the previous pixel.  Normally this is the pixel
  1465.         # size in bytes, but when this is smaller than 1, the previous
  1466.         # byte is used instead.
  1467.         fu = max(1, self.psize)
  1468.  
  1469.         # For the first line of a pass, synthesize a dummy previous
  1470.         # line.  An alternative approach would be to observe that on the
  1471.         # first line 'up' is the same as 'null', 'paeth' is the same
  1472.         # as 'sub', with only 'average' requiring any special case.
  1473.         if not previous:
  1474.             previous = bytearray([0] * len(scanline))
  1475.  
  1476.         # Call appropriate filter algorithm.  Note that 0 has already
  1477.         # been dealt with.
  1478.         fn = (None,
  1479.               undo_filter_sub,
  1480.               undo_filter_up,
  1481.               undo_filter_average,
  1482.               undo_filter_paeth)[filter_type]
  1483.         fn(fu, scanline, previous, result)
  1484.         return result
  1485.  
  1486.     def _deinterlace(self, raw):
  1487.         """
  1488.        Read raw pixel data, undo filters, deinterlace, and flatten.
  1489.        Return a single array of values.
  1490.        """
  1491.  
  1492.         # Values per row (of the target image)
  1493.         vpr = self.width * self.planes
  1494.  
  1495.         # Values per image
  1496.         vpi = vpr * self.height
  1497.         # Interleaving writes to the output array randomly
  1498.         # (well, not quite), so the entire output array must be in memory.
  1499.         # Make a result array, and make it big enough.
  1500.         if self.bitdepth > 8:
  1501.             a = array('H', [0] * vpi)
  1502.         else:
  1503.             a = bytearray([0] * vpi)
  1504.         source_offset = 0
  1505.  
  1506.         for lines in adam7_generate(self.width, self.height):
  1507.             # The previous (reconstructed) scanline.
  1508.             # `None` at the beginning of a pass
  1509.             # to indicate that there is no previous line.
  1510.             recon = None
  1511.             for x, y, xstep in lines:
  1512.                 # Pixels per row (reduced pass image)
  1513.                 ppr = int(math.ceil((self.width - x) / float(xstep)))
  1514.                 # Row size in bytes for this pass.
  1515.                 row_size = int(math.ceil(self.psize * ppr))
  1516.  
  1517.                 filter_type = raw[source_offset]
  1518.                 source_offset += 1
  1519.                 scanline = raw[source_offset: source_offset + row_size]
  1520.                 source_offset += row_size
  1521.                 recon = self.undo_filter(filter_type, scanline, recon)
  1522.                 # Convert so that there is one element per pixel value
  1523.                 flat = self._bytes_to_values(recon, width=ppr)
  1524.                 if xstep == 1:
  1525.                     assert x == 0
  1526.                     offset = y * vpr
  1527.                     a[offset: offset + vpr] = flat
  1528.                 else:
  1529.                     offset = y * vpr + x * self.planes
  1530.                     end_offset = (y + 1) * vpr
  1531.                     skip = self.planes * xstep
  1532.                     for i in range(self.planes):
  1533.                         a[offset + i: end_offset: skip] = \
  1534.                             flat[i:: self.planes]
  1535.  
  1536.         return a
  1537.  
  1538.     def _iter_bytes_to_values(self, byte_rows):
  1539.         """
  1540.        Iterator that yields each scanline;
  1541.        each scanline being a sequence of values.
  1542.        `byte_rows` should be an iterator that yields
  1543.        the bytes of each row in turn.
  1544.        """
  1545.  
  1546.         for row in byte_rows:
  1547.             yield self._bytes_to_values(row)
  1548.  
  1549.     def _bytes_to_values(self, bs, width=None):
  1550.         """Convert a packed row of bytes into a row of values.
  1551.        Result will be a freshly allocated object,
  1552.        not shared with the argument.
  1553.        """
  1554.  
  1555.         if self.bitdepth == 8:
  1556.             return bytearray(bs)
  1557.         if self.bitdepth == 16:
  1558.             return array('H',
  1559.                          struct.unpack('!%dH' % (len(bs) // 2), bs))
  1560.  
  1561.         assert self.bitdepth < 8
  1562.         if width is None:
  1563.             width = self.width
  1564.         # Samples per byte
  1565.         spb = 8 // self.bitdepth
  1566.         out = bytearray()
  1567.         mask = 2**self.bitdepth - 1
  1568.         shifts = [self.bitdepth * i
  1569.                   for i in reversed(list(range(spb)))]
  1570.         for o in bs:
  1571.             out.extend([mask & (o >> i) for i in shifts])
  1572.         return out[:width]
  1573.  
  1574.     def _iter_straight_packed(self, byte_blocks):
  1575.         """Iterator that undoes the effect of filtering;
  1576.        yields each row as a sequence of packed bytes.
  1577.        Assumes input is straightlaced.
  1578.        `byte_blocks` should be an iterable that yields the raw bytes
  1579.        in blocks of arbitrary size.
  1580.        """
  1581.  
  1582.         # length of row, in bytes
  1583.         rb = self.row_bytes
  1584.         a = bytearray()
  1585.         # The previous (reconstructed) scanline.
  1586.         # None indicates first line of image.
  1587.         recon = None
  1588.         for some_bytes in byte_blocks:
  1589.             a.extend(some_bytes)
  1590.             while len(a) >= rb + 1:
  1591.                 filter_type = a[0]
  1592.                 scanline = a[1: rb + 1]
  1593.                 del a[: rb + 1]
  1594.                 recon = self.undo_filter(filter_type, scanline, recon)
  1595.                 yield recon
  1596.         if len(a) != 0:
  1597.             # :file:format We get here with a file format error:
  1598.             # when the available bytes (after decompressing) do not
  1599.             # pack into exact rows.
  1600.             raise FormatError('Wrong size for decompressed IDAT chunk.')
  1601.         assert len(a) == 0
  1602.  
  1603.     def validate_signature(self):
  1604.         """
  1605.        If signature (header) has not been read then read and
  1606.        validate it; otherwise do nothing.
  1607.        """
  1608.  
  1609.         if self.signature:
  1610.             return
  1611.         self.signature = self.file.read(8)
  1612.         if self.signature != signature:
  1613.             raise FormatError("PNG file has invalid signature.")
  1614.  
  1615.     def preamble(self, lenient=False):
  1616.         """
  1617.        Extract the image metadata by reading
  1618.        the initial part of the PNG file up to
  1619.        the start of the ``IDAT`` chunk.
  1620.        All the chunks that precede the ``IDAT`` chunk are
  1621.        read and either processed for metadata or discarded.
  1622.  
  1623.        If the optional `lenient` argument evaluates to `True`,
  1624.        checksum failures will raise warnings rather than exceptions.
  1625.        """
  1626.  
  1627.         self.validate_signature()
  1628.  
  1629.         while True:
  1630.             if not self.atchunk:
  1631.                 self.atchunk = self._chunk_len_type()
  1632.                 if self.atchunk is None:
  1633.                     raise FormatError('This PNG file has no IDAT chunks.')
  1634.             if self.atchunk[1] == b'IDAT':
  1635.                 return
  1636.             self.process_chunk(lenient=lenient)
  1637.  
  1638.     def _chunk_len_type(self):
  1639.         """
  1640.        Reads just enough of the input to
  1641.        determine the next chunk's length and type;
  1642.        return a (*length*, *type*) pair where *type* is a byte sequence.
  1643.        If there are no more chunks, ``None`` is returned.
  1644.        """
  1645.  
  1646.         x = self.file.read(8)
  1647.         if not x:
  1648.             return None
  1649.         if len(x) != 8:
  1650.             raise FormatError(
  1651.                 'End of file whilst reading chunk length and type.')
  1652.         length, type = struct.unpack('!I4s', x)
  1653.         if length > 2 ** 31 - 1:
  1654.             raise FormatError('Chunk %s is too large: %d.' % (type, length))
  1655.         # Check that all bytes are in valid ASCII range.
  1656.         # https://www.w3.org/TR/2003/REC-PNG-20031110/#5Chunk-layout
  1657.         type_bytes = set(bytearray(type))
  1658.         if not(type_bytes <= set(range(65, 91)) | set(range(97, 123))):
  1659.             raise FormatError(
  1660.                 'Chunk %r has invalid Chunk Type.'
  1661.                 % list(type))
  1662.         return length, type
  1663.  
  1664.     def process_chunk(self, lenient=False):
  1665.         """
  1666.        Process the next chunk and its data.
  1667.        This only processes the following chunk types:
  1668.        ``IHDR``, ``PLTE``, ``bKGD``, ``tRNS``, ``gAMA``, ``sBIT``, ``pHYs``.
  1669.        All other chunk types are ignored.
  1670.  
  1671.        If the optional `lenient` argument evaluates to `True`,
  1672.        checksum failures will raise warnings rather than exceptions.
  1673.        """
  1674.  
  1675.         type, data = self.chunk(lenient=lenient)
  1676.         method = '_process_' + type.decode('ascii')
  1677.         m = getattr(self, method, None)
  1678.         if m:
  1679.             m(data)
  1680.  
  1681.     def _process_IHDR(self, data):
  1682.         # http://www.w3.org/TR/PNG/#11IHDR
  1683.         if len(data) != 13:
  1684.             raise FormatError('IHDR chunk has incorrect length.')
  1685.         (self.width, self.height, self.bitdepth, self.color_type,
  1686.          self.compression, self.filter,
  1687.          self.interlace) = struct.unpack("!2I5B", data)
  1688.  
  1689.         check_bitdepth_colortype(self.bitdepth, self.color_type)
  1690.  
  1691.         if self.compression != 0:
  1692.             raise FormatError(
  1693.                 "Unknown compression method %d" % self.compression)
  1694.         if self.filter != 0:
  1695.             raise FormatError(
  1696.                 "Unknown filter method %d,"
  1697.                 " see http://www.w3.org/TR/2003/REC-PNG-20031110/#9Filters ."
  1698.                 % self.filter)
  1699.         if self.interlace not in (0, 1):
  1700.             raise FormatError(
  1701.                 "Unknown interlace method %d, see "
  1702.                 "http://www.w3.org/TR/2003/REC-PNG-20031110/#8InterlaceMethods"
  1703.                 " ."
  1704.                 % self.interlace)
  1705.  
  1706.         # Derived values
  1707.         # http://www.w3.org/TR/PNG/#6Colour-values
  1708.         colormap = bool(self.color_type & 1)
  1709.         greyscale = not(self.color_type & 2)
  1710.         alpha = bool(self.color_type & 4)
  1711.         color_planes = (3, 1)[greyscale or colormap]
  1712.         planes = color_planes + alpha
  1713.  
  1714.         self.colormap = colormap
  1715.         self.greyscale = greyscale
  1716.         self.alpha = alpha
  1717.         self.color_planes = color_planes
  1718.         self.planes = planes
  1719.         self.psize = float(self.bitdepth) / float(8) * planes
  1720.         if int(self.psize) == self.psize:
  1721.             self.psize = int(self.psize)
  1722.         self.row_bytes = int(math.ceil(self.width * self.psize))
  1723.         # Stores PLTE chunk if present, and is used to check
  1724.         # chunk ordering constraints.
  1725.         self.plte = None
  1726.         # Stores tRNS chunk if present, and is used to check chunk
  1727.         # ordering constraints.
  1728.         self.trns = None
  1729.         # Stores sBIT chunk if present.
  1730.         self.sbit = None
  1731.  
  1732.     def _process_PLTE(self, data):
  1733.         # http://www.w3.org/TR/PNG/#11PLTE
  1734.         if self.plte:
  1735.             warnings.warn("Multiple PLTE chunks present.")
  1736.         self.plte = data
  1737.         if len(data) % 3 != 0:
  1738.             raise FormatError(
  1739.                 "PLTE chunk's length should be a multiple of 3.")
  1740.         if len(data) > (2 ** self.bitdepth) * 3:
  1741.             raise FormatError("PLTE chunk is too long.")
  1742.         if len(data) == 0:
  1743.             raise FormatError("Empty PLTE is not allowed.")
  1744.  
  1745.     def _process_bKGD(self, data):
  1746.         try:
  1747.             if self.colormap:
  1748.                 if not self.plte:
  1749.                     warnings.warn(
  1750.                         "PLTE chunk is required before bKGD chunk.")
  1751.                 self.background = struct.unpack('B', data)
  1752.             else:
  1753.                 self.background = struct.unpack("!%dH" % self.color_planes,
  1754.                                                 data)
  1755.         except struct.error:
  1756.             raise FormatError("bKGD chunk has incorrect length.")
  1757.  
  1758.     def _process_tRNS(self, data):
  1759.         # http://www.w3.org/TR/PNG/#11tRNS
  1760.         self.trns = data
  1761.         if self.colormap:
  1762.             if not self.plte:
  1763.                 warnings.warn("PLTE chunk is required before tRNS chunk.")
  1764.             else:
  1765.                 if len(data) > len(self.plte) / 3:
  1766.                     # Was warning, but promoted to Error as it
  1767.                     # would otherwise cause pain later on.
  1768.                     raise FormatError("tRNS chunk is too long.")
  1769.         else:
  1770.             if self.alpha:
  1771.                 raise FormatError(
  1772.                     "tRNS chunk is not valid with colour type %d." %
  1773.                     self.color_type)
  1774.             try:
  1775.                 self.transparent = \
  1776.                     struct.unpack("!%dH" % self.color_planes, data)
  1777.             except struct.error:
  1778.                 raise FormatError("tRNS chunk has incorrect length.")
  1779.  
  1780.     def _process_gAMA(self, data):
  1781.         try:
  1782.             self.gamma = struct.unpack("!L", data)[0] / 100000.0
  1783.         except struct.error:
  1784.             raise FormatError("gAMA chunk has incorrect length.")
  1785.  
  1786.     def _process_sBIT(self, data):
  1787.         self.sbit = data
  1788.         if (self.colormap and len(data) != 3 or
  1789.                 not self.colormap and len(data) != self.planes):
  1790.             raise FormatError("sBIT chunk has incorrect length.")
  1791.  
  1792.     def _process_pHYs(self, data):
  1793.         # http://www.w3.org/TR/PNG/#11pHYs
  1794.         self.phys = data
  1795.         fmt = "!LLB"
  1796.         if len(data) != struct.calcsize(fmt):
  1797.             raise FormatError("pHYs chunk has incorrect length.")
  1798.         self.x_pixels_per_unit, self.y_pixels_per_unit, unit = \
  1799.             struct.unpack(fmt, data)
  1800.         self.unit_is_meter = bool(unit)
  1801.  
  1802.     def read(self, lenient=False):
  1803.         """
  1804.        Read the PNG file and decode it.
  1805.        Returns (`width`, `height`, `rows`, `info`).
  1806.  
  1807.        May use excessive memory.
  1808.  
  1809.        `rows` is a sequence of rows;
  1810.        each row is a sequence of values.
  1811.  
  1812.        If the optional `lenient` argument evaluates to True,
  1813.        checksum failures will raise warnings rather than exceptions.
  1814.        """
  1815.  
  1816.         def iteridat():
  1817.             """Iterator that yields all the ``IDAT`` chunks as strings."""
  1818.             while True:
  1819.                 type, data = self.chunk(lenient=lenient)
  1820.                 if type == b'IEND':
  1821.                     # http://www.w3.org/TR/PNG/#11IEND
  1822.                     break
  1823.                 if type != b'IDAT':
  1824.                     continue
  1825.                 # type == b'IDAT'
  1826.                 # http://www.w3.org/TR/PNG/#11IDAT
  1827.                 if self.colormap and not self.plte:
  1828.                     warnings.warn("PLTE chunk is required before IDAT chunk")
  1829.                 yield data
  1830.  
  1831.         self.preamble(lenient=lenient)
  1832.         raw = decompress(iteridat())
  1833.  
  1834.         if self.interlace:
  1835.             def rows_from_interlace():
  1836.                 """Yield each row from an interlaced PNG."""
  1837.                 # It's important that this iterator doesn't read
  1838.                 # IDAT chunks until it yields the first row.
  1839.                 bs = bytearray(itertools.chain(*raw))
  1840.                 arraycode = 'BH'[self.bitdepth > 8]
  1841.                 # Like :meth:`group` but
  1842.                 # producing an array.array object for each row.
  1843.                 values = self._deinterlace(bs)
  1844.                 vpr = self.width * self.planes
  1845.                 for i in range(0, len(values), vpr):
  1846.                     row = array(arraycode, values[i:i+vpr])
  1847.                     yield row
  1848.             rows = rows_from_interlace()
  1849.         else:
  1850.             rows = self._iter_bytes_to_values(self._iter_straight_packed(raw))
  1851.         info = dict()
  1852.         for attr in 'greyscale alpha planes bitdepth interlace'.split():
  1853.             info[attr] = getattr(self, attr)
  1854.         info['size'] = (self.width, self.height)
  1855.         for attr in 'gamma transparent background'.split():
  1856.             a = getattr(self, attr, None)
  1857.             if a is not None:
  1858.                 info[attr] = a
  1859.         if getattr(self, 'x_pixels_per_unit', None):
  1860.             info['physical'] = Resolution(self.x_pixels_per_unit,
  1861.                                           self.y_pixels_per_unit,
  1862.                                           self.unit_is_meter)
  1863.         if self.plte:
  1864.             info['palette'] = self.palette()
  1865.         return self.width, self.height, rows, info
  1866.  
  1867.     def read_flat(self):
  1868.         """
  1869.        Read a PNG file and decode it into a single array of values.
  1870.        Returns (*width*, *height*, *values*, *info*).
  1871.  
  1872.        May use excessive memory.
  1873.  
  1874.        `values` is a single array.
  1875.  
  1876.        The :meth:`read` method is more stream-friendly than this,
  1877.        because it returns a sequence of rows.
  1878.        """
  1879.  
  1880.         x, y, pixel, info = self.read()
  1881.         arraycode = 'BH'[info['bitdepth'] > 8]
  1882.         pixel = array(arraycode, itertools.chain(*pixel))
  1883.         return x, y, pixel, info
  1884.  
  1885.     def palette(self, alpha='natural'):
  1886.         """
  1887.        Returns a palette that is a sequence of 3-tuples or 4-tuples,
  1888.        synthesizing it from the ``PLTE`` and ``tRNS`` chunks.
  1889.        These chunks should have already been processed (for example,
  1890.        by calling the :meth:`preamble` method).
  1891.        All the tuples are the same size:
  1892.        3-tuples if there is no ``tRNS`` chunk,
  1893.        4-tuples when there is a ``tRNS`` chunk.
  1894.  
  1895.        Assumes that the image is colour type
  1896.        3 and therefore a ``PLTE`` chunk is required.
  1897.  
  1898.        If the `alpha` argument is ``'force'`` then an alpha channel is
  1899.        always added, forcing the result to be a sequence of 4-tuples.
  1900.        """
  1901.  
  1902.         if not self.plte:
  1903.             raise FormatError(
  1904.                 "Required PLTE chunk is missing in colour type 3 image.")
  1905.         plte = group(array('B', self.plte), 3)
  1906.         if self.trns or alpha == 'force':
  1907.             trns = array('B', self.trns or [])
  1908.             trns.extend([255] * (len(plte) - len(trns)))
  1909.             plte = list(map(operator.add, plte, group(trns, 1)))
  1910.         return plte
  1911.  
  1912.     def asDirect(self):
  1913.         """
  1914.        Returns the image data as a direct representation of
  1915.        an ``x * y * planes`` array.
  1916.        This removes the need for callers to deal with
  1917.        palettes and transparency themselves.
  1918.        Images with a palette (colour type 3) are converted to RGB or RGBA;
  1919.        images with transparency (a ``tRNS`` chunk) are converted to
  1920.        LA or RGBA as appropriate.
  1921.        When returned in this format the pixel values represent
  1922.        the colour value directly without needing to refer
  1923.        to palettes or transparency information.
  1924.  
  1925.        Like the :meth:`read` method this method returns a 4-tuple:
  1926.  
  1927.        (*width*, *height*, *rows*, *info*)
  1928.  
  1929.        This method normally returns pixel values with
  1930.        the bit depth they have in the source image, but
  1931.        when the source PNG has an ``sBIT`` chunk it is inspected and
  1932.        can reduce the bit depth of the result pixels;
  1933.        pixel values will be reduced according to the bit depth
  1934.        specified in the ``sBIT`` chunk.
  1935.        PNG nerds should note a single result bit depth is
  1936.        used for all channels:
  1937.        the maximum of the ones specified in the ``sBIT`` chunk.
  1938.        An RGB565 image will be rescaled to 6-bit RGB666.
  1939.  
  1940.        The *info* dictionary that is returned reflects
  1941.        the `direct` format and not the original source image.
  1942.        For example, an RGB source image with a ``tRNS`` chunk
  1943.        to represent a transparent colour,
  1944.        will start with ``planes=3`` and ``alpha=False`` for the
  1945.        source image,
  1946.        but the *info* dictionary returned by this method
  1947.        will have ``planes=4`` and ``alpha=True`` because
  1948.        an alpha channel is synthesized and added.
  1949.  
  1950.        *rows* is a sequence of rows;
  1951.        each row being a sequence of values
  1952.        (like the :meth:`read` method).
  1953.  
  1954.        All the other aspects of the image data are not changed.
  1955.        """
  1956.  
  1957.         self.preamble()
  1958.  
  1959.         # Simple case, no conversion necessary.
  1960.         if not self.colormap and not self.trns and not self.sbit:
  1961.             return self.read()
  1962.  
  1963.         x, y, pixels, info = self.read()
  1964.  
  1965.         if self.colormap:
  1966.             info['colormap'] = False
  1967.             info['alpha'] = bool(self.trns)
  1968.             info['bitdepth'] = 8
  1969.             info['planes'] = 3 + bool(self.trns)
  1970.             plte = self.palette()
  1971.  
  1972.             def iterpal(pixels):
  1973.                 for row in pixels:
  1974.                     row = [plte[x] for x in row]
  1975.                     yield array('B', itertools.chain(*row))
  1976.             pixels = iterpal(pixels)
  1977.         elif self.trns:
  1978.             # It would be nice if there was some reasonable way
  1979.             # of doing this without generating a whole load of
  1980.             # intermediate tuples.  But tuples does seem like the
  1981.             # easiest way, with no other way clearly much simpler or
  1982.             # much faster.  (Actually, the L to LA conversion could
  1983.             # perhaps go faster (all those 1-tuples!), but I still
  1984.             # wonder whether the code proliferation is worth it)
  1985.             it = self.transparent
  1986.             maxval = 2 ** info['bitdepth'] - 1
  1987.             planes = info['planes']
  1988.             info['alpha'] = True
  1989.             info['planes'] += 1
  1990.             typecode = 'BH'[info['bitdepth'] > 8]
  1991.  
  1992.             def itertrns(pixels):
  1993.                 for row in pixels:
  1994.                     # For each row we group it into pixels, then form a
  1995.                     # characterisation vector that says whether each
  1996.                     # pixel is opaque or not.  Then we convert
  1997.                     # True/False to 0/maxval (by multiplication),
  1998.                     # and add it as the extra channel.
  1999.                     row = group(row, planes)
  2000.                     opa = map(it.__ne__, row)
  2001.                     opa = map(maxval.__mul__, opa)
  2002.                     opa = list(zip(opa))    # convert to 1-tuples
  2003.                     yield array(
  2004.                         typecode,
  2005.                         itertools.chain(*map(operator.add, row, opa)))
  2006.             pixels = itertrns(pixels)
  2007.         targetbitdepth = None
  2008.         if self.sbit:
  2009.             sbit = struct.unpack('%dB' % len(self.sbit), self.sbit)
  2010.             targetbitdepth = max(sbit)
  2011.             if targetbitdepth > info['bitdepth']:
  2012.                 raise Error('sBIT chunk %r exceeds bitdepth %d' %
  2013.                             (sbit, self.bitdepth))
  2014.             if min(sbit) <= 0:
  2015.                 raise Error('sBIT chunk %r has a 0-entry' % sbit)
  2016.         if targetbitdepth:
  2017.             shift = info['bitdepth'] - targetbitdepth
  2018.             info['bitdepth'] = targetbitdepth
  2019.  
  2020.             def itershift(pixels):
  2021.                 for row in pixels:
  2022.                     yield [p >> shift for p in row]
  2023.             pixels = itershift(pixels)
  2024.         return x, y, pixels, info
  2025.  
  2026.     def _as_rescale(self, get, targetbitdepth):
  2027.         """Helper used by :meth:`asRGB8` and :meth:`asRGBA8`."""
  2028.  
  2029.         width, height, pixels, info = get()
  2030.         maxval = 2**info['bitdepth'] - 1
  2031.         targetmaxval = 2**targetbitdepth - 1
  2032.         factor = float(targetmaxval) / float(maxval)
  2033.         info['bitdepth'] = targetbitdepth
  2034.  
  2035.         def iterscale():
  2036.             for row in pixels:
  2037.                 yield [int(round(x * factor)) for x in row]
  2038.         if maxval == targetmaxval:
  2039.             return width, height, pixels, info
  2040.         else:
  2041.             return width, height, iterscale(), info
  2042.  
  2043.     def asRGB8(self):
  2044.         """
  2045.        Return the image data as an RGB pixels with 8-bits per sample.
  2046.        This is like the :meth:`asRGB` method except that
  2047.        this method additionally rescales the values so that
  2048.        they are all between 0 and 255 (8-bit).
  2049.        In the case where the source image has a bit depth < 8
  2050.        the transformation preserves all the information;
  2051.        where the source image has bit depth > 8, then
  2052.        rescaling to 8-bit values loses precision.
  2053.        No dithering is performed.
  2054.        Like :meth:`asRGB`,
  2055.        an alpha channel in the source image will raise an exception.
  2056.  
  2057.        This function returns a 4-tuple:
  2058.        (*width*, *height*, *rows*, *info*).
  2059.        *width*, *height*, *info* are as per the :meth:`read` method.
  2060.  
  2061.        *rows* is the pixel data as a sequence of rows.
  2062.        """
  2063.  
  2064.         return self._as_rescale(self.asRGB, 8)
  2065.  
  2066.     def asRGBA8(self):
  2067.         """
  2068.        Return the image data as RGBA pixels with 8-bits per sample.
  2069.        This method is similar to :meth:`asRGB8` and :meth:`asRGBA`:
  2070.        The result pixels have an alpha channel, *and*
  2071.        values are rescaled to the range 0 to 255.
  2072.        The alpha channel is synthesized if necessary
  2073.        (with a small speed penalty).
  2074.        """
  2075.  
  2076.         return self._as_rescale(self.asRGBA, 8)
  2077.  
  2078.     def asRGB(self):
  2079.         """
  2080.        Return image as RGB pixels.
  2081.        RGB colour images are passed through unchanged;
  2082.        greyscales are expanded into RGB triplets
  2083.        (there is a small speed overhead for doing this).
  2084.  
  2085.        An alpha channel in the source image will raise an exception.
  2086.  
  2087.        The return values are as for the :meth:`read` method except that
  2088.        the *info* reflect the returned pixels, not the source image.
  2089.        In particular,
  2090.        for this method ``info['greyscale']`` will be ``False``.
  2091.        """
  2092.  
  2093.         width, height, pixels, info = self.asDirect()
  2094.         if info['alpha']:
  2095.             raise Error("will not convert image with alpha channel to RGB")
  2096.         if not info['greyscale']:
  2097.             return width, height, pixels, info
  2098.         info['greyscale'] = False
  2099.         info['planes'] = 3
  2100.  
  2101.         if info['bitdepth'] > 8:
  2102.             def newarray():
  2103.                 return array('H', [0])
  2104.         else:
  2105.             def newarray():
  2106.                 return bytearray([0])
  2107.  
  2108.         def iterrgb():
  2109.             for row in pixels:
  2110.                 a = newarray() * 3 * width
  2111.                 for i in range(3):
  2112.                     a[i::3] = row
  2113.                 yield a
  2114.         return width, height, iterrgb(), info
  2115.  
  2116.     def asRGBA(self):
  2117.         """
  2118.        Return image as RGBA pixels.
  2119.        Greyscales are expanded into RGB triplets;
  2120.        an alpha channel is synthesized if necessary.
  2121.        The return values are as for the :meth:`read` method except that
  2122.        the *info* reflect the returned pixels, not the source image.
  2123.        In particular, for this method
  2124.        ``info['greyscale']`` will be ``False``, and
  2125.        ``info['alpha']`` will be ``True``.
  2126.        """
  2127.  
  2128.         width, height, pixels, info = self.asDirect()
  2129.         if info['alpha'] and not info['greyscale']:
  2130.             return width, height, pixels, info
  2131.         typecode = 'BH'[info['bitdepth'] > 8]
  2132.         maxval = 2**info['bitdepth'] - 1
  2133.         maxbuffer = struct.pack('=' + typecode, maxval) * 4 * width
  2134.  
  2135.         if info['bitdepth'] > 8:
  2136.             def newarray():
  2137.                 return array('H', maxbuffer)
  2138.         else:
  2139.             def newarray():
  2140.                 return bytearray(maxbuffer)
  2141.  
  2142.         if info['alpha'] and info['greyscale']:
  2143.             # LA to RGBA
  2144.             def convert():
  2145.                 for row in pixels:
  2146.                     # Create a fresh target row, then copy L channel
  2147.                     # into first three target channels, and A channel
  2148.                     # into fourth channel.
  2149.                     a = newarray()
  2150.                     convert_la_to_rgba(row, a)
  2151.                     yield a
  2152.         elif info['greyscale']:
  2153.             # L to RGBA
  2154.             def convert():
  2155.                 for row in pixels:
  2156.                     a = newarray()
  2157.                     convert_l_to_rgba(row, a)
  2158.                     yield a
  2159.         else:
  2160.             assert not info['alpha'] and not info['greyscale']
  2161.             # RGB to RGBA
  2162.  
  2163.             def convert():
  2164.                 for row in pixels:
  2165.                     a = newarray()
  2166.                     convert_rgb_to_rgba(row, a)
  2167.                     yield a
  2168.         info['alpha'] = True
  2169.         info['greyscale'] = False
  2170.         info['planes'] = 4
  2171.         return width, height, convert(), info
  2172.  
  2173.  
  2174. def decompress(data_blocks):
  2175.     """
  2176.    `data_blocks` should be an iterable that
  2177.    yields the compressed data (from the ``IDAT`` chunks).
  2178.    This yields decompressed byte strings.
  2179.    """
  2180.  
  2181.     # Currently, with no max_length parameter to decompress,
  2182.     # this routine will do one yield per IDAT chunk: Not very
  2183.     # incremental.
  2184.     d = zlib.decompressobj()
  2185.     # Each IDAT chunk is passed to the decompressor, then any
  2186.     # remaining state is decompressed out.
  2187.     for data in data_blocks:
  2188.         # :todo: add a max_length argument here to limit output size.
  2189.         yield bytearray(d.decompress(data))
  2190.     yield bytearray(d.flush())
  2191.  
  2192.  
  2193. def check_bitdepth_colortype(bitdepth, colortype):
  2194.     """
  2195.    Check that `bitdepth` and `colortype` are both valid,
  2196.    and specified in a valid combination.
  2197.    Returns (None) if valid, raise an Exception if not valid.
  2198.    """
  2199.  
  2200.     if bitdepth not in (1, 2, 4, 8, 16):
  2201.         raise FormatError("invalid bit depth %d" % bitdepth)
  2202.     if colortype not in (0, 2, 3, 4, 6):
  2203.         raise FormatError("invalid colour type %d" % colortype)
  2204.     # Check indexed (palettized) images have 8 or fewer bits
  2205.     # per pixel; check only indexed or greyscale images have
  2206.     # fewer than 8 bits per pixel.
  2207.     if colortype & 1 and bitdepth > 8:
  2208.         raise FormatError(
  2209.             "Indexed images (colour type %d) cannot"
  2210.             " have bitdepth > 8 (bit depth %d)."
  2211.             " See http://www.w3.org/TR/2003/REC-PNG-20031110/#table111 ."
  2212.             % (bitdepth, colortype))
  2213.     if bitdepth < 8 and colortype not in (0, 3):
  2214.         raise FormatError(
  2215.             "Illegal combination of bit depth (%d)"
  2216.             " and colour type (%d)."
  2217.             " See http://www.w3.org/TR/2003/REC-PNG-20031110/#table111 ."
  2218.             % (bitdepth, colortype))
  2219.  
  2220.  
  2221. def is_natural(x):
  2222.     """A non-negative integer."""
  2223.     try:
  2224.         is_integer = int(x) == x
  2225.     except (TypeError, ValueError):
  2226.         return False
  2227.     return is_integer and x >= 0
  2228.  
  2229.  
  2230. def undo_filter_sub(filter_unit, scanline, previous, result):
  2231.     """Undo sub filter."""
  2232.  
  2233.     ai = 0
  2234.     # Loops starts at index fu.  Observe that the initial part
  2235.     # of the result is already filled in correctly with
  2236.     # scanline.
  2237.     for i in range(filter_unit, len(result)):
  2238.         x = scanline[i]
  2239.         a = result[ai]
  2240.         result[i] = (x + a) & 0xff
  2241.         ai += 1
  2242.  
  2243.  
  2244. def undo_filter_up(filter_unit, scanline, previous, result):
  2245.     """Undo up filter."""
  2246.  
  2247.     for i in range(len(result)):
  2248.         x = scanline[i]
  2249.         b = previous[i]
  2250.         result[i] = (x + b) & 0xff
  2251.  
  2252.  
  2253. def undo_filter_average(filter_unit, scanline, previous, result):
  2254.     """Undo up filter."""
  2255.  
  2256.     ai = -filter_unit
  2257.     for i in range(len(result)):
  2258.         x = scanline[i]
  2259.         if ai < 0:
  2260.             a = 0
  2261.         else:
  2262.             a = result[ai]
  2263.         b = previous[i]
  2264.         result[i] = (x + ((a + b) >> 1)) & 0xff
  2265.         ai += 1
  2266.  
  2267.  
  2268. def undo_filter_paeth(filter_unit, scanline, previous, result):
  2269.     """Undo Paeth filter."""
  2270.  
  2271.     # Also used for ci.
  2272.     ai = -filter_unit
  2273.     for i in range(len(result)):
  2274.         x = scanline[i]
  2275.         if ai < 0:
  2276.             a = c = 0
  2277.         else:
  2278.             a = result[ai]
  2279.             c = previous[ai]
  2280.         b = previous[i]
  2281.         p = a + b - c
  2282.         pa = abs(p - a)
  2283.         pb = abs(p - b)
  2284.         pc = abs(p - c)
  2285.         if pa <= pb and pa <= pc:
  2286.             pr = a
  2287.         elif pb <= pc:
  2288.             pr = b
  2289.         else:
  2290.             pr = c
  2291.         result[i] = (x + pr) & 0xff
  2292.         ai += 1
  2293.  
  2294.  
  2295. def convert_la_to_rgba(row, result):
  2296.     for i in range(3):
  2297.         result[i::4] = row[0::2]
  2298.     result[3::4] = row[1::2]
  2299.  
  2300.  
  2301. def convert_l_to_rgba(row, result):
  2302.     """
  2303.    Convert a grayscale image to RGBA.
  2304.    This method assumes the alpha channel in result is
  2305.    already correctly initialized.
  2306.    """
  2307.     for i in range(3):
  2308.         result[i::4] = row
  2309.  
  2310.  
  2311. def convert_rgb_to_rgba(row, result):
  2312.     """
  2313.    Convert an RGB image to RGBA.
  2314.    This method assumes the alpha channel in result is
  2315.    already correctly initialized.
  2316.    """
  2317.     for i in range(3):
  2318.         result[i::4] = row[i::3]
  2319.  
  2320.  
  2321. # -------------------------------------------------------------------------------------------
  2322. # -------------------------------------------------------------------------------------------
  2323. # -------------------------------------------------------------------------------------------
  2324. # -------------------------------------------------------------------------------------------
  2325. # -------------------------------------------------------------------------------------------
  2326. # -------------------------------------------------------------------------------------------
  2327. # -------------------------------------------------------------------------------------------
  2328. # -------------------------------------------------------------------------------------------
  2329. # -------------------------------------------------------------------------------------------
  2330. """
  2331. Pokemon Crystal data de/compression.
  2332. """
  2333.  
  2334. """
  2335. A rundown of Pokemon Crystal's compression scheme:
  2336.  
  2337. Control commands occupy bits 5-7.
  2338. Bits 0-4 serve as the first parameter <n> for each command.
  2339. """
  2340. lz_commands = {
  2341.     'literal':   0, # n values for n bytes
  2342.     'iterate':   1, # one value for n bytes
  2343.     'alternate': 2, # alternate two values for n bytes
  2344.     'blank':     3, # zero for n bytes
  2345. }
  2346.  
  2347. """
  2348. Repeater commands repeat any data that was just decompressed.
  2349. They take an additional signed parameter <s> to mark a relative starting point.
  2350. These wrap around (positive from the start, negative from the current position).
  2351. """
  2352. lz_commands.update({
  2353.     'repeat':    4, # n bytes starting from s
  2354.     'flip':      5, # n bytes in reverse bit order starting from s
  2355.     'reverse':   6, # n bytes backwards starting from s
  2356. })
  2357.  
  2358. """
  2359. The long command is used when 5 bits aren't enough. Bits 2-4 contain a new control code.
  2360. Bits 0-1 are appended to a new byte as 8-9, allowing a 10-bit parameter.
  2361. """
  2362. lz_commands.update({
  2363.     'long':      7, # n is now 10 bits for a new control code
  2364. })
  2365. max_length = 1 << 10 # can't go higher than 10 bits
  2366. lowmax     = 1 <<  5 # standard 5-bit param
  2367.  
  2368. """
  2369. If 0xff is encountered instead of a command, decompression ends.
  2370. """
  2371. lz_end = 0xff
  2372.  
  2373.  
  2374. bit_flipped = [
  2375.     sum(((byte >> i) & 1) << (7 - i) for i in xrange(8))
  2376.     for byte in xrange(0x100)
  2377. ]
  2378.  
  2379.  
  2380. def fbitstream(f):
  2381.     while 1:
  2382.         char = f.read(1)
  2383.         if not char:
  2384.             break
  2385.         byte = ord(char)
  2386.         for i in range(7, -1, -1):
  2387.             yield (byte >> i) & 1
  2388.  
  2389. def bitstream(b):
  2390.     for byte in b:
  2391.         for i in range(7, -1, -1):
  2392.             yield (byte >> i) & 1
  2393.  
  2394. def readint(bs, count):
  2395.     n = 0
  2396.     while count:
  2397.         n <<= 1
  2398.         n |= next(bs)
  2399.         count -= 1
  2400.     return n
  2401.  
  2402. def bitgroups_to_bytes(bits):
  2403.     l = []
  2404.     for i in range(0, len(bits) - 3, 4):
  2405.         n = ((bits[i + 0] << 6)
  2406.            | (bits[i + 1] << 4)
  2407.            | (bits[i + 2] << 2)
  2408.            | (bits[i + 3] << 0))
  2409.         l.append(n)
  2410.     return bytearray(l)
  2411.  
  2412.  
  2413. def bytes_to_bits(bytelist):
  2414.     return list(bitstream(bytelist))
  2415.  
  2416. class Compressed:
  2417.  
  2418.     """
  2419.     Usage:
  2420.         lz = Compressed(data).output
  2421.     or
  2422.         lz = Compressed().compress(data)
  2423.     or
  2424.         c = Compressed()
  2425.         c.data = data
  2426.         lz = c.compress()
  2427.  
  2428.     There are some issues with reproducing the target compressor.
  2429.     Some notes are listed here:
  2430.         - the criteria for detecting a lookback is inconsistent
  2431.             - sometimes lookbacks that are mostly 0s are pruned, sometimes not
  2432.         - target appears to skip ahead if it can use a lookback soon, stopping the current command short or in some cases truncating it with literals.
  2433.             - this has been implemented, but the specifics are unknown
  2434.         - self.min_scores: It's unknown if blank's minimum score should be 1 or 2. Most likely it's 1, with some other hack to account for edge cases.
  2435.             - may be related to the above
  2436.         - target does not appear to compress backwards
  2437.     """
  2438.  
  2439.     def __init__(self, *args, **kwargs):
  2440.  
  2441.         self.min_scores = {
  2442.             'blank':     1,
  2443.             'iterate':   2,
  2444.             'alternate': 3,
  2445.             'repeat':    3,
  2446.             'reverse':   3,
  2447.             'flip':      3,
  2448.         }
  2449.  
  2450.         self.preference = [
  2451.             'repeat',
  2452.             'blank',
  2453.             'flip',
  2454.             'reverse',
  2455.             'iterate',
  2456.             'alternate',
  2457.             #'literal',
  2458.         ]
  2459.  
  2460.         self.lookback_methods = 'repeat', 'reverse', 'flip'
  2461.  
  2462.         self.__dict__.update({
  2463.             'data': None,
  2464.             'commands': lz_commands,
  2465.             'debug': False,
  2466.             'literal_only': False,
  2467.         })
  2468.  
  2469.         self.arg_names = 'data', 'commands', 'debug', 'literal_only'
  2470.  
  2471.         self.__dict__.update(kwargs)
  2472.         self.__dict__.update(dict(zip(self.arg_names, args)))
  2473.  
  2474.         if self.data is not None:
  2475.             self.compress()
  2476.  
  2477.     def compress(self, data=None):
  2478.         if data is not None:
  2479.             self.data = data
  2480.  
  2481.         self.data = list(bytearray(self.data))
  2482.  
  2483.         self.indexes = {}
  2484.         self.lookbacks = {}
  2485.         for method in self.lookback_methods:
  2486.             self.lookbacks[method] = {}
  2487.  
  2488.         self.address = 0
  2489.         self.end     = len(self.data)
  2490.         self.output  = []
  2491.         self.literal = None
  2492.  
  2493.         while self.address < self.end:
  2494.  
  2495.             if self.score():
  2496.                 self.do_literal()
  2497.                 self.do_winner()
  2498.  
  2499.             else:
  2500.                 if self.literal == None:
  2501.                     self.literal = self.address
  2502.                 self.address += 1
  2503.  
  2504.         self.do_literal()
  2505.  
  2506.         self.output += [lz_end]
  2507.         return self.output
  2508.    
  2509.     def reset_scores(self):
  2510.         self.scores = {}
  2511.         self.offsets = {}
  2512.         self.helpers = {}
  2513.         for method in self.min_scores.iterkeys():
  2514.             self.scores[method] = 0
  2515.    
  2516.     def bit_flip(self, byte):
  2517.         return bit_flipped[byte]
  2518.    
  2519.     def do_literal(self):
  2520.         if self.literal != None:
  2521.             length = abs(self.address - self.literal)
  2522.             start  = min(self.literal, self.address + 1)
  2523.             self.helpers['literal'] = self.data[start:start+length]
  2524.             self.do_cmd('literal', length)
  2525.             self.literal = None
  2526.    
  2527.     def score(self):
  2528.         self.reset_scores()
  2529.        
  2530.         map(self.score_literal, ['iterate', 'alternate', 'blank'])
  2531.        
  2532.         for method in self.lookback_methods:
  2533.             self.scores[method], self.offsets[method] = self.find_lookback(method, self.address)
  2534.        
  2535.         self.stop_short()
  2536.        
  2537.         return any(
  2538.             score
  2539.           > self.min_scores[method] + int(score > lowmax)
  2540.             for method, score in self.scores.iteritems()
  2541.         )
  2542.    
  2543.     def stop_short(self):
  2544.         """
  2545.         If a lookback is close, reduce the scores of other commands.
  2546.         """
  2547.         best_method, best_score = max(
  2548.             self.scores.items(),
  2549.             key = lambda x: (
  2550.                 x[1],
  2551.                 -self.preference.index(x[0])
  2552.             )
  2553.         )
  2554.         for method in self.lookback_methods:
  2555.             min_score = self.min_scores[method]
  2556.             for address in xrange(self.address+1, self.address+best_score):
  2557.                 length, index = self.find_lookback(method, address)
  2558.                 if length > max(min_score, best_score):
  2559.                     # BUG: lookbacks can reduce themselves. This appears to be a bug in the target also.
  2560.                     for m, score in self.scores.items():
  2561.                         self.scores[m] = min(score, address - self.address)
  2562.    
  2563.    
  2564.     def read(self, address=None):
  2565.         if address is None:
  2566.             address = self.address
  2567.         if 0 <= address < len(self.data):
  2568.             return self.data[address]
  2569.         return None
  2570.    
  2571.     def find_all_lookbacks(self):
  2572.         for method in self.lookback_methods:
  2573.             for address, byte in enumerate(self.data):
  2574.                 self.find_lookback(method, address)
  2575.    
  2576.     def find_lookback(self, method, address=None):
  2577.         #Temporarily stubbed, because the real function doesn't run in polynomial time.
  2578.         return (0, None)
  2579.  
  2580.     def broken_find_lookback(self, method, address=None):
  2581.         if address is None:
  2582.             address = self.address
  2583.  
  2584.         existing = self.lookbacks.get(method, {}).get(address)
  2585.         if existing != None:
  2586.             return existing
  2587.  
  2588.         lookback = 0, None
  2589.  
  2590.         # Better to not carelessly optimize at the moment.
  2591.         """
  2592.         if address < 2:
  2593.             return lookback
  2594.         """
  2595.  
  2596.         byte = self.read(address)
  2597.         if byte is None:
  2598.             return lookback
  2599.  
  2600.         direction, mutate = {
  2601.             'repeat':  ( 1, int),
  2602.             'reverse': (-1, int),
  2603.             'flip':    ( 1, self.bit_flip),
  2604.         }[method]
  2605.  
  2606.         # Doesn't seem to help
  2607.         """
  2608.         if mutate == self.bit_flip:
  2609.             if byte == 0:
  2610.                 self.lookbacks[method][address] = lookback
  2611.                 return lookback
  2612.         """
  2613.  
  2614.         data_len = len(self.data)
  2615.         is_two_byte_index = lambda index: int(index < address - 0x7f)
  2616.  
  2617.         for index in self.get_indexes(mutate(byte)):
  2618.  
  2619.             if index >= address:
  2620.                 break
  2621.  
  2622.             old_length, old_index = lookback
  2623.             if direction == 1:
  2624.                 if old_length > data_len - index: break
  2625.             else:
  2626.                 if old_length > index: continue
  2627.  
  2628.             if self.read(index) in [None]: continue
  2629.  
  2630.             length = 1 # we know there's at least one match, or we wouldn't be checking this index
  2631.             while 1:
  2632.                 this_byte = self.read(address + length)
  2633.                 that_byte = self.read(index + length * direction)
  2634.                 if that_byte == None or this_byte != mutate(that_byte):
  2635.                     break
  2636.                 length += 1
  2637.  
  2638.             score = length - is_two_byte_index(index)
  2639.             old_score = old_length - is_two_byte_index(old_index)
  2640.             if score >= old_score or (score == old_score and length > old_length):
  2641.                 # XXX maybe avoid two-byte indexes when possible
  2642.                 if score >= lookback[0] - is_two_byte_index(lookback[1]):
  2643.                     lookback = length, index
  2644.  
  2645.         self.lookbacks[method][address] = lookback
  2646.         return lookback
  2647.  
  2648.     def get_indexes(self, byte):
  2649.         if not self.indexes.has_key(byte):
  2650.             self.indexes[byte] = []
  2651.             index = -1
  2652.             while 1:
  2653.                 try:
  2654.                     index = self.data.index(byte, index + 1)
  2655.                 except ValueError:
  2656.                     break
  2657.                 self.indexes[byte].append(index)
  2658.         return self.indexes[byte]
  2659.  
  2660.     def score_literal(self, method):
  2661.         address = self.address
  2662.  
  2663.         compare = {
  2664.             'blank': [0],
  2665.             'iterate': [self.read(address)],
  2666.             'alternate': [self.read(address), self.read(address + 1)],
  2667.         }[method]
  2668.  
  2669.         # XXX may or may not be correct
  2670.         if method == 'alternate' and compare[0] == 0:
  2671.             return
  2672.  
  2673.         length = 0
  2674.         while self.read(address + length) == compare[length % len(compare)]:
  2675.             length += 1
  2676.  
  2677.         self.scores[method] = length
  2678.         self.helpers[method] = compare
  2679.  
  2680.     def do_winner(self):
  2681.         winners = filter(
  2682.             lambda method, score:
  2683.                 score
  2684.               > self.min_scores[method] + int(score > lowmax),
  2685.             self.scores.iteritems()
  2686.         )
  2687.         winners.sort(
  2688.             key = lambda method, score: (
  2689.                 -(score - self.min_scores[method] - int(score > lowmax)),
  2690.                 self.preference.index(method)
  2691.             )
  2692.         )
  2693.         winner, score = winners[0]
  2694.  
  2695.         length = min(score, max_length)
  2696.         self.do_cmd(winner, length)
  2697.         self.address += length
  2698.  
  2699.     def do_cmd(self, cmd, length):
  2700.         start_address = self.address
  2701.  
  2702.         cmd_length = length - 1
  2703.  
  2704.         output = []
  2705.  
  2706.         if length > lowmax:
  2707.             output.append(
  2708.                 (self.commands['long'] << 5)
  2709.               + (self.commands[cmd] << 2)
  2710.               + (cmd_length >> 8)
  2711.             )
  2712.             output.append(
  2713.                 cmd_length & 0xff
  2714.             )
  2715.         else:
  2716.             output.append(
  2717.                 (self.commands[cmd] << 5)
  2718.               + cmd_length
  2719.             )
  2720.  
  2721.         self.helpers['blank'] = [] # quick hack
  2722.         output += self.helpers.get(cmd, [])
  2723.  
  2724.         if cmd in self.lookback_methods:
  2725.             offset = self.offsets[cmd]
  2726.             # Negative offsets are one byte.
  2727.             # Positive offsets are two.
  2728.             if 0 < start_address - offset - 1 <= 0x7f:
  2729.                 offset = (start_address - offset - 1) | 0x80
  2730.                 output += [offset]
  2731.             else:
  2732.                 output += [offset / 0x100, offset % 0x100] # big endian
  2733.  
  2734.         if self.debug:
  2735.             print(' '.join(map(str, [
  2736.                   cmd, length, '\t',
  2737.                   ' '.join(map('{:02x}'.format, output)),
  2738.                   self.data[start_address:start_address+length] if cmd in self.lookback_methods else '',
  2739.             ])))
  2740.         self.output += output
  2741.  
  2742.  
  2743.  
  2744. class Decompressed:
  2745.     """
  2746.     Interpret and decompress lz-compressed data, usually 2bpp.
  2747.     """
  2748.  
  2749.     """
  2750.     Usage:
  2751.         data = Decompressed(lz).output
  2752.     or
  2753.         data = Decompressed().decompress(lz)
  2754.     or
  2755.         d = Decompressed()
  2756.         d.lz = lz
  2757.         data = d.decompress()
  2758.  
  2759.     To decompress from offset 0x80000 in a rom:
  2760.         data = Decompressed(rom, start=0x80000).output
  2761.     """
  2762.  
  2763.     lz = None
  2764.     start = 0
  2765.     commands = lz_commands
  2766.     debug = False
  2767.  
  2768.     arg_names = 'lz', 'start', 'commands', 'debug'
  2769.  
  2770.     def __init__(self, *args, **kwargs):
  2771.         self.__dict__.update(dict(zip(self.arg_names, args)))
  2772.         self.__dict__.update(kwargs)
  2773.  
  2774.         self.command_names = dict(map(reversed, self.commands.items()))
  2775.         self.address = self.start
  2776.  
  2777.         if self.lz is not None:
  2778.             self.decompress()
  2779.  
  2780.         if self.debug: print( self.command_list() )
  2781.  
  2782.  
  2783.     def command_list(self):
  2784.         """
  2785.         Print a list of commands that were used. Useful for debugging.
  2786.         """
  2787.  
  2788.         text = ''
  2789.  
  2790.         output_address = 0
  2791.         for name, attrs in self.used_commands:
  2792.             length     = attrs['length']
  2793.             address    = attrs['address']
  2794.             offset     = attrs['offset']
  2795.             direction  = attrs['direction']
  2796.  
  2797.             text += '{2:03x} {0}: {1}'.format(name, length, output_address)
  2798.             text += '\t' + ' '.join(
  2799.                 '{:02x}'.format(int(byte))
  2800.                 for byte in self.lz[ address : address + attrs['cmd_length'] ]
  2801.             )
  2802.  
  2803.             if offset is not None:
  2804.                 repeated_data = self.output[ offset : offset + length * direction : direction ]
  2805.                 if name == 'flip':
  2806.                     repeated_data = map(bit_flipped.__getitem__, repeated_data)
  2807.                 text += ' [' + ' '.join(map('{:02x}'.format, repeated_data)) + ']'
  2808.  
  2809.             text += '\n'
  2810.             output_address += length
  2811.  
  2812.         return text
  2813.  
  2814.  
  2815.     def decompress(self, lz=None):
  2816.  
  2817.         if lz is not None:
  2818.             self.lz = lz
  2819.  
  2820.         self.lz = bytearray(self.lz)
  2821.  
  2822.         self.used_commands = []
  2823.         self.output = []
  2824.  
  2825.         while 1:
  2826.  
  2827.             cmd_address = self.address
  2828.             self.offset = None
  2829.             self.direction = None
  2830.  
  2831.             if (self.byte == lz_end):
  2832.                 self.next()
  2833.                 break
  2834.  
  2835.             self.cmd = (self.byte & 0b11100000) >> 5
  2836.  
  2837.             if self.cmd_name == 'long':
  2838.                 # 10-bit length
  2839.                 self.cmd = (self.byte & 0b00011100) >> 2
  2840.                 self.length = (self.next() & 0b00000011) * 0x100
  2841.                 self.length += self.next() + 1
  2842.             else:
  2843.                 # 5-bit length
  2844.                 self.length = (self.next() & 0b00011111) + 1
  2845.  
  2846.             self.__class__.__dict__[self.cmd_name](self)
  2847.  
  2848.             self.used_commands += [(
  2849.                 self.cmd_name,
  2850.                 {
  2851.                     'length':     self.length,
  2852.                     'address':    cmd_address,
  2853.                     'offset':     self.offset,
  2854.                     'cmd_length': self.address - cmd_address,
  2855.                     'direction':  self.direction,
  2856.                 }
  2857.             )]
  2858.  
  2859.         # Keep track of the data we just decompressed.
  2860.         self.compressed_data = self.lz[self.start : self.address]
  2861.  
  2862.  
  2863.     @property
  2864.     def byte(self):
  2865.         return self.lz[ self.address ]
  2866.  
  2867.     def next(self):
  2868.         byte = self.byte
  2869.         self.address += 1
  2870.         return byte
  2871.  
  2872.     @property
  2873.     def cmd_name(self):
  2874.         return self.command_names.get(self.cmd)
  2875.  
  2876.  
  2877.     def get_offset(self):
  2878.  
  2879.         if self.byte >= 0x80: # negative
  2880.             # negative
  2881.             offset = self.next() & 0x7f
  2882.             offset = len(self.output) - offset - 1
  2883.         else:
  2884.             # positive
  2885.             offset =  self.next() * 0x100
  2886.             offset += self.next()
  2887.  
  2888.         self.offset = offset
  2889.  
  2890.  
  2891.     def literal(self):
  2892.         """
  2893.         Copy data directly.
  2894.         """
  2895.         self.output  += self.lz[ self.address : self.address + self.length ]
  2896.         self.address += self.length
  2897.  
  2898.     def iterate(self):
  2899.         """
  2900.         Write one byte repeatedly.
  2901.         """
  2902.         self.output += [self.next()] * self.length
  2903.  
  2904.     def alternate(self):
  2905.         """
  2906.         Write alternating bytes.
  2907.         """
  2908.         alts = [self.next(), self.next()]
  2909.         self.output += [ alts[x & 1] for x in xrange(self.length) ]
  2910.  
  2911.     def blank(self):
  2912.         """
  2913.         Write zeros.
  2914.         """
  2915.         self.output += [0] * self.length
  2916.  
  2917.     def flip(self):
  2918.         """
  2919.         Repeat flipped bytes from output.
  2920.  
  2921.         Example: 11100100 -> 00100111
  2922.         """
  2923.         self._repeat(table=bit_flipped)
  2924.  
  2925.     def reverse(self):
  2926.         """
  2927.         Repeat reversed bytes from output.
  2928.         """
  2929.         self._repeat(direction=-1)
  2930.  
  2931.     def repeat(self):
  2932.         """
  2933.         Repeat bytes from output.
  2934.         """
  2935.         self._repeat()
  2936.  
  2937.     def _repeat(self, direction=1, table=None):
  2938.         self.get_offset()
  2939.         self.direction = direction
  2940.         # Note: appends must be one at a time (this way, repeats can draw from themselves if required)
  2941.         for i in xrange(self.length):
  2942.             byte = self.output[ self.offset + i * direction ]
  2943.             self.output.append( table[byte] if table else byte )
  2944.  
  2945.  
  2946.  
  2947. def connect(tiles):
  2948.     """
  2949.     Combine 8x8 tiles into a 2bpp image.
  2950.     """
  2951.     return [byte for tile in tiles for byte in tile]
  2952.  
  2953. def transpose(tiles, width=None):
  2954.     """
  2955.     Transpose a tile arrangement along line y=-x.
  2956.  
  2957.       00 01 02 03 04 05     00 06 0c 12 18 1e
  2958.       06 07 08 09 0a 0b     01 07 0d 13 19 1f
  2959.       0c 0d 0e 0f 10 11 <-> 02 08 0e 14 1a 20
  2960.       12 13 14 15 16 17     03 09 0f 15 1b 21
  2961.       18 19 1a 1b 1c 1d     04 0a 10 16 1c 22
  2962.       1e 1f 20 21 22 23     05 0b 11 17 1d 23
  2963.  
  2964.       00 01 02 03     00 04 08
  2965.       04 05 06 07 <-> 01 05 09
  2966.       08 09 0a 0b     02 06 0a
  2967.                       03 07 0b
  2968.     """
  2969.     if width == None:
  2970.         width = int(sqrt(len(tiles))) # assume square image
  2971.     tiles = sorted(enumerate(tiles), key= lambda i_tile: i_tile[0] % width)
  2972.     return [tile for i, tile in tiles]
  2973.  
  2974. def transpose_tiles(image, width=None):
  2975.     return connect(transpose(get_tiles(image), width))
  2976.  
  2977. def bitflip(x, n):
  2978.     r = 0
  2979.     while n:
  2980.         r = (r << 1) | (x & 1)
  2981.         x >>= 1
  2982.         n -= 1
  2983.     return r
  2984.  
  2985.  
  2986. class Decompressor:
  2987.     """
  2988.     pokered pic decompression.
  2989.  
  2990.     Ported to python 2.7 from the python 3 code at https://github.com/magical/pokemon-sprites-rby.
  2991.     """
  2992.  
  2993.     table1 = [(2 << i) - 1 for i in range(16)]
  2994.     table2 = [
  2995.         [0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5, 0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa],
  2996.         [0xf, 0xe, 0xc, 0xd, 0x8, 0x9, 0xb, 0xa, 0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5], # prev ^ 0xf
  2997.         [0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa, 0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5],
  2998.         [0xf, 0x7, 0x3, 0xb, 0x1, 0x9, 0xd, 0x5, 0x0, 0x8, 0xc, 0x4, 0xe, 0x6, 0x2, 0xa], # prev ^ 0xf
  2999.     ]
  3000.     table3 = [bitflip(i, 4) for i in range(16)]
  3001.  
  3002.     tilesize = 8
  3003.  
  3004.  
  3005.     def __init__(self, f=None, d=None, mirror=False, planar=True):
  3006.         if f is not None:
  3007.             self.bs = fbitstream( f )
  3008.         elif d is not None:
  3009.             self.bs = bitstream( d )
  3010.         else:
  3011.             print("No decompressed data specified")
  3012.             raise
  3013.         self.mirror = mirror
  3014.         self.planar = planar
  3015.         self.data = None
  3016.    
  3017.     def decompress(self):
  3018.         rams = [[], []]
  3019.  
  3020.         self.sizex  = self._readint(4) * self.tilesize
  3021.         self.sizey  = self._readint(4)
  3022.  
  3023.         self.size = self.sizex * self.sizey
  3024.  
  3025.         self.ramorder = self._readbit()
  3026.  
  3027.         r1 = self.ramorder
  3028.         r2 = self.ramorder ^ 1
  3029.  
  3030.         self._fillram(rams[r1])
  3031.         mode = self._readbit()
  3032.         if mode:
  3033.             mode += self._readbit()
  3034.         self._fillram(rams[r2])
  3035.  
  3036.         rams[0] = bytearray(bitgroups_to_bytes(rams[0]))
  3037.         rams[1] = bytearray(bitgroups_to_bytes(rams[1]))
  3038.  
  3039.         if mode == 0:
  3040.             self._decode(rams[0])
  3041.             self._decode(rams[1])
  3042.         elif mode == 1:
  3043.             self._decode(rams[r1])
  3044.             self._xor(rams[r1], rams[r2])
  3045.         elif mode == 2:
  3046.             self._decode(rams[r2], mirror=False)
  3047.             self._decode(rams[r1])
  3048.             self._xor(rams[r1], rams[r2])
  3049.         else:
  3050.             raise Exception("Invalid deinterlace mode!")
  3051.  
  3052.         data = []
  3053.         if self.planar:
  3054.             for a, b in zip(rams[0], rams[1]):
  3055.                 data += [a, b]
  3056.             self.data = bytearray(data)
  3057.         else:
  3058.             for a, b in zip(bitstream(rams[0]), bitstream(rams[1])):
  3059.                 data.append(a | (b << 1))
  3060.             self.data = bitgroups_to_bytes(data)
  3061.  
  3062.     def _fillram(self, ram):
  3063.         mode = ['rle', 'data'][self._readbit()]
  3064.         size = self.size * 4
  3065.         while len(ram) < size:
  3066.             if mode == 'rle':
  3067.                 self._read_rle_chunk(ram)
  3068.                 mode = 'data'
  3069.             elif mode == 'data':
  3070.                 self._read_data_chunk(ram, size)
  3071.                 mode = 'rle'
  3072.         '''if len(ram) > size:
  3073.             #ram = ram[:size]
  3074.             raise ValueError(size, len(ram))
  3075.         '''
  3076.         ram[:] = self._deinterlace_bitgroups(ram)
  3077.  
  3078.     def _read_rle_chunk(self, ram):
  3079.  
  3080.         i = 0
  3081.         while self._readbit():
  3082.             i += 1
  3083.  
  3084.         n = self.table1[i]
  3085.         a = self._readint(i + 1)
  3086.         n += a
  3087.  
  3088.         for i in range(n):
  3089.             ram.append(0)
  3090.  
  3091.     def _read_data_chunk(self, ram, size):
  3092.         while 1:
  3093.             bitgroup = self._readint(2)
  3094.             if bitgroup == 0:
  3095.                 break
  3096.             ram.append(bitgroup)
  3097.  
  3098.             if size <= len(ram):
  3099.                 break
  3100.  
  3101.     def _decode(self, ram, mirror=None):
  3102.         if mirror is None:
  3103.             mirror = self.mirror
  3104.  
  3105.         for x in range(self.sizex):
  3106.             bit = 0
  3107.             for y in range(self.sizey):
  3108.                 i = y * self.sizex + x
  3109.                 a = (ram[i] >> 4) & 0xf
  3110.                 b = ram[i] & 0xf
  3111.  
  3112.                 a = self.table2[bit][a]
  3113.                 bit = a & 1
  3114.                 if mirror:
  3115.                     a = self.table3[a]
  3116.  
  3117.                 b = self.table2[bit][b]
  3118.                 bit = b & 1
  3119.                 if mirror:
  3120.                     b = self.table3[b]
  3121.  
  3122.                 ram[i] = (a << 4) | b
  3123.  
  3124.     def _xor(self, ram1, ram2, mirror=None):
  3125.         if mirror is None:
  3126.             mirror = self.mirror
  3127.  
  3128.         for i in range(len(ram2)):
  3129.             if mirror:
  3130.                 a = (ram2[i] >> 4) & 0xf
  3131.                 b = ram2[i] & 0xf
  3132.                 a = self.table3[a]
  3133.                 b = self.table3[b]
  3134.                 ram2[i] = (a << 4) | b
  3135.  
  3136.             ram2[i] ^= ram1[i]
  3137.  
  3138.     def _deinterlace_bitgroups(self, bits):
  3139.         l = []
  3140.         for y in range(self.sizey):
  3141.             for x in range(self.sizex):
  3142.                 i = 4 * y * self.sizex + x
  3143.                 for j in range(4):
  3144.                     l.append(bits[i])
  3145.                     i += self.sizex
  3146.         return l
  3147.  
  3148.  
  3149.     def _readbit(self):
  3150.         return next(self.bs)
  3151.  
  3152.     def _readint(self, count):
  3153.         return readint(self.bs, count)
  3154.  
  3155. def decompress(f, offset=None, mirror=False):
  3156.     """
  3157.     Decompress a pic given a file object. Return a planar 2bpp image.
  3158.  
  3159.     Optional: offset (for roms).
  3160.     """
  3161.     if offset is not None:
  3162.         f.seek(offset)
  3163.     dcmp = Decompressor(f, mirror=mirror)
  3164.     dcmp.decompress()
  3165.     return dcmp.data
  3166.  
  3167. def decomp_main( in_path, out_path ):
  3168.     if os.path.exists( out_path ):
  3169.         print( out_path + " already exists, use a different name" )
  3170.     else:
  3171.         if not os.path.exists( in_path ):
  3172.             print( in_path + " does not exist" )
  3173.         else:
  3174.             with open( in_path, 'rb' ) as in_file:
  3175.                 sprite = decompress( in_file )
  3176.                 with open( out_path, 'wb' ) as out_file:
  3177.                     out_file.write( sprite )
  3178.  
  3179. # hex converter
  3180.  
  3181. def str2byt(x):
  3182.     done = False
  3183.     out = b''
  3184.     byte = 0
  3185.     for nybble in x:
  3186.         if not done:
  3187.             byte = int( nybble, 16 )
  3188.             byte <<= 4
  3189.         else:
  3190.             byte |= ( int( nybble, 16 ) & 0xF )
  3191.             out = out + bytes([byte])
  3192.         done = not done
  3193.     return out
  3194.  
  3195. hx1 = lambda n: hex(n)[2:].zfill(2)
  3196.  
  3197. REGEX_DEC = re.compile(r'^  d[bw]   ((?:[0-9]{3},?)+)\s*(?:;.*)?$')
  3198. REGEX_HEX = re.compile(r'^  d[bw]   ((?:0(?:[A-Fa-f0-9]{2})+h,?)+)\s*(?:;.*)?$')
  3199.  
  3200. def divhx(x):
  3201.     pcs = []
  3202.     for i in range(len(x) // 2):
  3203.         pcs.append(x[i*2:i*2+2])
  3204.     return ''.join(pcs[::-1])
  3205.    
  3206. def divhxl(x):
  3207.     out = ''
  3208.     for k in x:
  3209.         out += divhx(k)
  3210.     return out
  3211.  
  3212. def hex_convert_main( input_string, data_type=None, v=None ):
  3213.     lins = input_string.splitlines()
  3214.     if not data_type: # we have to guess the type
  3215.         r1 = REGEX_DEC
  3216.         r2 = REGEX_HEX
  3217.         for i in range( len(lins) ):
  3218.             if r1.match(lins[i]):
  3219.                 data_type = 0
  3220.                 break
  3221.             elif r2.match(lins[i]):
  3222.                 data_type = 1
  3223.                 break
  3224.         else:
  3225.             raise Exception
  3226.     if data_type == 0: # decimal
  3227.         reg = REGEX_DEC
  3228.     else:
  3229.         reg = REGEX_HEX
  3230.     datout = ''
  3231.     uuct = 0
  3232.     for i,l in enumerate(lins):
  3233.         mm = reg.match(l)
  3234.         if mm:
  3235.             cur_dats = mm.groups()[0].split(',')
  3236.             if data_type == 0:
  3237.                 cur_dats_2 = [hx1(int(x[1:-1])) for x in cur_dats]
  3238.             else:
  3239.                 cur_dats_2 = [x[1:-1] for x in cur_dats]
  3240.             datout += ''.join(cur_dats_2)
  3241.         else:
  3242.             if v:
  3243.                 print('Line does not match pattern:')
  3244.                 print(i+1, l)
  3245.                 uuct += 1
  3246.                 if uuct > 10:
  3247.                     input()
  3248.                     uuct = 0
  3249.    
  3250.     try:
  3251.         return bytes.fromhex( datout )
  3252.     except AttributeError:
  3253.         return bytearray.fromhex( datout )
  3254.  
  3255.  
  3256. # makeimg.py
  3257.  
  3258. GBGRY = ((232, 232, 232), (160, 160, 160), (88, 88, 88), (16, 16, 16))
  3259. GBGRH = ((255, 255, 255), (176, 176, 176), (104, 104, 104), (0, 0, 0))
  3260. GBGRN = ((224, 248, 208), (136, 192, 112), (52, 104, 86), (8, 24, 32))
  3261. # GBGRHR = ((255, 255, 255), (104, 104, 104), (176, 176, 176), (0, 0, 0))
  3262. GBGRHR = ((255, 255, 255), (85, 85, 85), (170, 170, 170), (0, 0, 0))
  3263. #          255 * 1       ,  255 * (1/3),   255 * (2/3)  ,  255 * 0
  3264.  
  3265. def makeimg_main(args):
  3266.     dat = None
  3267.     cbb = False
  3268.     _2bpp = False
  3269.     nothing = False
  3270.    
  3271.     if not args.out and not args.width and not args.depth and not args.scale and not args.reverse and not args.sprite and not args.subblk \
  3272.     and not args.view and not args.palette and not args.horizontal and not args.square and not args.extradat and not args.flag_iscbb \
  3273.     and not args.flag_is2bpp and not args.flag_ishex and not args.flag_decimal and not args.flag_guess and not args.flag_compressed:
  3274.         nothing = True
  3275.    
  3276.     output_name = ''
  3277.     # is args.data a data string or a file name?
  3278.     if re.match(r'^(?:[A-Fa-f0-9]{2})+$', args.data): # is args.data a data string?
  3279.         if not os.path.exists( args.data ): # is it *REALLY* a data string and not a file path?
  3280.             dat = bytearray(args.data, hex_char_encoding) # if yes then use it directly
  3281.             if not args.out:
  3282.                 print('-out is required if the input is a bytestring')
  3283.                 exit()
  3284.     else: # if not, then open the filename and convert it to a datastring
  3285.         with open( args.data, 'rb' ) as f:
  3286.             raw = f.read()
  3287.             try:
  3288.                 dat = bytearray( raw, hex_char_encoding )
  3289.             except TypeError:
  3290.                 dat = bytearray( raw )
  3291.             if ".cbb" in args.data.lower():
  3292.                 cbb = True
  3293.             if '.2bpp' in args.data.lower():
  3294.                 _2bpp = True
  3295.             output_name = args.data.lower()
  3296.    
  3297.     if args.flag_ishex or nothing: # Decode a hex-coded file'
  3298.         if not args.flag_is2bpp and not cbb and not _2bpp:
  3299.             string = dat.decode( hex_char_encoding )
  3300.             if args.flag_guess:
  3301.                 dat = hex_convert_main( string )
  3302.             else:
  3303.                 if args.flag_decimal:
  3304.                     dat = hex_convert_main( string, data_type = 0 )
  3305.                 else:
  3306.                     dat = hex_convert_main( string, data_type = 1 )
  3307.    
  3308.     if args.flag_iscbb or cbb:
  3309.         if not _2bpp:
  3310.             i = 16
  3311.             while i < len(dat):
  3312.                 del dat[i:i+16]
  3313.                 i += 16
  3314.    
  3315.     if args.flag_compressed or ( nothing and (not cbb) ): # decompress a file
  3316.         if ( not _2bpp ) or args.flag_compressed:
  3317.             decomp = Decompressor( d = dat, mirror = False )
  3318.             decomp.decompress()
  3319.             dat = decomp.data
  3320.    
  3321.     # the PNG converter expects its input to be in a hex string format
  3322.     try:
  3323.         dat = dat.hex()
  3324.     except AttributeError:
  3325.         dat = "".join("%02x" % b for b in dat)
  3326.    
  3327.     # Convert data to a PNG
  3328.    
  3329.     # get various settings
  3330.     imglen = len(dat) // 2
  3331.     if (not args.width) or nothing: # autodetect width
  3332.         imgw = int((imglen // 16) ** 0.5)
  3333.         print( "Assuming image is square...\nwidth/height in tiles: " + str(imgw) )
  3334.     else:
  3335.         imgw = args.width
  3336.     output_name += '_' + str(imgw)+'x'+str(imgw)+'_'
  3337.    
  3338.     if args.depth:
  3339.         assert args.depth in (1, 2, 4, 8)
  3340.         bitd = args.depth
  3341.     else:
  3342.         bitd = 2
  3343.    
  3344.     if not args.palette:
  3345.         cols = GBGRHR
  3346.     if args.palette == 'green':
  3347.         cols = GBGRN
  3348.     elif args.palette == 'gray':
  3349.         cols = GBGRY
  3350.     elif args.palette == 'grayhi':
  3351.         cols = GBGRH
  3352.     elif args.palette == 'grayhir':
  3353.         cols = GBGRHR
  3354.     elif type(args.palette) == type(''): # guess shades
  3355.         # automatic gain control
  3356.         cols = []
  3357.         if args.sprite:
  3358.             for i in range(2 ** bitd - 1):
  3359.                 pcnt = i / (2 ** bitd - 2)
  3360.                 # ~ print(pcnt)
  3361.                 pcnt = math.sqrt(pcnt)
  3362.                 # ~ print(pcnt)
  3363.                 rgb = int(255 * pcnt)
  3364.                 cols.append( [rgb, rgb, rgb] )
  3365.             if args.reverse:
  3366.                 cols = cols[::-1]
  3367.             cols.insert(0, [255, 255, 255])
  3368.         else:
  3369.             for i in range(2 ** bitd):
  3370.                 pcnt = i / (2 ** bitd - 1)
  3371.                 rgb = int(255 * pcnt)
  3372.                 cols.append([rgb, rgb, rgb])
  3373.             if args.reverse:
  3374.                 cols = cols[::-1]
  3375.     imgh = imglen // bitd // imgw // 8
  3376.     if imgh * imgw * bitd * 8 != imglen:
  3377.         imgh += 1
  3378.    
  3379.     out_tmp = [[0 for x in range(imgw*8*3)] for y in range(imgh*8)]
  3380.    
  3381.     #out = Image.new('RGB', (imgw*8, imgh*8))
  3382.     #pxa = out.load()
  3383.     i = 0
  3384.     binimg = ''
  3385.     for nib in dat:
  3386.         binimg += bin(int(nib, 16))[2:].zfill(4)
  3387.     if bitd != 2:
  3388.         if args.subblk:
  3389.             for blky2 in range(imgh // 2):
  3390.                 for blkx2 in range(imgw // 2):
  3391.                     for blkyy in range(2):
  3392.                         for blkxx in range(2):
  3393.                             for yy in range(8):
  3394.                                 for xx in range(8):
  3395.                                     blkx = blkx2 * 2 + blkxx
  3396.                                     blky = blky2 * 2 + blkyy
  3397.                                     if nothing or not args.horizontal:
  3398.                                         x = blky * 8 + xx
  3399.                                         y = blkx * 8 + yy
  3400.                                     else:
  3401.                                         x = blkx * 8 + xx
  3402.                                         y = blky * 8 + yy
  3403.                                     if i >= len(binimg):
  3404.                                         continue
  3405.                                     colid = binimg[i*bitd:(i+1)*bitd]
  3406.                                     out_tmp[y][(x*3)+0] = cols[int(colid, 2)][0]
  3407.                                     out_tmp[y][(x*3)+1] = cols[int(colid, 2)][1]
  3408.                                     out_tmp[y][(x*3)+2] = cols[int(colid, 2)][2]
  3409.                                     i += 1
  3410.         else:
  3411.             for blky in range(imgh):
  3412.                 for blkx in range(imgw):
  3413.                     for yy in range(8):
  3414.                         for xx in range(8):
  3415.                             if nothing or not args.horizontal:
  3416.                                 x = blky * 8 + xx
  3417.                                 y = blkx * 8 + yy
  3418.                             else:
  3419.                                 x = blkx * 8 + xx
  3420.                                 y = blky * 8 + yy
  3421.                             if i >= len(binimg):
  3422.                                 continue
  3423.                             colid = binimg[i*bitd:(i+1)*bitd]
  3424.                             out_tmp[y][(x*3)+0] = cols[int(colid, 2)][0]
  3425.                             out_tmp[y][(x*3)+1] = cols[int(colid, 2)][1]
  3426.                             out_tmp[y][(x*3)+2] = cols[int(colid, 2)][2]
  3427.                             i += 1
  3428.     else:
  3429.         if args.subblk:
  3430.             for blky2 in range(imgh // 2):
  3431.                 for blkx2 in range(imgw // 2):
  3432.                     for blkyy in range(2):
  3433.                         for blkxx in range(2):
  3434.                             for yy in range(8):
  3435.                                 for xx in range(8):
  3436.                                     blkx = blkx2 * 2 + blkxx
  3437.                                     blky = blky2 * 2 + blkyy
  3438.                                     if nothing or not args.horizontal:
  3439.                                         x = blky * 8 + xx
  3440.                                         y = blkx * 8 + yy
  3441.                                     else:
  3442.                                         x = blkx * 8 + xx
  3443.                                         y = blky * 8 + yy
  3444.                                     if i >= len(binimg):
  3445.                                         continue
  3446.                                     colidhi = binimg[i]
  3447.                                     colidlo = binimg[i + 8]
  3448.                                     colid = colidhi + colidlo
  3449.                                     out_tmp[y][(x*3)+0] = cols[int(colid, 2)][0]
  3450.                                     out_tmp[y][(x*3)+1] = cols[int(colid, 2)][1]
  3451.                                     out_tmp[y][(x*3)+2] = cols[int(colid, 2)][2]
  3452.                                     i += 1
  3453.                                 i += 8
  3454.         else:
  3455.             for blky in range(imgh):
  3456.                 for blkx in range(imgw):
  3457.                     for yy in range(8):
  3458.                         for xx in range(8):
  3459.                             if nothing or not args.horizontal:
  3460.                                 x = blky * 8 + xx
  3461.                                 y = blkx * 8 + yy
  3462.                             else:
  3463.                                 x = blkx * 8 + xx
  3464.                                 y = blky * 8 + yy
  3465.                             if i >= len(binimg):
  3466.                                 continue
  3467.                             try:
  3468.                                 colidhi = binimg[i]
  3469.                                 colidlo = binimg[i + 8]
  3470.                                 colid = colidhi + colidlo
  3471.                                 out_tmp[y][(x*3)+0] = cols[int(colid, 2)][0]
  3472.                                 out_tmp[y][(x*3)+1] = cols[int(colid, 2)][1]
  3473.                                 out_tmp[y][(x*3)+2] = cols[int(colid, 2)][2]
  3474.                             except IndexError:
  3475.                                 pass
  3476.                             i += 1
  3477.                         i += 8
  3478.     if args.scale:
  3479.         print("Sprite scaling is not implemented")
  3480.     #   oldw, oldh = out.size
  3481.     #   neww = oldw * args.scale
  3482.     #   newh = oldh * args.scale
  3483.     #   out = out.resize((neww, newh))
  3484.     #   # ~ out = transform.scale(out, (neww, newh))
  3485.     out = from_array( out_tmp, 'RGB' )
  3486.     if not args.out:
  3487.         output_name += '.png'
  3488.         out.save(output_name)
  3489.     else:
  3490.         out.save(args.out)
  3491. #
  3492.  
  3493. if __name__ == '__main__':
  3494.     parser = argparse.ArgumentParser(
  3495.         formatter_class=argparse.RawDescriptionHelpFormatter,
  3496.         description=textwrap.dedent("""Make a PNG from a given GB graphic file
  3497.  
  3498. Examples:
  3499. makeimg_standalone.py SPRITE.DAT
  3500. makeimg_standalone.py SPRITE.CBB
  3501. makeimg_standalone.py -cbb -C SPRITE.CBB
  3502. makeimg_standalone.py -H SPRITE.DAT
  3503. makeimg_standalone.py -H -C SPRITE.DAT
  3504. makeimg_standalone.py -H -C -d 2 -l grayhir -tq SPRITE.DAT
  3505. makeimg_standalone.py -H -C -s 6 -d 2 -l grayhir -q SPRITE.DAT
  3506. makeimg_standalone.py -H -C -s 6 -d 2 -l grayhir -q -out sprite_out.png SPRITE.DAT
  3507. """)
  3508.     )
  3509.     parser.add_argument('data', type=str, help='byte string OR file to read')
  3510.     parser.add_argument('-out', type=str, help='filename to write (required if using a byte string)')
  3511.     parser.add_argument('-w', dest='width', type=int, help='width (in 8x8 tiles) of the image')
  3512.     parser.add_argument('-d', dest='depth', type=int, help='bit depth (1, 2, 4, 8)')
  3513.     parser.add_argument('-s', dest='scale', type=int, help='<- make the sprite this much bigger (not implemented)')
  3514.     parser.add_argument('-r', dest='reverse', action='store_true', help='reverse colors')
  3515.     parser.add_argument('-p', dest='sprite', action='store_true', help='use sprite colors')
  3516.     parser.add_argument('-k', dest='subblk', action='store_true', help='subblk')
  3517.     parser.add_argument('-v', dest='view', action='store_true', help='View the output file automatically')
  3518.     parser.add_argument('-l', dest='palette', type=str, help='The palette to use ("green"/"gray"/"grayhi"/"grayhir")')
  3519.     parser.add_argument('-t', dest='horizontal', action='store_true', help='Horizontal tile order?')
  3520.     parser.add_argument('-q', dest='square', action='store_true', help=  'Assume square image, autodetect width')
  3521.     parser.add_argument('-x', dest='extradat', action='store_true', help='extradat (?)')
  3522.     parser.add_argument('-cbb',  dest='flag_iscbb',  action='store_true', help='The input file is a CBB graphic')
  3523.     parser.add_argument('-2bpp', dest='flag_is2bpp', action='store_true', help='The input file is a raw 2bpp graphic')
  3524.     parser.add_argument('-H', dest='flag_ishex', action='store_true', help='The input file is a text-encoded binary')
  3525.     parser.add_argument('-D', dest='flag_decimal', action='store_true', help='(if -H) Input data type is decimal')
  3526.     parser.add_argument('-G', dest='flag_guess', action='store_true', help='(if -H) Try to guess input data type')
  3527.     parser.add_argument('-C', dest='flag_compressed', action='store_true', help='The input file is compressed')
  3528.     #parser.add_argument('-simple', dest="flag_simple", action='store_true', help='The same as -H -C -s 1')
  3529.    
  3530.     args = parser.parse_args()
  3531.    
  3532.     makeimg_main(args)
  3533. #
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement