Add search by color in the search API (for static themes only)

docs/topics/api/addons.rst

@@ -45,6 +45,7 @@ This endpoint allows you to search through public add-ons.
     :query string appversion: Filter by application version compatibility. Pass the full version as a string, e.g. ``46.0``. Only valid when the ``app`` parameter is also present.
     :query string author: Filter by exact (listed) author username or user id. Multiple author usernames or ids can be specified, separated by comma(s), in which case add-ons with at least one matching author are returned.
     :query string category: Filter by :ref:`category slug <category-list>`. ``app`` and ``type`` parameters need to be set, otherwise this parameter is ignored.
+    :query string color: (Experimental) Filter by color in RGB hex format, trying to find themes that approximately match the specified color. Only works for static themes.
     :query string exclude_addons: Exclude add-ons by ``slug`` or ``id``. Multiple add-ons can be specified, separated by comma(s).
     :query boolean featured: Filter to only featured add-ons.  Only ``featured=true`` is supported.
         If ``app`` is provided as a parameter then only featured collections targeted to that application are taken into account.

requirements/prod_common.txt

@@ -147,6 +147,8 @@ chardet==3.0.4 \
 click==7.0 \
     --hash=sha256:2335065e6395b9e67ca716de5f7526736bfa6ceead690adf616d925bdc622b13 \
     --hash=sha256:5b94b49521f6456670fdb30cd82a4eca9412788a93fa6dd6df72c94d5a8ff2d7
+colorgram.py==1.1.0 \
+    --hash=sha256:35f1692a2f070f65eae53645f7c3eebc36488cafeb7434838beb26def2910868
 # contextlib2 is required by raven
 contextlib2==0.5.5 \
     --hash=sha256:f5260a6e679d2ff42ec91ec5252f4eeffdcf21053db9113bd0a8e4d953769c00 \

src/olympia/addons/indexers.py

@@ -22,6 +22,7 @@ class AddonIndexer(BaseSearchIndexer):
     hidden_fields = (
         '*.raw',
         'boost',
+        'colors',
         'hotness',
         # Translated content that is used for filtering purposes is stored
         # under 3 different fields:
@@ -99,6 +100,15 @@ class AddonIndexer(BaseSearchIndexer):
                     'bayesian_rating': {'type': 'double'},
                     'boost': {'type': 'float', 'null_value': 1.0},
                     'category': {'type': 'integer'},
+                    'colors': {
+                        'type': 'nested',
+                        'properties': {
+                            'h': {'type': 'integer'},
+                            's': {'type': 'integer'},
+                            'l': {'type': 'integer'},
+                            'ratio': {'type': 'double'},
+                        },
+                    },
                     'contributions': {'type': 'text'},
                     'created': {'type': 'date'},
                     'current_version': version_mapping,
@@ -284,6 +294,7 @@ class AddonIndexer(BaseSearchIndexer):
                  'status', 'type', 'view_source', 'weekly_downloads')
         data = {attr: getattr(obj, attr) for attr in attrs}
 
+        data['colors'] = None
         if obj.type == amo.ADDON_PERSONA:
             # Personas are compatible with all platforms. They don't have files
             # so we have to fake the info to be consistent with the rest of the
@@ -320,6 +331,12 @@ class AddonIndexer(BaseSearchIndexer):
                                      obj.current_version.supported_platforms]
             data['has_theme_rereview'] = None
 
+            # Extract dominant colors from static themes.
+            if obj.type == amo.ADDON_STATICTHEME:
+                first_preview = obj.current_previews.first()
+                if first_preview:
+                    data['colors'] = first_preview.colors
+
         data['app'] = [app.id for app in obj.compatible_apps.keys()]
         # Boost by the number of users on a logarithmic scale.
         data['boost'] = float(data['average_daily_users'] ** .2)

src/olympia/addons/management/commands/process_addons.py

@@ -9,6 +9,7 @@ from olympia.addons.models import Addon
 from olympia.addons.tasks import (
     add_dynamic_theme_tag, add_firefox57_tag, bump_appver_for_legacy_addons,
     disable_legacy_files,
+    extract_colors_from_static_themes,
     find_inconsistencies_between_es_and_db,
     migrate_legacy_dictionaries_to_webextension,
     migrate_lwts_to_static_themes,
@@ -115,6 +116,10 @@ tasks = {
         ],
         'pre': lambda values_qs: values_qs.distinct(),
     },
