Repository: home-assistant/core
File: homeassistant/components/websocket_api/automation.py (Python)

Code before:
def _async_get_automation_components_for_target(
    hass: HomeAssistant,
    target_selection: ConfigType,
    expand_group: bool,
    component_descriptions: Mapping[str, Mapping[str, Any] | None],
) -> set[str]:
    """Get automation components (triggers/conditions/services) for a target.

    Returns all components that can be used on any entity that are currently part of a target.
    """
    extracted = target_helpers.async_extract_referenced_entity_ids(
        hass,
        target_helpers.TargetSelectorData(target_selection),
        expand_group=expand_group,
    )
    _LOGGER.debug("Extracted entities for lookup: %s", extracted)

    # Build lookup structure: domain -> list of trigger/condition/service lookup data
    domain_components: dict[str | None, list[_AutomationComponentLookupData]] = {}
    component_count = 0
    for component, description in component_descriptions.items():
        if description is None or CONF_TARGET not in description:
            _LOGGER.debug("Skipping component %s without target description", component)
            continue
        domains = _get_automation_component_domains(description[CONF_TARGET])
        lookup_data = _AutomationComponentLookupData.create(
            component, description[CONF_TARGET]
        )
        for domain in domains:
            domain_components.setdefault(domain, []).append(lookup_data)
        component_count += 1

    _LOGGER.debug("Automation components per domain: %s", domain_components)

    entity_infos = entity_sources(hass)
    matched_components: set[str] = set()
    for entity_id in extracted.referenced | extracted.indirectly_referenced:
        if component_count == len(matched_components):
            # All automation components matched already, so we don't need to iterate further
            break

        entity_info = entity_infos.get(entity_id)
        if entity_info is None:
            _LOGGER.debug("No entity source found for %s", entity_id)
            continue

        entity_domain = entity_id.split(".")[0]
        entity_integration = entity_info["domain"]
        for domain in (entity_domain, entity_integration, None):
            for component_data in domain_components.get(domain, []):
                if component_data.component in matched_components:

Review comment (performance):
Looking at this 2 weeks later, I'm not sure why I didn't just suggest this:
```suggestion
    cache = hass.data.setdefault(AUTOMATION_COMPONENT_LOOKUP_CACHE, {})
```

The drawback is that setdefault unfortunately doesn't allow passing a function, so we need to create and discard a dict every single time except for the first one. 

I don't mind which version; what's in the PR now (because it's abit more efficient) or this updated proposal (because it's maybe a bit more pythonic)

Code after:
            domains.add(None)
            continue

        if filter_integration:
            domains.add(filter_integration)

        for domain in filter_domains:
            domains.add(domain)

    return domains


def _get_automation_component_lookup_table(
    hass: HomeAssistant,
    component_type: str,
    component_descriptions: Mapping[str, Mapping[str, Any] | None],
) -> _AutomationComponentLookupTable:
    """Get a dict of automation components keyed by domain, along with the total number of components.

    Returns a cached object if available.
    """

    try:
        cache = hass.data[AUTOMATION_COMPONENT_LOOKUP_CACHE]
    except KeyError:
        cache = hass.data[AUTOMATION_COMPONENT_LOOKUP_CACHE] = {}

    if (cached := cache.get(component_type)) is not None:
        cached_descriptions, cached_lookup = cached
        if cached_descriptions is component_descriptions:
            _LOGGER.debug(
                "Using cached automation component lookup data for %s", component_type
            )
            return cached_lookup

    lookup_table = _AutomationComponentLookupTable(
        domain_components={}, component_count=0
    )
    for component, description in component_descriptions.items():
        if description is None or CONF_TARGET not in description:
            _LOGGER.debug("Skipping component %s without target description", component)
            continue
        domains = _get_automation_component_domains(description[CONF_TARGET])
        lookup_data = _AutomationComponentLookupData.create(
            component, description[CONF_TARGET]
        )
        for domain in domains:
            lookup_table.domain_components.setdefault(domain, []).append(lookup_data)
        lookup_table.component_count += 1

    cache[component_type] = (component_descriptions, lookup_table)