Add gas support to energy (#54560)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/energy/data.py

@@ -88,7 +88,25 @@ class BatterySourceType(TypedDict):
     stat_energy_to: str
 
 
-SourceType = Union[GridSourceType, SolarSourceType, BatterySourceType]
+class GasSourceType(TypedDict):
+    """Dictionary holding the source of gas storage."""
+
+    type: Literal["gas"]
+
+    stat_energy_from: str
+
+    # statistic_id of costs ($) incurred from the energy meter
+    # If set to None and entity_energy_from and entity_energy_price are configured,
+    # an EnergyCostSensor will be automatically created
+    stat_cost: str | None
+
+    # Used to generate costs if stat_cost is set to None
+    entity_energy_from: str | None  # entity_id of an gas meter (m³), entity_id of the gas meter for stat_energy_from
+    entity_energy_price: str | None  # entity_id of an entity providing price ($/m³)
+    number_energy_price: float | None  # Price for energy ($/m³)
+
+
+SourceType = Union[GridSourceType, SolarSourceType, BatterySourceType, GasSourceType]
 
 
 class DeviceConsumption(TypedDict):
@@ -193,6 +211,16 @@ BATTERY_SOURCE_SCHEMA = vol.Schema(
         vol.Required("stat_energy_to"): str,
     }
 )
+GAS_SOURCE_SCHEMA = vol.Schema(
+    {
+        vol.Required("type"): "gas",
+        vol.Required("stat_energy_from"): str,
+        vol.Optional("stat_cost"): vol.Any(str, None),
+        vol.Optional("entity_energy_from"): vol.Any(str, None),
+        vol.Optional("entity_energy_price"): vol.Any(str, None),
+        vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
+    }
+)
 
 
 def check_type_limits(value: list[SourceType]) -> list[SourceType]:
@@ -214,6 +242,7 @@ ENERGY_SOURCE_SCHEMA = vol.All(
                     "grid": GRID_SOURCE_SCHEMA,
                     "solar": SOLAR_SOURCE_SCHEMA,
                     "battery": BATTERY_SOURCE_SCHEMA,
+                    "gas": GAS_SOURCE_SCHEMA,
                 },
             )
         ]

homeassistant/components/energy/sensor.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from functools import partial
 import logging
 from typing import Any, Final, Literal, TypeVar, cast
 
@@ -16,6 +15,7 @@ from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     ENERGY_KILO_WATT_HOUR,
     ENERGY_WATT_HOUR,