+    'extract_colors_from_static_themes': {
+        'method': extract_colors_from_static_themes,
+        'qs': [Q(type=amo.ADDON_STATICTHEME)]
+    }
 }
 
 

src/olympia/addons/models.py

@@ -1350,7 +1350,7 @@ class Addon(OnChangeMixin, ModelBase):
         if self.has_per_version_previews:
             if self.current_version:
                 return self.current_version.previews.all()
-            return []
+            return VersionPreview.objects.none()
         else:
             return self._all_previews
 

src/olympia/addons/tasks.py

@@ -28,7 +28,8 @@ from olympia.amo.decorators import set_modified_on, use_primary_db
 from olympia.amo.storage_utils import rm_stored_dir
 from olympia.amo.templatetags.jinja_helpers import user_media_path
 from olympia.amo.utils import (
-    ImageCheck, LocalFileStorage, cache_ns_key, pngcrush_image, StopWatch)
+    ImageCheck, LocalFileStorage, cache_ns_key, extract_colors_from_image,
+    pngcrush_image, StopWatch)
 from olympia.applications.models import AppVersion
 from olympia.constants.categories import CATEGORIES
 from olympia.constants.licenses import (
@@ -903,3 +904,21 @@ def remove_amo_links_in_url_fields(ids, **kw):
                 ).update(localized_string=u'', localized_string_clean=u'')
             if settings.DOMAIN.lower() in addon.contributions.lower():
                 addon.update(contributions=u'')
+
+
+@task
+@use_primary_db
+def extract_colors_from_static_themes(ids, **kw):
+    """Extract and store colors from existing static themes."""
+    log.info('Extracting static themes colors %d-%d [%d].', ids[0], ids[-1],
+             len(ids))
+    addons = Addon.objects.filter(id__in=ids)
+    extracted = []
+    for addon in addons:
+        first_preview = addon.current_previews.first()
+        if first_preview and not first_preview.colors:
+            colors = extract_colors_from_image(first_preview.thumbnail_path)
+            addon.current_previews.update(colors=colors)
+            extracted.append(addon.pk)
+    if extracted:
+        index_addons.delay(extracted)

src/olympia/addons/tests/test_commands.py

