Add category registry (#110897)

* Add category registry

* Add entity registry support

* Update homeassistant/components/config/entity_registry.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Use ulid instead

* Add tests for adding same name in different scopes

* Handle keyerror on update

* Lookup tweak

* Omit categories from entity registry snapshots

* Use base registry

* Update snapshots

* Update snapshots

---------

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/bootstrap.py

@@ -62,6 +62,7 @@
 from .exceptions import HomeAssistantError
 from .helpers import (
     area_registry,
+    category_registry,
     config_validation as cv,
     device_registry,
     entity,
@@ -342,6 +343,7 @@ def _cache_uname_processor() -> None:
     template.async_setup(hass)
     await asyncio.gather(
         create_eager_task(area_registry.async_load(hass)),
+        create_eager_task(category_registry.async_load(hass)),
         create_eager_task(device_registry.async_load(hass)),
         create_eager_task(entity_registry.async_load(hass)),
         create_eager_task(floor_registry.async_load(hass)),

homeassistant/components/config/__init__.py

@@ -14,6 +14,7 @@
     auth,
     auth_provider_homeassistant,
     automation,
+    category_registry,
     config_entries,
     core,
     device_registry,
@@ -30,6 +31,7 @@
     auth,
     auth_provider_homeassistant,
     automation,
+    category_registry,
     config_entries,
     core,
     device_registry,

homeassistant/components/config/category_registry.py

@@ -0,0 +1,134 @@
+"""Websocket API to interact with the category registry."""
+from typing import Any
+
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+from homeassistant.components.websocket_api.connection import ActiveConnection
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import category_registry as cr, config_validation as cv
+
+
+@callback
+def async_setup(hass: HomeAssistant) -> bool:
+    """Register the category registry WS commands."""
+    websocket_api.async_register_command(hass, websocket_list_categories)
+    websocket_api.async_register_command(hass, websocket_create_category)
+    websocket_api.async_register_command(hass, websocket_delete_category)
+    websocket_api.async_register_command(hass, websocket_update_category)
+    return True
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/category_registry/list",
+        vol.Required("scope"): str,
+    }
+)
+@callback
+def websocket_list_categories(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Handle list categories command."""
+    category_registry = cr.async_get(hass)
+    connection.send_result(
+        msg["id"],
+        [
+            _entry_dict(entry)
+            for entry in category_registry.async_list_categories(scope=msg["scope"])
+        ],
+    )
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/category_registry/create",
+        vol.Required("scope"): str,
+        vol.Required("name"): str,
+        vol.Optional("icon"): vol.Any(cv.icon, None),
+    }
+)
+@websocket_api.require_admin
+@callback
+def websocket_create_category(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Create category command."""
+    category_registry = cr.async_get(hass)
+
+    data = dict(msg)
+    data.pop("type")
+    data.pop("id")
+
+    try:
+        entry = category_registry.async_create(**data)
+    except ValueError as err:
+        connection.send_error(msg["id"], "invalid_info", str(err))
+    else:
+        connection.send_result(msg["id"], _entry_dict(entry))
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/category_registry/delete",
+        vol.Required("scope"): str,
+        vol.Required("category_id"): str,
+    }
+)
+@websocket_api.require_admin
+@callback
+def websocket_delete_category(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Delete category command."""
+    category_registry = cr.async_get(hass)
+
+    try:
+        category_registry.async_delete(
+            scope=msg["scope"], category_id=msg["category_id"]
+        )
+    except KeyError:
+        connection.send_error(msg["id"], "invalid_info", "Category ID doesn't exist")
+    else:
+        connection.send_result(msg["id"])
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/category_registry/update",
+        vol.Required("scope"): str,
+        vol.Required("category_id"): str,
+        vol.Optional("name"): str,
+        vol.Optional("icon"): vol.Any(cv.icon, None),
+    }
+)
+@websocket_api.require_admin
+@callback
+def websocket_update_category(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Handle update category websocket command."""
+    category_registry = cr.async_get(hass)
+
+    data = dict(msg)
+    data.pop("type")
+    data.pop("id")
+
+    try:
+        entry = category_registry.async_update(**data)
+    except ValueError as err:
+        connection.send_error(msg["id"], "invalid_info", str(err))
+    except KeyError:
+        connection.send_error(msg["id"], "invalid_info", "Category ID doesn't exist")
+    else:
+        connection.send_result(msg["id"], _entry_dict(entry))
+
+
+@callback
+def _entry_dict(entry: cr.CategoryEntry) -> dict[str, Any]:
+    """Convert entry to API format."""
+    return {
+        "category_id": entry.category_id,
+        "icon": entry.icon,
+        "name": entry.name,
+    }

homeassistant/components/config/entity_registry.py

@@ -153,6 +153,16 @@ def websocket_get_entities(
         # If passed in, we update value. Passing None will remove old value.
         vol.Optional("aliases"): list,
         vol.Optional("area_id"): vol.Any(str, None),
+        # Categories is a mapping of key/value (scope/category_id) pairs.
+        # If passed in, we update/adjust only the provided scope(s).
+        # Other category scopes in the entity, are left as is.
+        #
+        # Categorized items such as entities
+        # can only be in 1 category ID per scope at a time.
+        # Therefore, passing in a category ID will either add or move
+        # the entity to that specific category. Passing in None will
+        # remove the entity from the category.
+        vol.Optional("categories"): cv.schema_with_slug_keys(vol.Any(str, None)),
         vol.Optional("device_class"): vol.Any(str, None),
         vol.Optional("icon"): vol.Any(str, None),
         vol.Optional("name"): vol.Any(str, None),
@@ -227,6 +237,18 @@ def websocket_update_entity(
                 )
                 return
 
+    # Update the categories if provided
+    if "categories" in msg:
+        categories = entity_entry.categories.copy()
+        for scope, category_id in msg["categories"].items():
+            if scope in categories and category_id is None:
+                # Remove the category from the scope as it was unset
+                del categories[scope]
+            elif category_id is not None:
+                # Add or update the category for the given scope
+                categories[scope] = category_id
+        changes["categories"] = categories
+
     try:
         if changes:
             entity_entry = registry.async_update_entity(entity_id, **changes)