Replace statistics _get_unit_class function with a dict lookup (#124452)

home-assistant · Aug 23, 2024 · 7579ebc · 7579ebc
1 parent 61ac4c7
commit 7579ebc
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/homeassistant/components/recorder/statistics.py b/homeassistant/components/recorder/statistics.py
@@ -148,6 +148,12 @@
     **{unit: VolumeFlowRateConverter for unit in VolumeFlowRateConverter.VALID_UNITS},
 }
 
+
+UNIT_CLASSES = {
+    unit: converter.UNIT_CLASS
+    for unit, converter in STATISTIC_UNIT_TO_UNIT_CONVERTER.items()
+}
+
 DATA_SHORT_TERM_STATISTICS_RUN_CACHE = "recorder_short_term_statistics_run_cache"
 
 
@@ -211,13 +217,6 @@ class StatisticsRow(BaseStatisticsRow, total=False):
     change: float | None
 
 
-def _get_unit_class(unit: str | None) -> str | None:
-    """Get corresponding unit class from from the statistics unit."""
-    if converter := STATISTIC_UNIT_TO_UNIT_CONVERTER.get(unit):
-        return converter.UNIT_CLASS
-    return None
-
-
 def get_display_unit(
     hass: HomeAssistant,
     statistic_id: str,
@@ -807,7 +806,7 @@ def _statistic_by_id_from_metadata(
             "has_sum": meta["has_sum"],
             "name": meta["name"],
             "source": meta["source"],
-            "unit_class": _get_unit_class(meta["unit_of_measurement"]),
+            "unit_class": UNIT_CLASSES.get(meta["unit_of_measurement"]),
             "unit_of_measurement": meta["unit_of_measurement"],
         }
         for _, meta in metadata.values()
@@ -881,7 +880,7 @@ def list_statistic_ids(
                     "has_sum": meta["has_sum"],
                     "name": meta["name"],
                     "source": meta["source"],
-                    "unit_class": _get_unit_class(meta["unit_of_measurement"]),
+                    "unit_class": UNIT_CLASSES.get(meta["unit_of_measurement"]),
                     "unit_of_measurement": meta["unit_of_measurement"],
                 }