Make all translations cacheable (#42892)

homeassistant/helpers/translation.py

@@ -1,10 +1,10 @@
 """Translation string lookup helpers."""
 import asyncio
+from collections import ChainMap
 import logging
-from typing import Any, Dict, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
-from homeassistant.const import EVENT_COMPONENT_LOADED
+from homeassistant.core import callback
-from homeassistant.core import Event, callback
 from homeassistant.loader import (
     Integration,
     async_get_config_flows,
@@ -19,6 +19,7 @@ _LOGGER = logging.getLogger(__name__)
 
 TRANSLATION_LOAD_LOCK = "translation_load_lock"
 TRANSLATION_FLATTEN_CACHE = "translation_flatten_cache"
+LOCALE_EN = "en"
 
 
 def recursive_flatten(prefix: Any, data: Dict) -> Dict[str, Any]:
@@ -32,11 +33,6 @@ def recursive_flatten(prefix: Any, data: Dict) -> Dict[str, Any]:
     return output
 
 
-def flatten(data: Dict) -> Dict[str, Any]:
-    """Return a flattened representation of dict data."""
-    return recursive_flatten("", data)
-
-
 @callback
 def component_translation_path(
     component: str, language: str, integration: Integration
@@ -91,7 +87,7 @@ def load_translations_files(
     return loaded
 
 
-def merge_resources(
+def _merge_resources(
     translation_strings: Dict[str, Dict[str, Any]],
     components: Set[str],
     category: str,
@@ -120,57 +116,31 @@ def merge_resources(
         if new_value is None:
             continue
 
-        cur_value = domain_resources.get(category)
+        if isinstance(new_value, dict):
-
+            domain_resources.update(new_value)
-        # If not exists, set value.
-        if cur_value is None:
-            domain_resources[category] = new_value
-
-        # If exists, and a list, append
-        elif isinstance(cur_value, list):
-            cur_value.append(new_value)
-
-        # If exists, and a dict make it a list with 2 entries.
-        else:
-            domain_resources[category] = [cur_value, new_value]
-
-    # Merge all the lists
-    for domain, domain_resources in list(resources.items()):
-        if not isinstance(domain_resources.get(category), list):
-            continue
-
-        merged = {}
-        for entry in domain_resources[category]:
-            if isinstance(entry, dict):
-                merged.update(entry)
         else:
             _LOGGER.error(
                 "An integration providing translations for %s provided invalid data: %s",
                 domain,
-                    entry,
+                new_value,
             )
-        domain_resources[category] = merged
 
-    return {"component": resources}
+    return resources
 
 
-def build_resources(
+def _build_resources(
     translation_strings: Dict[str, Dict[str, Any]],
     components: Set[str],
     category: str,
 ) -> Dict[str, Dict[str, Any]]:
     """Build the resources response for the given components."""
     # Build response
-    resources: Dict[str, Dict[str, Any]] = {}
+    return {
-    for component in components:
+        component: translation_strings[component][category]
-        new_value = translation_strings[component].get(category)
+        for component in components
-
+        if category in translation_strings[component]
-        if new_value is None:
+        and translation_strings[component][category] is not None
-            continue
+    }
-
-        resources[component] = {category: new_value}
-
-    return {"component": resources}
 
 
 async def async_get_component_strings(
@@ -226,35 +196,83 @@ async def async_get_component_strings(
     return translations
 
 
-class FlatCache:
+class _TranslationCache:
     """Cache for flattened translations."""
 
     def __init__(self, hass: HomeAssistantType) -> None:
         """Initialize the cache."""
         self.hass = hass
-        self.cache: Dict[str, Dict[str, Dict[str, str]]] = {}
+        self.loaded: Dict[str, Set[str]] = {}
+        self.cache: Dict[str, Dict[str, Dict[str, Any]]] = {}
+
+    async def async_fetch(
+        self,
+        language: str,
+        category: str,
+        components: Set,
+    ) -> List[Dict[str, Dict[str, Any]]]:
+        """Load resources into the cache."""
+        components_to_load = components - self.loaded.setdefault(language, set())
+
+        if components_to_load:
+            await self._async_load(language, components_to_load)
+
+        cached = self.cache.get(language, {})
+
+        return [cached.get(component, {}).get(category, {}) for component in components]
+
+    async def _async_load(self, language: str, components: Set) -> None:
+        """Populate the cache for a given set of components."""
+        _LOGGER.debug(
+            "Cache miss for %s: %s",
+            language,
+            ", ".join(components),
+        )
+        # Fetch the English resources, as a fallback for missing keys
+        languages = [LOCALE_EN] if language == LOCALE_EN else [LOCALE_EN, language]
+        for translation_strings in await asyncio.gather(
+            *[
+                async_get_component_strings(self.hass, lang, components)
+                for lang in languages
+            ]
+        ):
+            self._build_category_cache(language, components, translation_strings)
+
+        self.loaded[language].update(components)
 
     @callback
-    def async_setup(self) -> None:
+    def _build_category_cache(
-        """Initialize the cache clear listeners."""
+        self,
-        self.hass.bus.async_listen(EVENT_COMPONENT_LOADED, self._async_component_loaded)
+        language: str,
-
+        components: Set,
-    @callback
+        translation_strings: Dict[str, Dict[str, Any]],
-    def _async_component_loaded(self, event: Event) -> None:
-        """Clear cache when a new component is loaded."""
-        self.cache = {}
-
-    @callback
-    def async_get_cache(self, language: str, category: str) -> Optional[Dict[str, str]]:
-        """Get cache."""
-        return self.cache.setdefault(language, {}).get(category)
-
-    @callback
-    def async_set_cache(
-        self, language: str, category: str, data: Dict[str, str]
     ) -> None:
-        """Set cache."""
+        """Extract resources into the cache."""
-        self.cache.setdefault(language, {})[category] = data
+        cached = self.cache.setdefault(language, {})
+        categories: Set[str] = set()
+        for resource in translation_strings.values():
+            categories.update(resource)
+
+        for category in categories:
+            resource_func = (
+                _merge_resources if category == "state" else _build_resources
+            )
+            new_resources = resource_func(translation_strings, components, category)
+
+            for component, resource in new_resources.items():
+                category_cache: Dict[str, Any] = cached.setdefault(
+                    component, {}
+                ).setdefault(category, {})
+
+                if isinstance(resource, dict):
+                    category_cache.update(
+                        recursive_flatten(
+                            f"component.{component}.{category}.",
+                            resource,
+                        )
+                    )
+                else:
+                    category_cache[f"component.{component}.{category}"] = resource
 
 
 @bind_hass
@@ -271,71 +289,22 @@ async def async_get_translations(
     Otherwise default to loaded intgrations combined with config flow
     integrations if config_flow is true.
     """
-    lock = hass.data.get(TRANSLATION_LOAD_LOCK)
+    lock = hass.data.setdefault(TRANSLATION_LOAD_LOCK, asyncio.Lock())
-    if lock is None:
-        lock = hass.data[TRANSLATION_LOAD_LOCK] = asyncio.Lock()
 
     if integration is not None:
         components = {integration}
     elif config_flow:
-        # When it's a config flow, we're going to merge the cached loaded component results
-        # with the integrations that have not been loaded yet. We merge this at the end.
-        # We can't cache with config flow, as we can't monitor it during runtime.
         components = (await async_get_config_flows(hass)) - hass.config.components
-    else:
+    elif category == "state":
-        # Only 'state' supports merging, so remove platforms from selection
-        if category == "state":
         components = set(hass.config.components)
     else:
+        # Only 'state' supports merging, so remove platforms from selection
         components = {
-                component
+            component for component in hass.config.components if "." not in component
-                for component in hass.config.components
-                if "." not in component
         }
 
     async with lock:
-        use_cache = integration is None and not config_flow
+        cache = hass.data.setdefault(TRANSLATION_FLATTEN_CACHE, _TranslationCache(hass))
-        if use_cache:
+        cached = await cache.async_fetch(language, category, components)
-            cache = hass.data.get(TRANSLATION_FLATTEN_CACHE)
-            if cache is None:
-                cache = hass.data[TRANSLATION_FLATTEN_CACHE] = FlatCache(hass)
-                cache.async_setup()
 
-            cached_translations = cache.async_get_cache(language, category)
+    return dict(ChainMap(*cached))
-
-            if cached_translations is not None:
-                return cached_translations
-
-        tasks = [async_get_component_strings(hass, language, components)]
-
-        # Fetch the English resources, as a fallback for missing keys
-        if language != "en":
-            tasks.append(async_get_component_strings(hass, "en", components))
-
-        _LOGGER.debug(
-            "Cache miss for %s, %s: %s", language, category, ", ".join(components)
-        )
-
-        results = await asyncio.gather(*tasks)
-
-        if category == "state":
-            resource_func = merge_resources
-        else:
-            resource_func = build_resources
-
-        resources = flatten(resource_func(results[0], components, category))
-
-        if language != "en":
-            base_resources = flatten(resource_func(results[1], components, category))
-            resources = {**base_resources, **resources}
-
-        # The cache must be set while holding the lock
-        if use_cache:
-            assert cache is not None
-            cache.async_set_cache(language, category, resources)
-
-    if config_flow:
-        loaded_comp_resources = await async_get_translations(hass, language, category)
-        resources.update(loaded_comp_resources)
-
-    return resources

tests/helpers/test_translation.py

@@ -5,7 +5,6 @@ import pathlib
 
 import pytest
 
-from homeassistant.const import EVENT_COMPONENT_LOADED
 from homeassistant.generated import config_flows
 from homeassistant.helpers import translation
 from homeassistant.loader import async_get_integration
@@ -22,16 +21,16 @@ def mock_config_flows():
         yield flows
 
 
-def test_flatten():
+def test_recursive_flatten():
     """Test the flatten function."""
     data = {"parent1": {"child1": "data1", "child2": "data2"}, "parent2": "data3"}
 
-    flattened = translation.flatten(data)
+    flattened = translation.recursive_flatten("prefix.", data)
 
     assert flattened == {
-        "parent1.child1": "data1",
+        "prefix.parent1.child1": "data1",
-        "parent1.child2": "data2",
+        "prefix.parent1.child2": "data2",
-        "parent2": "data3",
+        "prefix.parent2": "data3",
     }
 
 
@@ -149,21 +148,62 @@ async def test_get_translations_loads_config_flows(hass, mock_config_flows):
         return_value="bla.json",
     ), patch(
         "homeassistant.helpers.translation.load_translations_files",
-        return_value={"component1": {"hello": "world"}},
+        return_value={"component1": {"title": "world"}},
     ), patch(
         "homeassistant.helpers.translation.async_get_integration",
         return_value=integration,
     ):
         translations = await translation.async_get_translations(
-            hass, "en", "hello", config_flow=True
+            hass, "en", "title", config_flow=True
+        )
+        translations_again = await translation.async_get_translations(
+            hass, "en", "title", config_flow=True
         )
 
+        assert translations == translations_again
+
     assert translations == {
-        "component.component1.hello": "world",
+        "component.component1.title": "world",
     }
 
     assert "component1" not in hass.config.components
 
+    mock_config_flows.append("component2")
+    integration = Mock(file_path=pathlib.Path(__file__))
+    integration.name = "Component 2"
+
+    with patch(
+        "homeassistant.helpers.translation.component_translation_path",
+        return_value="bla.json",
+    ), patch(
+        "homeassistant.helpers.translation.load_translations_files",
+        return_value={"component2": {"title": "world"}},
+    ), patch(
+        "homeassistant.helpers.translation.async_get_integration",
+        return_value=integration,
+    ):
+        translations = await translation.async_get_translations(
+            hass, "en", "title", config_flow=True
+        )
+        translations_again = await translation.async_get_translations(
+            hass, "en", "title", config_flow=True
+        )
+
+        assert translations == translations_again
+
+    assert translations == {
+        "component.component1.title": "world",
+        "component.component2.title": "world",
+    }
+
+    translations_all_cached = await translation.async_get_translations(
+        hass, "en", "title", config_flow=True
+    )
+    assert translations == translations_all_cached
+
+    assert "component1" not in hass.config.components
+    assert "component2" not in hass.config.components
+
 
 async def test_get_translations_while_loading_components(hass):
     """Test the get translations helper loads config flow translations."""
@@ -178,7 +218,7 @@ async def test_get_translations_while_loading_components(hass):
         load_count += 1
         # Mimic race condition by loading a component during setup
         setup_component(hass, "persistent_notification", {})
-        return {"component1": {"hello": "world"}}
+        return {"component1": {"title": "world"}}
 
     with patch(
         "homeassistant.helpers.translation.component_translation_path",
@@ -191,12 +231,12 @@ async def test_get_translations_while_loading_components(hass):
         return_value=integration,
     ):
         tasks = [
-            translation.async_get_translations(hass, "en", "hello") for _ in range(5)
+            translation.async_get_translations(hass, "en", "title") for _ in range(5)
         ]
         all_translations = await asyncio.gather(*tasks)
 
     assert all_translations[0] == {
-        "component.component1.hello": "world",
+        "component.component1.title": "world",
     }
     assert load_count == 1
 
@@ -218,17 +258,13 @@ async def test_get_translation_categories(hass):
 async def test_translation_merging(hass, caplog):
     """Test we merge translations of two integrations."""
     hass.config.components.add("sensor.moon")
-    hass.config.components.add("sensor.season")
     hass.config.components.add("sensor")
 
     translations = await translation.async_get_translations(hass, "en", "state")
 
     assert "component.sensor.state.moon__phase.first_quarter" in translations
-    assert "component.sensor.state.season__season.summer" in translations
 
-    # Clear cache
+    hass.config.components.add("sensor.season")
-    hass.bus.async_fire(EVENT_COMPONENT_LOADED)
-    await hass.async_block_till_done()
 
     # Patch in some bad translation data
 
@@ -254,27 +290,91 @@ async def test_translation_merging(hass, caplog):
     )
 
 
+async def test_translation_merging_loaded_apart(hass, caplog):
+    """Test we merge translations of two integrations when they are not loaded at the same time."""
+    hass.config.components.add("sensor")
+
+    translations = await translation.async_get_translations(hass, "en", "state")
+
+    assert "component.sensor.state.moon__phase.first_quarter" not in translations
+
+    hass.config.components.add("sensor.moon")
+
+    translations = await translation.async_get_translations(hass, "en", "state")
+
+    assert "component.sensor.state.moon__phase.first_quarter" in translations
+
+    translations = await translation.async_get_translations(
+        hass, "en", "state", integration="sensor"
+    )
+
+    assert "component.sensor.state.moon__phase.first_quarter" in translations
+
+
 async def test_caching(hass):
     """Test we cache data."""
     hass.config.components.add("sensor")
+    hass.config.components.add("light")
 
     # Patch with same method so we can count invocations
     with patch(
-        "homeassistant.helpers.translation.merge_resources",
+        "homeassistant.helpers.translation._merge_resources",
-        side_effect=translation.merge_resources,
+        side_effect=translation._merge_resources,
     ) as mock_merge:
-        await translation.async_get_translations(hass, "en", "state")
+        load1 = await translation.async_get_translations(hass, "en", "state")
         assert len(mock_merge.mock_calls) == 1
 
-        await translation.async_get_translations(hass, "en", "state")
+        load2 = await translation.async_get_translations(hass, "en", "state")
         assert len(mock_merge.mock_calls) == 1
 
-        # This event clears the cache so we should record another call
+        assert load1 == load2
-        hass.bus.async_fire(EVENT_COMPONENT_LOADED)
-        await hass.async_block_till_done()
 
-        await translation.async_get_translations(hass, "en", "state")
+        for key in load1:
-        assert len(mock_merge.mock_calls) == 2
+            assert key.startswith("component.sensor.state.") or key.startswith(
+                "component.light.state."
+            )
+
+    load_sensor_only = await translation.async_get_translations(
+        hass, "en", "state", integration="sensor"
+    )
+    assert load_sensor_only
+    for key in load_sensor_only:
+        assert key.startswith("component.sensor.state.")
+
+    load_light_only = await translation.async_get_translations(
+        hass, "en", "state", integration="light"
+    )
+    assert load_light_only
+    for key in load_light_only:
+        assert key.startswith("component.light.state.")
+
+    hass.config.components.add("media_player")
+
+    # Patch with same method so we can count invocations
+    with patch(
+        "homeassistant.helpers.translation._build_resources",
+        side_effect=translation._build_resources,
+    ) as mock_build:
+        load_sensor_only = await translation.async_get_translations(
+            hass, "en", "title", integration="sensor"
+        )
+        assert load_sensor_only
+        for key in load_sensor_only:
+            assert key == "component.sensor.title"
+        assert len(mock_build.mock_calls) == 0
+
+        assert await translation.async_get_translations(
+            hass, "en", "title", integration="sensor"
+        )
+        assert len(mock_build.mock_calls) == 0
+
+        load_light_only = await translation.async_get_translations(
+            hass, "en", "title", integration="media_player"
+        )
+        assert load_light_only
+        for key in load_light_only:
+            assert key == "component.media_player.title"
+        assert len(mock_build.mock_calls) > 1
 
 
 async def test_custom_component_translations(hass):