@@ -19,7 +19,7 @@ from olympia.amo.tests import (
 from olympia.applications.models import AppVersion
 from olympia.files.models import FileValidation, WebextPermission
 from olympia.reviewers.models import AutoApprovalSummary, ReviewerScore
-from olympia.versions.models import ApplicationsVersions
+from olympia.versions.models import ApplicationsVersions, VersionPreview
 
 
 # Where to monkeypatch "lib.crypto.tasks.sign_addons" so it's correctly mocked.
@@ -751,3 +751,19 @@ class TestRemoveAMOLinksInURLFields(TestCase):
         translation.activate('fr')
         addon.reload()
         assert addon.support_url == u''
+
+
+class TestExtractColorsFromStaticThemes(TestCase):
+    @mock.patch('olympia.addons.tasks.extract_colors_from_image')
+    def test_basic(self, extract_colors_from_image_mock):
+        addon = addon_factory(type=amo.ADDON_STATICTHEME)
+        preview = VersionPreview.objects.create(version=addon.current_version)
+        extract_colors_from_image_mock.return_value = [
+            {'h': 4, 's': 8, 'l': 15, 'ratio': .16}
+        ]
+        call_command(
+            'process_addons', task='extract_colors_from_static_themes')
+        preview.reload()
+        assert preview.colors == [
+            {'h': 4, 's': 8, 'l': 15, 'ratio': .16}
+        ]

src/olympia/addons/tests/test_indexers.py

@@ -50,8 +50,8 @@ class TestAddonIndexer(TestCase):
         # exist on the model, or it has a different name, or the value we need
         # to store in ES differs from the one in the db.
         complex_fields = [
-            'app', 'boost', 'category', 'current_version', 'description',
-            'featured_for', 'has_eula', 'has_privacy_policy',
+            'app', 'boost', 'category', 'colors', 'current_version',
+            'description', 'featured_for', 'has_eula', 'has_privacy_policy',
             'has_theme_rereview', 'is_featured', 'listed_authors', 'name',
             'platforms', 'previews', 'public_stats', 'ratings', 'summary',
             'tags',
@@ -191,6 +191,7 @@ class TestAddonIndexer(TestCase):
         assert extracted['has_eula'] is True
         assert extracted['has_privacy_policy'] is True
         assert extracted['is_featured'] is False
+        assert extracted['colors'] is None
 
     def test_extract_is_featured(self):
         collection = collection_factory()
@@ -482,15 +483,36 @@ class TestAddonIndexer(TestCase):
         self.addon.update(type=amo.ADDON_STATICTHEME)
         current_preview = VersionPreview.objects.create(
             version=self.addon.current_version,
-            sizes={'thumbnail': [56, 78], 'image': [91, 234]})
+            colors=[{'h': 1, 's': 2, 'l': 3, 'ratio': 0.9}],
+            sizes={'thumbnail': [56, 78], 'image': [91, 234]}, position=1)
+        second_preview = VersionPreview.objects.create(
+            version=self.addon.current_version,
+            sizes={'thumbnail': [12, 34], 'image': [56, 78]}, position=2)
         extracted = self._extract()
         assert extracted['previews']
-        assert len(extracted['previews']) == 1
+        assert len(extracted['previews']) == 2
         assert 'caption_translations' not in extracted['previews'][0]
         assert extracted['previews'][0]['id'] == current_preview.pk
         assert extracted['previews'][0]['modified'] == current_preview.modified
         assert extracted['previews'][0]['sizes'] == current_preview.sizes == {
             'thumbnail': [56, 78], 'image': [91, 234]}
+        assert 'caption_translations' not in extracted['previews'][1]
+        assert extracted['previews'][1]['id'] == second_preview.pk
+        assert extracted['previews'][1]['modified'] == second_preview.modified
+        assert extracted['previews'][1]['sizes'] == second_preview.sizes == {
+            'thumbnail': [12, 34], 'image': [56, 78]}
+
+        # Make sure we extract colors from the first preview.
+        assert extracted['colors'] == [{'h': 1, 's': 2, 'l': 3, 'ratio': 0.9}]
+
+    def test_extract_staticthemes_somehow_no_previews(self):
+        # Extracting a static theme with no previews should not fail.
+        self.addon.update(type=amo.ADDON_STATICTHEME)
+
+        extracted = self._extract()
+        assert extracted['id'] == self.addon.pk
+        assert extracted['previews'] == []
+        assert extracted['colors'] is None
 
 
 class TestAddonIndexerWithES(ESTestCase):

src/olympia/addons/tests/test_views.py

@@ -2215,7 +2215,7 @@ class TestAddonSearchView(ESTestCase):
         qset = view.get_queryset()
 
         assert set(qset.to_dict()['_source']['excludes']) == set(
-            ('*.raw', 'boost', 'hotness', 'name', 'description',
+            ('*.raw', 'boost', 'colors', 'hotness', 'name', 'description',
              'name_l10n_*', 'description_l10n_*', 'summary', 'summary_l10n_*')
         )
 

src/olympia/amo/tests/test_utils_.py

@@ -1,5 +1,6 @@
 import collections
 import datetime
+import os.path
 import tempfile
 
 from django.conf import settings
@@ -15,8 +16,9 @@ from olympia import amo
 from olympia.addons.models import Addon
 from olympia.amo.tests import TestCase, addon_factory
 from olympia.amo.utils import (
-    attach_trans_dict, get_locale_from_lang, pngcrush_image,
-    translations_for_field, utc_millesecs_from_epoch, walkfiles)
+    attach_trans_dict, extract_colors_from_image, get_locale_from_lang,
+    pngcrush_image, translations_for_field, utc_millesecs_from_epoch,
+    walkfiles)
 from olympia.versions.models import Version
 
 
@@ -241,3 +243,18 @@ def test_utc_millesecs_from_epoch():
     new_timestamp = utc_millesecs_from_epoch(
         future_now + datetime.timedelta(milliseconds=42))
     assert new_timestamp == timestamp + 42
+
+
+def test_extract_colors_from_image():
+    path = os.path.join(
+        settings.ROOT,
+        'src/olympia/versions/tests/static_themes/weta.png')
+    expected = [
+        {'h': 45, 'l': 158, 'ratio': 0.40547158773994313, 's': 34},
+        {'h': 44, 'l': 94, 'ratio': 0.2812929380875291, 's': 28},
+        {'h': 68, 'l': 99, 'ratio': 0.13200103391513734, 's': 19},
+        {'h': 43, 'l': 177, 'ratio': 0.06251105336906689, 's': 93},
+        {'h': 47, 'l': 115, 'ratio': 0.05938209966397758, 's': 60},
+        {'h': 40, 'l': 201, 'ratio': 0.05934128722434598, 's': 83}
+    ]
+    assert extract_colors_from_image(path) == expected

src/olympia/amo/utils.py

@@ -35,6 +35,7 @@ from django.utils.encoding import force_bytes, force_text
 from django.utils.http import _urlparse as django_urlparse, quote_etag
 
 import bleach
+import colorgram
 import html5lib
 import jinja2
 import pytz
@@ -1002,6 +1003,20 @@ def utc_millesecs_from_epoch(for_datetime=None):
     return int(seconds * 1000)
 
 
+def extract_colors_from_image(path):
+    try:
+        image_colors = colorgram.extract(path, 6)
+        colors = [{
+            'h': color.hsl.h,
+            's': color.hsl.s,
+            'l': color.hsl.l,
+            'ratio': color.proportion
+        } for color in image_colors]
+    except IOError:
+        colors = None
+    return colors
+
+
 class AMOJSONEncoder(JSONEncoder):
     def default(self, obj):
         if isinstance(obj, Translation):

src/olympia/migrations/1067-add-version-preview-color.sql

@@ -0,0 +1 @@
+ALTER TABLE `version_previews` ADD COLUMN `colors` longtext NULL;

src/olympia/search/filters.py

@@ -1,6 +1,7 @@
 from django.utils import translation
 from django.utils.translation import ugettext
 
+import colorgram
 from elasticsearch_dsl import Q, query
 from rest_framework import serializers
 from rest_framework.filters import BaseFilterBackend
@@ -303,6 +304,58 @@ class AddonFeaturedQueryParam(AddonQueryParam):
             filter=clauses))]
 
 