+    VOLUME_CUBIC_METERS,
 )
 from homeassistant.core import HomeAssistant, State, callback, split_entity_id
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -36,22 +36,19 @@ async def async_setup_platform(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the energy sensors."""
-    manager = await async_get_manager(hass)
+    sensor_manager = SensorManager(await async_get_manager(hass), async_add_entities)
-    process_now = partial(_process_manager_data, hass, manager, async_add_entities, {})
+    await sensor_manager.async_start()
-    manager.async_listen_updates(process_now)
-
-    if manager.data:
-        await process_now()
 
 
 T = TypeVar("T")
 
 
 @dataclass
-class FlowAdapter:
+class SourceAdapter:
-    """Adapter to allow flows to be used as sensors."""
+    """Adapter to allow sources and their flows to be used as sensors."""
 
-    flow_type: Literal["flow_from", "flow_to"]
+    source_type: Literal["grid", "gas"]
+    flow_type: Literal["flow_from", "flow_to", None]
     stat_energy_key: Literal["stat_energy_from", "stat_energy_to"]
     entity_energy_key: Literal["entity_energy_from", "entity_energy_to"]
     total_money_key: Literal["stat_cost", "stat_compensation"]
@@ -59,8 +56,9 @@ class FlowAdapter:
     entity_id_suffix: str
 
 
-FLOW_ADAPTERS: Final = (
+SOURCE_ADAPTERS: Final = (
-    FlowAdapter(
+    SourceAdapter(
+        "grid",
         "flow_from",
         "stat_energy_from",
         "entity_energy_from",
@@ -68,7 +66,8 @@ FLOW_ADAPTERS: Final = (
         "Cost",
         "cost",
     ),
-    FlowAdapter(
+    SourceAdapter(
+        "grid",
         "flow_to",
         "stat_energy_to",
         "entity_energy_to",
@@ -76,67 +75,112 @@ FLOW_ADAPTERS: Final = (
         "Compensation",
         "compensation",
     ),
+    SourceAdapter(
+        "gas",
+        None,
+        "stat_energy_from",
+        "entity_energy_from",
+        "stat_cost",
+        "Cost",
+        "cost",
+    ),
 )
 
 
-async def _process_manager_data(
+class SensorManager:
-    hass: HomeAssistant,
+    """Class to handle creation/removal of sensor data."""
-    manager: EnergyManager,
+
-    async_add_entities: AddEntitiesCallback,
+    def __init__(
-    current_entities: dict[tuple[str, str], EnergyCostSensor],
+        self, manager: EnergyManager, async_add_entities: AddEntitiesCallback
     ) -> None:
-    """Process updated data."""
+        """Initialize sensor manager."""
+        self.manager = manager
+        self.async_add_entities = async_add_entities
+        self.current_entities: dict[tuple[str, str | None, str], EnergyCostSensor] = {}
+
+    async def async_start(self) -> None:
+        """Start."""
+        self.manager.async_listen_updates(self._process_manager_data)
+
+        if self.manager.data:
+            await self._process_manager_data()
+
+    async def _process_manager_data(self) -> None:
+        """Process manager data."""
         to_add: list[SensorEntity] = []
-    to_remove = dict(current_entities)
+        to_remove = dict(self.current_entities)
 
         async def finish() -> None:
             if to_add:
-            async_add_entities(to_add)
+                self.async_add_entities(to_add)
 
             for key, entity in to_remove.items():
-            current_entities.pop(key)
+                self.current_entities.pop(key)
                 await entity.async_remove()
 
-    if not manager.data:
+        if not self.manager.data:
             await finish()
             return
 
-    for energy_source in manager.data["energy_sources"]:
+        for energy_source in self.manager.data["energy_sources"]:
-        if energy_source["type"] != "grid":
+            for adapter in SOURCE_ADAPTERS:
+                if adapter.source_type != energy_source["type"]:
                     continue
 
-        for adapter in FLOW_ADAPTERS:
+                if adapter.flow_type is None:
-            for flow in energy_source[adapter.flow_type]:
+                    self._process_sensor_data(
+                        adapter,
                         # Opting out of the type complexity because can't get it to work
-                untyped_flow = cast(dict, flow)
+                        energy_source,  # type: ignore
-
+                        to_add,
-                # No need to create an entity if we already have a cost stat
+                        to_remove,
-                if untyped_flow.get(adapter.total_money_key) is not None:
+                    )
                     continue
 
-                # This is unique among all flow_from's
+                for flow in energy_source[adapter.flow_type]:  # type: ignore
-                key = (adapter.flow_type, untyped_flow[adapter.stat_energy_key])
+                    self._process_sensor_data(
+                        adapter,
+                        # Opting out of the type complexity because can't get it to work
+                        flow,  # type: ignore
+                        to_add,
+                        to_remove,
+                    )
+
+        await finish()
+
+    @callback
+    def _process_sensor_data(
+        self,
+        adapter: SourceAdapter,
+        config: dict,
+        to_add: list[SensorEntity],
+        to_remove: dict[tuple[str, str | None, str], EnergyCostSensor],
+    ) -> None:
+        """Process sensor data."""
+        # No need to create an entity if we already have a cost stat
+        if config.get(adapter.total_money_key) is not None:
+            return
+
+        key = (adapter.source_type, adapter.flow_type, config[adapter.stat_energy_key])
 
         # Make sure the right data is there
         # If the entity existed, we don't pop it from to_remove so it's removed
-                if untyped_flow.get(adapter.entity_energy_key) is None or (
+        if config.get(adapter.entity_energy_key) is None or (
-                    untyped_flow.get("entity_energy_price") is None
+            config.get("entity_energy_price") is None
-                    and untyped_flow.get("number_energy_price") is None
+            and config.get("number_energy_price") is None
         ):
-                    continue
+            return
 
         current_entity = to_remove.pop(key, None)
         if current_entity:
-                    current_entity.update_config(untyped_flow)
+            current_entity.update_config(config)
-                    continue
+            return
 
-                current_entities[key] = EnergyCostSensor(
+        self.current_entities[key] = EnergyCostSensor(
             adapter,
-                    untyped_flow,
+            config,
         )
-                to_add.append(current_entities[key])
+        to_add.append(self.current_entities[key])
-
-    await finish()
 
 
 class EnergyCostSensor(SensorEntity):
@@ -148,17 +192,19 @@ class EnergyCostSensor(SensorEntity):
 
     def __init__(
         self,
-        adapter: FlowAdapter,
+        adapter: SourceAdapter,
-        flow: dict,
+        config: dict,
     ) -> None:
         """Initialize the sensor."""
         super().__init__()
 
         self._adapter = adapter
-        self.entity_id = f"{flow[adapter.entity_energy_key]}_{adapter.entity_id_suffix}"
+        self.entity_id = (
+            f"{config[adapter.entity_energy_key]}_{adapter.entity_id_suffix}"
+        )
         self._attr_device_class = DEVICE_CLASS_MONETARY
         self._attr_state_class = STATE_CLASS_MEASUREMENT
-        self._flow = flow
+        self._config = config
         self._last_energy_sensor_state: State | None = None
         self._cur_value = 0.0
 
@@ -174,7 +220,7 @@ class EnergyCostSensor(SensorEntity):
     def _update_cost(self) -> None:
         """Update incurred costs."""
         energy_state = self.hass.states.get(
-            cast(str, self._flow[self._adapter.entity_energy_key])
+            cast(str, self._config[self._adapter.entity_energy_key])
         )
 
         if energy_state is None or ATTR_LAST_RESET not in energy_state.attributes:
@@ -186,8 +232,10 @@ class EnergyCostSensor(SensorEntity):
             return
 
         # Determine energy price
-        if self._flow["entity_energy_price"] is not None:
+        if self._config["entity_energy_price"] is not None:
-            energy_price_state = self.hass.states.get(self._flow["entity_energy_price"])
+            energy_price_state = self.hass.states.get(
+                self._config["entity_energy_price"]
+            )
 
             if energy_price_state is None:
                 return
@@ -197,14 +245,17 @@ class EnergyCostSensor(SensorEntity):
             except ValueError:
                 return
 
-            if energy_price_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT, "").endswith(
+            if (
-                f"/{ENERGY_WATT_HOUR}"
+                self._adapter.source_type == "grid"
+                and energy_price_state.attributes.get(
+                    ATTR_UNIT_OF_MEASUREMENT, ""
+                ).endswith(f"/{ENERGY_WATT_HOUR}")
             ):
                 energy_price *= 1000.0
 
         else:
             energy_price_state = None
-            energy_price = cast(float, self._flow["number_energy_price"])
+            energy_price = cast(float, self._config["number_energy_price"])
 
         if self._last_energy_sensor_state is None:
             # Initialize as it's the first time all required entities are in place.
@@ -213,9 +264,17 @@ class EnergyCostSensor(SensorEntity):
 
         energy_unit = energy_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
 
+        if self._adapter.source_type == "grid":
             if energy_unit == ENERGY_WATT_HOUR:
                 energy_price /= 1000
             elif energy_unit != ENERGY_KILO_WATT_HOUR:
+                energy_unit = None
+
+        elif self._adapter.source_type == "gas":
+            if energy_unit != VOLUME_CUBIC_METERS:
+                energy_unit = None
+
+        if energy_unit is None:
             _LOGGER.warning(
                 "Found unexpected unit %s for %s", energy_unit, energy_state.entity_id
             )
@@ -237,11 +296,13 @@ class EnergyCostSensor(SensorEntity):
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
-        energy_state = self.hass.states.get(self._flow[self._adapter.entity_energy_key])
+        energy_state = self.hass.states.get(
+            self._config[self._adapter.entity_energy_key]
+        )
         if energy_state:
             name = energy_state.name
         else:
-            name = split_entity_id(self._flow[self._adapter.entity_energy_key])[
+            name = split_entity_id(self._config[self._adapter.entity_energy_key])[
                 0
             ].replace("_", " ")
 
@@ -251,7 +312,7 @@ class EnergyCostSensor(SensorEntity):
 
         # Store stat ID in hass.data so frontend can look it up
         self.hass.data[DOMAIN]["cost_sensors"][
-            self._flow[self._adapter.entity_energy_key]
+            self._config[self._adapter.entity_energy_key]
         ] = self.entity_id
 
         @callback
@@ -263,7 +324,7 @@ class EnergyCostSensor(SensorEntity):
         self.async_on_remove(
             async_track_state_change_event(
                 self.hass,
-                cast(str, self._flow[self._adapter.entity_energy_key]),
+                cast(str, self._config[self._adapter.entity_energy_key]),
                 async_state_changed_listener,
             )
         )
@@ -271,14 +332,14 @@ class EnergyCostSensor(SensorEntity):
     async def async_will_remove_from_hass(self) -> None:
         """Handle removing from hass."""
         self.hass.data[DOMAIN]["cost_sensors"].pop(
-            self._flow[self._adapter.entity_energy_key]
+            self._config[self._adapter.entity_energy_key]
         )
         await super().async_will_remove_from_hass()
 
     @callback
-    def update_config(self, flow: dict) -> None:
+    def update_config(self, config: dict) -> None:
         """Update the config."""
-        self._flow = flow
+        self._config = config
 
     @property
     def native_unit_of_measurement(self) -> str | None:

tests/components/energy/test_sensor.py

@@ -18,6 +18,7 @@ from homeassistant.const import (
     DEVICE_CLASS_MONETARY,
     ENERGY_KILO_WATT_HOUR,
     ENERGY_WATT_HOUR,
+    VOLUME_CUBIC_METERS,
 )
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
@@ -295,3 +296,49 @@ async def test_cost_sensor_handle_wh(hass, hass_storage) -> None:
 
     state = hass.states.get("sensor.energy_consumption_cost")
     assert state.state == "5.0"
+
+
+async def test_cost_sensor_handle_gas(hass, hass_storage) -> None:
+    """Test gas cost price from sensor entity."""
+    energy_data = data.EnergyManager.default_preferences()
+    energy_data["energy_sources"].append(
+        {
+            "type": "gas",
+            "stat_energy_from": "sensor.gas_consumption",
+            "entity_energy_from": "sensor.gas_consumption",
+            "stat_cost": None,
+            "entity_energy_price": None,
+            "number_energy_price": 0.5,
+        }
+    )
+
+    hass_storage[data.STORAGE_KEY] = {
+        "version": 1,
+        "data": energy_data,
+    }
+
+    now = dt_util.utcnow()
+    last_reset = dt_util.utc_from_timestamp(0).isoformat()
+
+    hass.states.async_set(
+        "sensor.gas_consumption",
+        100,
+        {"last_reset": last_reset, "unit_of_measurement": VOLUME_CUBIC_METERS},
+    )
+
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        await setup_integration(hass)
+
+    state = hass.states.get("sensor.gas_consumption_cost")
+    assert state.state == "0.0"
+
+    # gas use bumped to 10 kWh
+    hass.states.async_set(
+        "sensor.gas_consumption",
+        200,
+        {"last_reset": last_reset, "unit_of_measurement": VOLUME_CUBIC_METERS},
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.gas_consumption_cost")
+    assert state.state == "50.0"