From 9d1985ab033920c2f4d965b576fbc9771b893a15 Mon Sep 17 00:00:00 2001
From: Kevin Worrel <37058192+dieselrabbit@users.noreply.github.com>
Date: Thu, 2 Dec 2021 18:36:31 -0800
Subject: [PATCH] Move screenlogic SCG levels to number platform (#60872)

Co-authored-by: J. Nick Koston <nick@koston.org>
---
 .coveragerc                                   |  1 +
 .../components/screenlogic/__init__.py        |  2 +-
 .../components/screenlogic/number.py          | 78 +++++++++++++++++++
 .../components/screenlogic/sensor.py          |  2 -
 4 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 homeassistant/components/screenlogic/number.py

diff --git a/.coveragerc b/.coveragerc
index bf26b41649d..fc033c77369 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -915,6 +915,7 @@ omit =
     homeassistant/components/screenlogic/binary_sensor.py
     homeassistant/components/screenlogic/climate.py
     homeassistant/components/screenlogic/light.py
+    homeassistant/components/screenlogic/number.py
     homeassistant/components/screenlogic/sensor.py
     homeassistant/components/screenlogic/services.py
     homeassistant/components/screenlogic/switch.py
diff --git a/homeassistant/components/screenlogic/__init__.py b/homeassistant/components/screenlogic/__init__.py
index bc7761857d5..72c1f162552 100644
--- a/homeassistant/components/screenlogic/__init__.py
+++ b/homeassistant/components/screenlogic/__init__.py
@@ -40,7 +40,7 @@ HEATER_COOLDOWN_DELAY = 6
 # These seem to be constant across all controller models
 PRIMARY_CIRCUIT_IDS = [500, 505]  # [Spa, Pool]
 
-PLATFORMS = ["binary_sensor", "climate", "light", "sensor", "switch"]
+PLATFORMS = ["binary_sensor", "climate", "light", "number", "sensor", "switch"]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
diff --git a/homeassistant/components/screenlogic/number.py b/homeassistant/components/screenlogic/number.py
new file mode 100644
index 00000000000..68b9f1d62c3
--- /dev/null
+++ b/homeassistant/components/screenlogic/number.py
@@ -0,0 +1,78 @@
+"""Support for a ScreenLogic number entity."""
+import logging
+
+from screenlogicpy.const import BODY_TYPE, DATA as SL_DATA, EQUIPMENT, SCG
+
+from homeassistant.components.number import NumberEntity
+
+from . import ScreenlogicEntity
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+SUPPORTED_SCG_NUMBERS = (
+    "scg_level1",
+    "scg_level2",
+)
+
+
+async def async_setup_entry(hass, config_entry, async_add_entities):
+    """Set up entry."""
+    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    equipment_flags = coordinator.data[SL_DATA.KEY_CONFIG]["equipment_flags"]
+    if equipment_flags & EQUIPMENT.FLAG_CHLORINATOR:
+        async_add_entities(
+            [
+                ScreenLogicNumber(coordinator, scg_level)
+                for scg_level in coordinator.data[SL_DATA.KEY_SCG]
+                if scg_level in SUPPORTED_SCG_NUMBERS
+            ]
+        )
+
+
+class ScreenLogicNumber(ScreenlogicEntity, NumberEntity):
+    """Class to represent a ScreenLogic Number."""
+
+    def __init__(self, coordinator, data_key, enabled=True):
+        """Initialize of the entity."""
+        super().__init__(coordinator, data_key, enabled)
+        self._body_type = SUPPORTED_SCG_NUMBERS.index(self._data_key)
+        self._attr_max_value = SCG.LIMIT_FOR_BODY[self._body_type]
+        self._attr_name = f"{self.gateway_name} {self.sensor['name']}"
+        self._attr_unit_of_measurement = self.sensor["unit"]
+
+    @property
+    def value(self) -> float:
+        """Return the current value."""
+        return self.sensor["value"]
+
+    async def async_set_value(self, value: float) -> None:
+        """Update the current value."""
+        # Need to set both levels at the same time, so we gather
+        # both existing level values and override the one that changed.
+        levels = {}
+        for level in SUPPORTED_SCG_NUMBERS:
+            levels[level] = self.coordinator.data[SL_DATA.KEY_SCG][level]["value"]
+        levels[self._data_key] = int(value)
+
+        if await self.coordinator.gateway.async_set_scg_config(
+            levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.POOL]],
+            levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.SPA]],
+        ):
+            _LOGGER.debug(
+                "Set SCG to %i, %i",
+                levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.POOL]],
+                levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.SPA]],
+            )
+            await self._async_refresh()
+        else:
+            _LOGGER.warning(
+                "Failed to set_scg to %i, %i",
+                levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.POOL]],
+                levels[SUPPORTED_SCG_NUMBERS[BODY_TYPE.SPA]],
+            )
+
+    @property
+    def sensor(self) -> dict:
+        """Shortcut to access the level sensor data."""
+        return self.coordinator.data[SL_DATA.KEY_SCG][self._data_key]
diff --git a/homeassistant/components/screenlogic/sensor.py b/homeassistant/components/screenlogic/sensor.py
index cd7ba068be0..7ab20164400 100644
--- a/homeassistant/components/screenlogic/sensor.py
+++ b/homeassistant/components/screenlogic/sensor.py
@@ -34,8 +34,6 @@ SUPPORTED_CHEM_SENSORS = (
 )
 
 SUPPORTED_SCG_SENSORS = (
-    "scg_level1",
-    "scg_level2",
     "scg_salt_ppm",
     "scg_super_chlor_timer",
 )