1 files changed, 31 insertions, 47 deletions
diff --git a/dedup/image.py b/dedup/image.py
index e05e7da..78f8bd9 100644
--- a/dedup/image.py
+++ b/dedup/image.py
@@ -1,9 +1,20 @@
-import io
 import struct
 
 import PIL.Image
 
-class ImageHash(object):
+from . import hashing
+
+def detectpng(content, maxpixels=1024 * 1024 * 32):
+    if len(content) < 33:
+        return None # defer decision
+    if not content.startswith(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR"):
+        return False
+    width, height = struct.unpack(">II", content[16:24])
+    if width * height > maxpixels:
+        return False
+    return True
+
+class ImageHash(hashing.StoredHash):
     """A hash on the contents of an image. This disregards mode, depth and meta
     information. Note that due to limitations in PIL and the image format
     (interlacing) the full contents are stored and decoded in hexdigest."""
@@ -16,52 +27,25 @@ class ImageHash(object):
         @param hashobj: a hashlib-like object
         """
         self.hashobj = hashobj
-        self.imagedetected = False
-        self.content = io.BytesIO()
-
-    def update(self, data):
-        self.content.write(data)
-        if self.content.tell() > self.maxsize:
-            raise ValueError("maximum image size exceeded")
-        if self.imagedetected:
-            return
-        if self.content.tell() < 33: # header + IHDR
-            return
-        curvalue = self.content.getvalue()
-        if curvalue.startswith(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR"):
-            width, height = struct.unpack(">II", curvalue[16:24])
-            if width * height > self.maxpixels:
-                raise ValueError("maximum image pixels exceeded")
-            self.imagedetected = True
-            return
-        raise ValueError("not a png image")
+        hashing.StoredHash.__init__(self, self.computehash, self.detect,
+                                    self.maxsize)
+        self.hashobj = hashobj
 
-    def copy(self):
-        new = ImageHash()
-        new.hashobj = self.hashobj.copy()
-        new.imagedetected = self.imagedetected
-        new.content = io.BytesIO(self.content.getvalue())
-        return new
+    def detect(self, bytesio):
+        return detectpng(bytesio.read(33), self.maxpixels)
 
-    def hexdigest(self):
-        if not self.imagedetected:
-            raise ValueError("not a png image")
+    def computehash(self, bytesio):
         hashobj = self.hashobj.copy()
-        pos = self.content.tell()
-        try:
-            self.content.seek(0)
-            img = PIL.Image.open(self.content)
-            width, height = img.size
-            pack = lambda elem: struct.pack("BBBB", *elem)
-            # special casing easy modes reduces memory usage
-            if img.mode == "L":
-                pack = lambda elem: struct.pack("BBBB", elem, elem, elem, 255)
-            elif img.mode == "RGB":
-                pack = lambda elem: struct.pack("BBBB", *(elem + (255,)))
-            elif img.mode != "RGBA":
-                img = img.convert("RGBA")
-            for elem in img.getdata():
-                hashobj.update(pack(elem))
-        finally:
-            self.content.seek(pos)
+        img = PIL.Image.open(self.content)
+        width, height = img.size
+        pack = lambda elem: struct.pack("BBBB", *elem)
+        # special casing easy modes reduces memory usage
+        if img.mode == "L":
+            pack = lambda elem: struct.pack("BBBB", elem, elem, elem, 255)
+        elif img.mode == "RGB":
+            pack = lambda elem: struct.pack("BBBB", *(elem + (255,)))
+        elif img.mode != "RGBA":
+            img = img.convert("RGBA")
+        for elem in img.getdata():
+            hashobj.update(pack(elem))
         return "%s%8.8x%8.8x" % (hashobj.hexdigest(), width, height)