Add trans support to PngImage

tmxlib/image_png.py

@@ -1,25 +1,34 @@
 
 from __future__ import division
 
-from six import BytesIO
+import itertools
 
+from six import BytesIO
 import png
+from array import array
 
 import tmxlib
 import tmxlib.image_base
 
 
+def _grouper(iterable, n):
+    "Collect data into fixed-length chunks, ignoring extras at the end"
+    # grouper('ABCDEFG', 3) --> ABC DEF
+    args = [iter(iterable)] * n
+    return itertools.izip(*args)
+
+
 class PngImage(tmxlib.image_base.Image):
     def load_image(self):
         """Load the image from self.data, and set self.size
         """
         try:
-            self._image_data
+            self._image_data_original
             return self.size
         except AttributeError:
             reader = png.Reader(bytes=self.data).asRGBA8()
             w, h, data, meta = reader
-            self._image_data = tuple(data)
+            self._image_data_original = tuple(data)
             if self._size:
                 assert (w, h) == self._size
             else:
@@ -32,8 +41,33 @@ def image_data(self):
             return self._image_data
         except AttributeError:
             self.load_image()
+            data = self._image_data_original
+            if self.trans:
+                xtrans = tuple(int(n * 255) for n in self.trans[:3])
+                new_data = []
+                for line in data:
+                    new_data.append(array(
+                        'B',
+                        itertools.chain.from_iterable(
+                            v[:3] + (0,) if tuple(v[:3]) == xtrans else v
+                            for v in _grouper(line, 4))))
+                self._image_data = new_data
+            else:
+                self._image_data = data
             return self._image_data
 
+    @property
+    def trans(self):
+        return self._trans
+
+    @trans.setter
+    def trans(self, new_trans):
+        self._trans = new_trans
+        try:
+            del self._image_data
+        except AttributeError:
+            pass
+
     def get_pixel(self, x, y):
         x, y = self._wrap_coords(x, y)
         return tuple(v / 255 for v in self.image_data[y][x * 4:(x + 1) * 4])
@@ -43,7 +77,9 @@ def _repr_png_(self, _crop_box=None):
 
         See: http://ipython.org/ipython-doc/stable/config/integrating.html
         """
-        if _crop_box:
+        if _crop_box or self.trans:
+            if not _crop_box:
+                _crop_box = 0, 0, self.width, self.height
             left, up, right, low = _crop_box
             data = [l[left * 4:right * 4] for l in self.image_data[up:low]]
             out = BytesIO()