+class AddonColorQueryParam(AddonQueryParam):
+    query_param = 'color'
+
+    def convert_to_hsl(self, hexvalue):
+        # The API is receiving color as a hex string. We store colors in HSL
+        # as colorgram generates it (which is on a 0 to 255 scale for each
+        # component), so some conversion is necessary.
+        if len(hexvalue) == 3:
+            hexvalue = ''.join(2 * c for c in hexvalue)
+        try:
+            rgb = tuple(bytearray.fromhex(hexvalue))
+        except ValueError:
+            rgb = (0, 0, 0)
+        return colorgram.colorgram.hsl(*rgb)
+
+    def get_value(self):
+        color = self.request.GET.get(self.query_param, '')
+        return self.convert_to_hsl(color.upper().lstrip('#'))
+
+    def get_es_query(self):
+        hsl = self.get_value()
+        # If we're given a color with a very low saturation, the user is
+        # searching for a black/white/grey and we need to take saturation and
+        # lightness into consideration, but ignore hue.
+        if hsl[1] <= 5:  # 2.5% saturation or lower.
+            clauses = [
+                Q('range', **{'colors.s': {
+                    'lte': 5,
+                }}),
+                Q('range', **{'colors.l': {
+                    'gte': max(min(hsl[2] - 64, 255), 0),
+                    'lte': max(min(hsl[2] + 64, 255), 0),
+                }})
+            ]
+        else:
+            # Otherwise, we want to do the opposite and just try to match the
+            # hue. The idea is to keep the UI simple, presenting the user with
+            # a limited set of colors that still allows them to find all
+            # themes.
+            clauses = [
+                Q('range', **{'colors.h': {
+                    'gte': (hsl[0] - 26) % 255,  # 10% less hue
+                    'lte': (hsl[0] + 26) % 255,  # 10% more hue
+                }}),
+            ]
+        # In any case, the color we're looking for needs to be present in at
+        # least 20% of the image.
+        clauses.append(Q('range', **{'colors.ratio': {'gte': 0.20}}))
+
+        return [Q('nested', path='colors', query=query.Bool(filter=clauses))]
+
+
 class SearchQueryFilter(BaseFilterBackend):
     """
     A django-rest-framework filter backend that performs an ES query according
@@ -632,6 +685,7 @@ class SearchParameterFilter(BaseFilterBackend):
         AddonPlatformQueryParam,
         AddonTagQueryParam,
         AddonTypeQueryParam,
+        AddonColorQueryParam,
     ]
 
     def get_applicable_clauses(self, request):
@@ -649,7 +703,8 @@ class SearchParameterFilter(BaseFilterBackend):
 
     def filter_queryset(self, request, qs, view):
         filters = self.get_applicable_clauses(request)
-        return qs.query(query.Bool(filter=filters)) if filters else qs
+        qs = qs.query(query.Bool(filter=filters)) if filters else qs
+        return qs
 
 
 class ReviewedContentFilter(BaseFilterBackend):

src/olympia/search/tests/test_filters.py

@@ -784,6 +784,49 @@ class TestSearchParameterFilter(FilterTestsBase):
         assert len(inner) == 1
         assert {'terms': {'featured_for.locales': ['fr', 'ALL']}} in inner
 
+    def test_search_by_color(self):
+        qs = self._filter(data={'color': 'ff0000'})
+        filter_ = qs['query']['bool']['filter']
+        assert len(filter_) == 1
+        inner = filter_[0]['nested']['query']['bool']['filter']
+        assert len(inner) == 2
+        assert inner == [
+            {'range': {'colors.h': {'gte': 229, 'lte': 26}}},
+            {'range': {'colors.ratio': {'gte': 0.2}}},
+        ]
+
+        qs = self._filter(data={'color': '#00ffff'})
+        filter_ = qs['query']['bool']['filter']
+        assert len(filter_) == 1
+        inner = filter_[0]['nested']['query']['bool']['filter']
+        assert len(inner) == 2
+        assert inner == [
+            {'range': {'colors.h': {'gte': 101, 'lte': 153}}},
+            {'range': {'colors.ratio': {'gte': 0.2}}},
+        ]
+
+        qs = self._filter(data={'color': '#f6f6f6'})
+        filter_ = qs['query']['bool']['filter']
+        assert len(filter_) == 1
+        inner = filter_[0]['nested']['query']['bool']['filter']
+        assert len(inner) == 3
+        assert inner == [
+            {'range': {'colors.s': {'lte': 5}}},
+            {'range': {'colors.l': {'gte': 182, 'lte': 255}}},
+            {'range': {'colors.ratio': {'gte': 0.2}}},
+        ]
+
+        qs = self._filter(data={'color': '333'})
+        filter_ = qs['query']['bool']['filter']
+        assert len(filter_) == 1
+        inner = filter_[0]['nested']['query']['bool']['filter']
+        assert len(inner) == 3
+        assert inner == [
+            {'range': {'colors.s': {'lte': 5}}},
+            {'range': {'colors.l': {'gte': 0, 'lte': 115}}},
+            {'range': {'colors.ratio': {'gte': 0.2}}},
+        ]
+
 
 class TestCombinedFilter(FilterTestsBase):
     """

src/olympia/versions/models.py

@@ -643,6 +643,7 @@ class VersionPreview(BasePreview, ModelBase):
     version = models.ForeignKey(Version, related_name='previews')
     position = models.IntegerField(default=0)
     sizes = JSONField(default={})
+    colors = JSONField(default=None, null=True)
     media_folder = 'version-previews'
 
     class Meta:

src/olympia/versions/tasks.py

@@ -10,7 +10,7 @@ from olympia import amo
 from olympia.addons.tasks import index_addons
 from olympia.amo.celery import task
 from olympia.amo.decorators import use_primary_db
-from olympia.amo.utils import pngcrush_image
+from olympia.amo.utils import extract_colors_from_image, pngcrush_image
 from olympia.devhub.tasks import resize_image
 from olympia.files.models import File
 from olympia.files.utils import get_background_images
@@ -70,6 +70,7 @@ def generate_static_theme_preview(theme_manifest, version_pk):
     sizes = sorted(
         amo.THEME_PREVIEW_SIZES.values(),
         lambda x, y: x['position'] - y['position'])
+    colors = None
     for size in sizes:
         # Create a Preview for this size.
         preview = VersionPreview.objects.create(
@@ -81,10 +82,19 @@ def generate_static_theme_preview(theme_manifest, version_pk):
             resize_image(
                 preview.image_path, preview.thumbnail_path, size['thumbnail'])
             pngcrush_image(preview.image_path)
-            preview_sizes = {}
-            preview_sizes['image'] = size['full']
-            preview_sizes['thumbnail'] = size['thumbnail']
-            preview.update(sizes=preview_sizes)
+            # Extract colors once and store it for all previews.
+            # Use the thumbnail for extra speed, we don't need to be super
+            # accurate.
+            if colors is None:
+                colors = extract_colors_from_image(preview.thumbnail_path)
+            data = {
+                'sizes': {
+                    'image': size['full'],
+                    'thumbnail': size['thumbnail'],
+                },
+                'colors': colors,
+            }
+            preview.update(**data)
     addon_id = Version.objects.values_list(
         'addon_id', flat=True).get(id=version_pk)
     index_addons.delay([addon_id])

src/olympia/versions/tests/test_tasks.py

@@ -58,10 +58,12 @@ def check_preview(preview_instance, theme_size_constant, write_svg_mock_args,
     assert thumb_path == preview_instance.thumbnail_path
     assert thumb_size == theme_size_constant['thumbnail']
     assert png_crush_mock_args[0] == preview_instance.image_path
+    assert preview_instance.colors
 
 
 @pytest.mark.django_db
 @mock.patch('olympia.versions.tasks.index_addons.delay')
+@mock.patch('olympia.versions.tasks.extract_colors_from_image')
 @mock.patch('olympia.versions.tasks.pngcrush_image')
 @mock.patch('olympia.versions.tasks.resize_image')
 @mock.patch('olympia.versions.tasks.write_svg_to_png')
@@ -81,9 +83,12 @@ def check_preview(preview_instance, theme_size_constant, write_svg_mock_args,
 )
 def test_generate_static_theme_preview(
         write_svg_to_png_mock, resize_image_mock, pngcrush_image_mock,
-        index_addons_mock,
+        extract_colors_from_image_mock, index_addons_mock,
         header_url, header_height, preserve_aspect_ratio, mimetype, valid_img):
     write_svg_to_png_mock.return_value = True
+    extract_colors_from_image_mock.return_value = [
+        {'h': 9, 's': 8, 'l': 7, 'ratio': 0.6}
+    ]
     theme_manifest = {
         "images": {
         },
@@ -164,13 +169,17 @@ def test_generate_static_theme_preview(
 
 @pytest.mark.django_db
 @mock.patch('olympia.versions.tasks.index_addons.delay')
+@mock.patch('olympia.versions.tasks.extract_colors_from_image')
 @mock.patch('olympia.versions.tasks.pngcrush_image')
 @mock.patch('olympia.versions.tasks.resize_image')
 @mock.patch('olympia.versions.tasks.write_svg_to_png')
 def test_generate_static_theme_preview_with_chrome_properties(
         write_svg_to_png_mock, resize_image_mock, pngcrush_image_mock,
-        index_addons_mock):
+        extract_colors_from_image_mock, index_addons_mock):
     write_svg_to_png_mock.return_value = True
+    extract_colors_from_image_mock.return_value = [
+        {'h': 9, 's': 8, 'l': 7, 'ratio': 0.6}
+    ]
     theme_manifest = {
         "images": {
             "theme_frame": "transparent.gif"
@@ -274,13 +283,17 @@ def check_render_additional(svg_content, inner_svg_width, colors):
 
 @pytest.mark.django_db
 @mock.patch('olympia.versions.tasks.index_addons.delay')
+@mock.patch('olympia.versions.tasks.extract_colors_from_image')
 @mock.patch('olympia.versions.tasks.pngcrush_image')
 @mock.patch('olympia.versions.tasks.resize_image')
 @mock.patch('olympia.versions.tasks.write_svg_to_png')
 def test_generate_preview_with_additional_backgrounds(
         write_svg_to_png_mock, resize_image_mock, pngcrush_image_mock,
-        index_addons_mock):
+        extract_colors_from_image_mock, index_addons_mock):
     write_svg_to_png_mock.return_value = True
+    extract_colors_from_image_mock.return_value = [
+        {'h': 9, 's': 8, 'l': 7, 'ratio': 0.6}
+    ]
 
     theme_manifest = {
         "images": {