From ca52dffdd11a52df3006d9197cbde4c96eae59e8 Mon Sep 17 00:00:00 2001
From: Sergey Lavrinenko <s@lavr.me>
Date: Sat, 4 Apr 2015 16:28:17 +0300
Subject: [PATCH] Transformer: do a callback to check load image or not. Fixes
 #50

---
 emails/testsuite/loader/test_loaders.py | 26 ++++++++++++++++++++
 emails/transformer.py                   | 32 +++++++++++++++++++------
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/emails/testsuite/loader/test_loaders.py b/emails/testsuite/loader/test_loaders.py
index f514fd1..c7de9d6 100644
--- a/emails/testsuite/loader/test_loaders.py
+++ b/emails/testsuite/loader/test_loaders.py
@@ -136,6 +136,32 @@ def test_loaders_with_params():
                 assert a.is_inline is True
 
 
+def test_loader_image_callback():
+
+    checked_images = []
+
+    def check_image_callback(el, **kwargs):
+        if hasattr(el, 'attrib'):
+            checked_images.append(el.attrib['src'])
+        elif hasattr(el, 'uri'):
+            checked_images.append(el.uri)
+        else:
+            assert 0, "el should be lxml.etree._Element or cssutils.css.value.URIValue"
+        return False
+
+    for message in load_messages(load_images=check_image_callback, **OLDORNAMENT_URLS):
+        # Check images not loaded
+        assert len(message.attachments.keys()) == 0
+
+    total_images = 0
+    for message in load_messages(**OLDORNAMENT_URLS):
+        # Check loaded images
+        assert len(message.attachments.keys()) == 13
+        total_images += len(message.attachments.keys())
+
+    assert len(checked_images) >= total_images
+
+
 def test_external_urls():
 
     # Load some real sites with complicated html and css.
diff --git a/emails/transformer.py b/emails/transformer.py
index fcb49a5..af9aa49 100644
--- a/emails/transformer.py
+++ b/emails/transformer.py
@@ -4,13 +4,15 @@ import logging
 import posixpath
 import re
 import warnings
+import functools
 
 from cssutils import CSSParser
+
 from lxml import etree
 from premailer import Premailer
 from premailer.premailer import ExternalNotFoundError
 
-from .compat import urlparse, to_unicode
+from .compat import urlparse, to_unicode, is_callable
 from .store import MemoryFileStore, LazyHTTPFile
 from .template.base import BaseTemplate
 from .loader.local_store import FileNotFound
@@ -216,16 +218,27 @@ class BaseTransformer(HTMLParser):
                and el.attrib.get(self.html_attribute_name) \
                or None
 
-    def _load_attachment_func(self, uri, element=None, **kw):
+    def _default_attachment_check(self, el, hints):
+        if hints['attrib'] == 'ignore':
+            return False
+        else:
+            return True
+
+    def _load_attachment_func(self, uri, element=None, callback=None, **kw):
 
         #
         # Load uri from remote url or from local_store
         # Return local uri
         #
 
-        # Ignore tags with attribute data-emails="ignore"
-        attribute_value = self._attribute_value(element)
-        if attribute_value == 'ignore':
+        if callback is None:
+            # Default callback: skip images with data-emails="ignore" attribute
+            callback = lambda _, hints: hints['attrib'] != 'ignore'
+
+        attribute_value = self._attribute_value(element) or ''
+
+        # If callback returns False, skip attachment loading
+        if not callback(element, hints={'attrib': attribute_value}):
             return uri
 
         attachment = self.attachment_store.by_uri(uri)
@@ -234,7 +247,7 @@ class BaseTransformer(HTMLParser):
                 uri=uri,
                 absolute_url=self.get_absolute_url(uri),
                 local_loader=self.local_loader,
-                content_disposition='inline' if attribute_value and 'inline' in attribute_value else None,
+                content_disposition='inline' if 'inline' in attribute_value else None,
                 requests_args=self.requests_params)
             self.attachment_store.add(attachment)
         return attachment.filename
@@ -285,8 +298,13 @@ class BaseTransformer(HTMLParser):
             self.get_premailer(**kw).transform()
 
         # 2. Load linked images and transform links
+        # If load_images is a function, use if as callback
         if load_images:
-            self.apply_to_images(self._load_attachment_func)
+            if is_callable(load_images):
+                func = functools.partial(self._load_attachment_func, callback=load_images)
+            else:
+                func = self._load_attachment_func
+            self.apply_to_images(func)
 
         # 3. Remove unsafe tags is requested
         if remove_unsafe_tags: