Add color_extractor service (#42129)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

CODEOWNERS

@@ -76,6 +76,7 @@ homeassistant/components/cisco_mobility_express/* @fbradyirl
 homeassistant/components/cisco_webex_teams/* @fbradyirl
 homeassistant/components/cloud/* @home-assistant/cloud
 homeassistant/components/cloudflare/* @ludeeus @ctalkington
+homeassistant/components/color_extractor/* @GenericStudent
 homeassistant/components/comfoconnect/* @michaelarnauts
 homeassistant/components/config/* @home-assistant/core
 homeassistant/components/configurator/* @home-assistant/core

homeassistant/components/color_extractor/__init__.py

@@ -0,0 +1,150 @@
+"""Module for color_extractor (RGB extraction from images) component."""
+import asyncio
+import io
+import logging
+
+from PIL import UnidentifiedImageError
+import aiohttp
+import async_timeout
+from colorthief import ColorThief
+import voluptuous as vol
+
+from homeassistant.components.color_extractor.const import (
+    ATTR_PATH,
+    ATTR_URL,
+    DOMAIN,
+    SERVICE_TURN_ON,
+)
+from homeassistant.components.light import (
+    ATTR_RGB_COLOR,
+    DOMAIN as LIGHT_DOMAIN,
+    LIGHT_TURN_ON_SCHEMA,
+    SERVICE_TURN_ON as LIGHT_SERVICE_TURN_ON,
+)
+from homeassistant.helpers import aiohttp_client
+import homeassistant.helpers.config_validation as cv
+
+_LOGGER = logging.getLogger(__name__)
+
+# Extend the existing light.turn_on service schema
+SERVICE_SCHEMA = cv.make_entity_service_schema(
+    {
+        **LIGHT_TURN_ON_SCHEMA,
+        vol.Exclusive(ATTR_PATH, "color_extractor"): cv.isfile,
+        vol.Exclusive(ATTR_URL, "color_extractor"): cv.url,
+    }
+)
+
+
+def _get_file(file_path):
+    """Get a PIL acceptable input file reference.
+
+    Allows us to mock patch during testing to make BytesIO stream.
+    """
+    return file_path
+
+
+def _get_color(file_handler) -> tuple:
+    """Given an image file, extract the predominant color from it."""
+    color_thief = ColorThief(file_handler)
+
+    # get_color returns a SINGLE RGB value for the given image
+    color = color_thief.get_color(quality=1)
+
+    _LOGGER.debug("Extracted RGB color %s from image", color)
+
+    return color
+
+
+async def async_setup(hass, hass_config):
+    """Set up services for color_extractor integration."""
+
+    async def async_handle_service(service_call):
+        """Decide which color_extractor method to call based on service."""
+        service_data = dict(service_call.data)
+
+        if ATTR_URL not in service_data and ATTR_PATH not in service_data:
+            _LOGGER.error("Missing either required %s or %s key", ATTR_URL, ATTR_PATH)
+            return
+
+        try:
+            if ATTR_URL in service_data:
+                image_type = "URL"
+                image_reference = service_data.pop(ATTR_URL)
+                color = await async_extract_color_from_url(image_reference)
+
+            elif ATTR_PATH in service_data:
+                image_type = "file path"
+                image_reference = service_data.pop(ATTR_PATH)
+                color = await hass.async_add_executor_job(
+                    extract_color_from_path, image_reference
+                )
+
+        except UnidentifiedImageError as ex:
+            _LOGGER.error(
+                "Bad image from %s '%s' provided, are you sure it's an image? %s",
+                image_type,
+                image_reference,
+                ex,
+            )
+            return
+
+        if color:
+            service_data[ATTR_RGB_COLOR] = color
+
+            await hass.services.async_call(
+                LIGHT_DOMAIN, LIGHT_SERVICE_TURN_ON, service_data, blocking=True
+            )
+
+    hass.services.async_register(
+        DOMAIN,
+        SERVICE_TURN_ON,
+        async_handle_service,
+        schema=SERVICE_SCHEMA,
+    )
+
+    async def async_extract_color_from_url(url):
+        """Handle call for URL based image."""
+        if not hass.config.is_allowed_external_url(url):
+            _LOGGER.error(
+                "External URL '%s' is not allowed, please add to 'allowlist_external_urls'",
+                url,
+            )
+            return
+
+        _LOGGER.debug("Getting predominant RGB from image URL '%s'", url)
+
+        # Download the image into a buffer for ColorThief to check against
+        try:
+            session = aiohttp_client.async_get_clientsession(hass)
+
+            with async_timeout.timeout(10):
+                response = await session.get(url)
+
+        except (asyncio.TimeoutError, aiohttp.ClientError) as err:
+            _LOGGER.error("Failed to get ColorThief image due to HTTPError: %s", err)
+            return
+
+        content = await response.content.read()
+
+        with io.BytesIO(content) as _file:
+            _file.name = "color_extractor.jpg"
+            _file.seek(0)
+
+            return _get_color(_file)
+
+    def extract_color_from_path(file_path):
+        """Handle call for local file based image."""
+        if not hass.config.is_allowed_path(file_path):
+            _LOGGER.error(
+                "File path '%s' is not allowed, please add to 'allowlist_external_dirs'",
+                file_path,
+            )
+            return
+
+        _LOGGER.debug("Getting predominant RGB from file path '%s'", file_path)
+
+        _file = _get_file(file_path)
+        return _get_color(_file)
+
+    return True

homeassistant/components/color_extractor/const.py

@@ -0,0 +1,7 @@
+"""Constants for the color_extractor component."""
+ATTR_PATH = "color_extract_path"
+ATTR_URL = "color_extract_url"
+
+DOMAIN = "color_extractor"
+
+SERVICE_TURN_ON = "turn_on"

homeassistant/components/color_extractor/manifest.json

@@ -0,0 +1,9 @@
+{
+    "domain": "color_extractor",
+    "name": "ColorExtractor",
+    "config_flow": false,
+    "documentation": "https://www.home-assistant.io/integrations/color_extractor",
+    "requirements": ["colorthief==0.2.1"],
+    "codeowners": ["@GenericStudent"]
+  }
+  

homeassistant/components/color_extractor/services.yaml

@@ -0,0 +1,12 @@
+turn_on:
+  description: Set the light RGB to the predominant color found in the image provided by url or file path.
+  fields:
+    color_extract_url:
+      description: The URL of the image we want to extract RGB values from. Must be allowed in allowlist_external_urls.
+      example: https://www.example.com/images/logo.png
+    color_extract_path:
+      description: The full system path to the image we want to extract RGB values from. Must be allowed in allowlist_external_dirs.
+      example: /opt/images/logo.png
+    entity_id:
+      description: The entity we want to set our RGB color on.
+      example: "light.living_room_shelves"

requirements_all.txt

@@ -433,6 +433,9 @@ coinmarketcap==5.0.3
 # homeassistant.scripts.check_config
 colorlog==4.2.1
 
+# homeassistant.components.color_extractor
+colorthief==0.2.1
+
 # homeassistant.components.concord232
 concord232==0.15
 

requirements_test_all.txt

@@ -224,6 +224,9 @@ coinmarketcap==5.0.3
 # homeassistant.scripts.check_config
 colorlog==4.2.1
 
+# homeassistant.components.color_extractor
+colorthief==0.2.1
+
 # homeassistant.components.eddystone_temperature
 # homeassistant.components.eq3btsmart
 # homeassistant.components.xiaomi_miio

tests/components/color_extractor/__init__.py

@@ -0,0 +1 @@
+"""Tests for the color_extractor component."""

tests/components/color_extractor/test_service.py

@@ -0,0 +1,315 @@
+"""Tests for color_extractor component service calls."""
+import base64
+import io
+
+import aiohttp
+import pytest
+
+from homeassistant.components.color_extractor import (
+    ATTR_PATH,
+    ATTR_URL,
+    DOMAIN,
+    SERVICE_TURN_ON,
+)
+from homeassistant.components.light import (
+    ATTR_BRIGHTNESS,
+    ATTR_BRIGHTNESS_PCT,
+    ATTR_RGB_COLOR,
+    DOMAIN as LIGHT_DOMAIN,
+    SERVICE_TURN_OFF as LIGHT_SERVICE_TURN_OFF,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
+from homeassistant.setup import async_setup_component
+import homeassistant.util.color as color_util
+
+from tests.async_mock import Mock, mock_open, patch
+from tests.common import load_fixture
+
+LIGHT_ENTITY = "light.kitchen_lights"
+CLOSE_THRESHOLD = 10
+
+
+def _close_enough(actual_rgb, testing_rgb):
+    """Validate the given RGB value is in acceptable tolerance."""
+    # Convert the given RGB values to hue / saturation and then back again
+    # as it wasn't reading the same RGB value set against it.
+    actual_hs = color_util.color_RGB_to_hs(*actual_rgb)
+    actual_rgb = color_util.color_hs_to_RGB(*actual_hs)
+
+    testing_hs = color_util.color_RGB_to_hs(*testing_rgb)
+    testing_rgb = color_util.color_hs_to_RGB(*testing_hs)
+
+    actual_red, actual_green, actual_blue = actual_rgb
+    testing_red, testing_green, testing_blue = testing_rgb
+
+    r_diff = abs(actual_red - testing_red)
+    g_diff = abs(actual_green - testing_green)
+    b_diff = abs(actual_blue - testing_blue)
+
+    return (
+        r_diff <= CLOSE_THRESHOLD
+        and g_diff <= CLOSE_THRESHOLD
+        and b_diff <= CLOSE_THRESHOLD
+    )
+
+
+@pytest.fixture(autouse=True)
+async def setup_light(hass):
+    """Configure our light component to work against for testing."""
+    assert await async_setup_component(
+        hass, LIGHT_DOMAIN, {LIGHT_DOMAIN: {"platform": "demo"}}
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+
+    # Validate starting values
+    assert state.state == STATE_ON
+    assert state.attributes.get(ATTR_BRIGHTNESS) == 180
+    assert state.attributes.get(ATTR_RGB_COLOR) == (255, 63, 111)
+
+    await hass.services.async_call(
+        LIGHT_DOMAIN,
+        LIGHT_SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: LIGHT_ENTITY},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get(LIGHT_ENTITY)
+
+    assert state
+    assert state.state == STATE_OFF
+
+
+async def test_missing_url_and_path(hass):
+    """Test that nothing happens when url and path are missing."""
+    # Load our color_extractor component
+    await async_setup_component(
+        hass,
+        DOMAIN,
+        {},
+    )
+    await hass.async_block_till_done()
+
+    # Validate pre service call
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+    # Missing url and path attributes, should cause error log
+    service_data = {
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+    }
+
+    await hass.services.async_call(DOMAIN, SERVICE_TURN_ON, service_data, blocking=True)
+    await hass.async_block_till_done()
+
+    # check light is still off, unchanged due to bad parameters on service call
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+
+async def _async_load_color_extractor_url(hass, service_data):
+    # Load our color_extractor component
+    await async_setup_component(
+        hass,
+        DOMAIN,
+        {},
+    )
+    await hass.async_block_till_done()
+
+    # Validate pre service call
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+    # Call the shared service, our above mock should return the base64 decoded fixture 1x1 pixel
+    assert await hass.services.async_call(
+        DOMAIN, SERVICE_TURN_ON, service_data, blocking=True
+    )
+
+    await hass.async_block_till_done()
+
+
+async def test_url_success(hass, aioclient_mock):
+    """Test that a successful image GET translate to light RGB."""
+    service_data = {
+        ATTR_URL: "http://example.com/images/logo.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+        # Standard light service data which we pass
+        ATTR_BRIGHTNESS_PCT: 50,
+    }
+
+    # Mock the HTTP Response with a base64 encoded 1x1 pixel
+    aioclient_mock.get(
+        url=service_data[ATTR_URL],
+        content=base64.b64decode(load_fixture("color_extractor_url.txt")),
+    )
+
+    # Allow access to this URL using the proper mechanism
+    hass.config.allowlist_external_urls.add("http://example.com/images/")
+
+    await _async_load_color_extractor_url(hass, service_data)
+
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+
+    # Ensure we turned it on
+    assert state.state == STATE_ON
+
+    # Brightness has changed, optional service call field
+    assert state.attributes[ATTR_BRIGHTNESS] == 128
+
+    # Ensure the RGB values are correct
+    assert _close_enough(state.attributes[ATTR_RGB_COLOR], (50, 100, 150))
+
+
+async def test_url_not_allowed(hass, aioclient_mock):
+    """Test that a not allowed external URL fails to turn light on."""
+    service_data = {
+        ATTR_URL: "http://denied.com/images/logo.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+    }
+
+    await _async_load_color_extractor_url(hass, service_data)
+
+    # Light has not been modified due to failure
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+
+async def test_url_exception(hass, aioclient_mock):
+    """Test that a HTTPError fails to turn light on."""
+    service_data = {
+        ATTR_URL: "http://example.com/images/logo.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+    }
+
+    # Don't let the URL not being allowed sway our exception test
+    hass.config.allowlist_external_urls.add("http://example.com/images/")
+
+    # Mock the HTTP Response with an HTTPError
+    aioclient_mock.get(url=service_data[ATTR_URL], exc=aiohttp.ClientError)
+
+    await _async_load_color_extractor_url(hass, service_data)
+
+    # Light has not been modified due to failure
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+
+async def test_url_error(hass, aioclient_mock):
+    """Test that a HTTP Error (non 200) doesn't turn light on."""
+    service_data = {
+        ATTR_URL: "http://example.com/images/logo.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+    }
+
+    # Don't let the URL not being allowed sway our exception test
+    hass.config.allowlist_external_urls.add("http://example.com/images/")
+
+    # Mock the HTTP Response with a 400 Bad Request error
+    aioclient_mock.get(url=service_data[ATTR_URL], status=400)
+
+    await _async_load_color_extractor_url(hass, service_data)
+
+    # Light has not been modified due to failure
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+
+@patch(
+    "builtins.open",
+    mock_open(read_data=base64.b64decode(load_fixture("color_extractor_file.txt"))),
+    create=True,
+)
+def _get_file_mock(file_path):
+    """Convert file to BytesIO for testing due to PIL UnidentifiedImageError."""
+    _file = None
+
+    with open(file_path) as file_handler:
+        _file = io.BytesIO(file_handler.read())
+
+    _file.name = "color_extractor.jpg"
+    _file.seek(0)
+
+    return _file
+
+
+@patch("os.path.isfile", Mock(return_value=True))
+@patch("os.access", Mock(return_value=True))
+async def test_file(hass):
+    """Test that the file only service reads a file and translates to light RGB."""
+    service_data = {
+        ATTR_PATH: "/opt/image.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+        # Standard light service data which we pass
+        ATTR_BRIGHTNESS_PCT: 100,
+    }
+
+    # Add our /opt/ path to the allowed list of paths
+    hass.config.allowlist_external_dirs.add("/opt/")
+
+    await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    # Verify pre service check
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+    # Mock the file handler read with our 1x1 base64 encoded fixture image
+    with patch("homeassistant.components.color_extractor._get_file", _get_file_mock):
+        await hass.services.async_call(DOMAIN, SERVICE_TURN_ON, service_data)
+        await hass.async_block_till_done()
+
+    state = hass.states.get(LIGHT_ENTITY)
+
+    assert state
+
+    # Ensure we turned it on
+    assert state.state == STATE_ON
+
+    # And set the brightness
+    assert state.attributes[ATTR_BRIGHTNESS] == 255
+
+    # Ensure the RGB values are correct
+    assert _close_enough(state.attributes[ATTR_RGB_COLOR], (25, 75, 125))
+
+
+@patch("os.path.isfile", Mock(return_value=True))
+@patch("os.access", Mock(return_value=True))
+async def test_file_denied_dir(hass):
+    """Test that the file only service fails to read an image in a dir not explicitly allowed."""
+    service_data = {
+        ATTR_PATH: "/path/to/a/dir/not/allowed/image.png",
+        ATTR_ENTITY_ID: LIGHT_ENTITY,
+        # Standard light service data which we pass
+        ATTR_BRIGHTNESS_PCT: 100,
+    }
+
+    await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    # Verify pre service check
+    state = hass.states.get(LIGHT_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+
+    # Mock the file handler read with our 1x1 base64 encoded fixture image
+    with patch("homeassistant.components.color_extractor._get_file", _get_file_mock):
+        await hass.services.async_call(DOMAIN, SERVICE_TURN_ON, service_data)
+        await hass.async_block_till_done()
+
+    state = hass.states.get(LIGHT_ENTITY)
+
+    assert state
+
+    # Ensure it's still off due to access error (dir not explicitly allowed)
+    assert state.state == STATE_OFF

tests/fixtures/color_extractor_file.txt

@@ -0,0 +1 @@
+iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQMAAAAl21bKAAAAA1BMVEUZS32JtD5qAAAACklEQVR4nGNiAAAABgADNjd8qAAAAABJRU5ErkJggg==

tests/fixtures/color_extractor_url.txt

@@ -0,0 +1 @@
+iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQMAAAAl21bKAAAAA1BMVEUyZJaYaJT2AAAACklEQVR4nGNiAAAABgADNjd8qAAAAABJRU5ErkJggg==