Move some code around (#76)

raman325 · Mar 7, 2024 · 4124740 · 4124740
1 parent e5f137f
commit 4124740
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 50 deletions.
diff --git a/custom_components/lock_code_manager/config_flow.py b/custom_components/lock_code_manager/config_flow.py
@@ -9,6 +9,7 @@
 import voluptuous as vol
 
 from homeassistant import config_entries
+from homeassistant.components.calendar import DOMAIN as CALENDAR_DOMAIN
 from homeassistant.components.lock import DOMAIN as LOCK_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ENABLED, CONF_NAME, CONF_PIN
@@ -23,19 +24,50 @@
 
 from . import providers
 from .const import (
+    CONF_CALENDAR,
     CONF_LOCKS,
     CONF_NUM_SLOTS,
+    CONF_NUMBER_OF_USES,
     CONF_SLOTS,
     CONF_START_SLOT,
     DEFAULT_NUM_SLOTS,
     DEFAULT_START,
     DOMAIN,
 )
 from .data import get_entry_data
-from .helpers import CODE_SLOT_SCHEMA, CODE_SLOTS_SCHEMA, UI_CODE_SLOT_SCHEMA
 
 _LOGGER = logging.getLogger(__name__)
 
+UI_CODE_SLOT_SCHEMA = vol.Schema(
+    {
+        vol.Optional(CONF_NAME): cv.string,
+        vol.Optional(CONF_PIN): cv.string,
+        vol.Required(CONF_ENABLED, default=True): cv.boolean,
+        vol.Optional(CONF_CALENDAR): sel.EntitySelector(
+            sel.EntitySelectorConfig(domain=CALENDAR_DOMAIN)
+        ),
+        vol.Optional(CONF_NUMBER_OF_USES): sel.TextSelector(
+            sel.TextSelectorConfig(type=sel.TextSelectorType.NUMBER)
+        ),
+    }
+)
+
+CODE_SLOT_SCHEMA = UI_CODE_SLOT_SCHEMA.extend(
+    {vol.Optional(CONF_NUMBER_OF_USES): vol.Coerce(int)}
+)
+
+
+def enabled_requires_pin(data: dict[str, Any]) -> dict[str, Any]:
+    """Validate that if enabled is True, pin is set."""
+    if any(val.get(CONF_ENABLED) and not val.get(CONF_PIN) for val in data.values()):
+        raise vol.Invalid("PIN must be set if enabled is True")
+    return data
+
+
+CODE_SLOTS_SCHEMA = vol.All(
+    vol.Schema({vol.Coerce(int): CODE_SLOT_SCHEMA}), enabled_requires_pin
+)
+
 LOCKS_FILTER_CONFIG = [
     sel.EntityFilterSelectorConfig(integration=integration, domain=LOCK_DOMAIN)
     for integration in [

diff --git a/custom_components/lock_code_manager/helpers.py b/custom_components/lock_code_manager/helpers.py
@@ -5,64 +5,18 @@
 import logging
 from typing import Any
 
-import voluptuous as vol
-
-from homeassistant.components.calendar import DOMAIN as CALENDAR_DOMAIN
 from homeassistant.components.lock import DOMAIN as LOCK_DOMAIN
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    ATTR_AREA_ID,
-    ATTR_DEVICE_ID,
-    ATTR_ENTITY_ID,
-    CONF_ENABLED,
-    CONF_NAME,
-    CONF_PIN,
-)
+from homeassistant.const import ATTR_AREA_ID, ATTR_DEVICE_ID, ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import (
-    config_validation as cv,
-    device_registry as dr,
-    entity_registry as er,
-    selector as sel,
-)
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from .const import CONF_CALENDAR, CONF_LOCKS, CONF_NUMBER_OF_USES, DOMAIN
+from .const import CONF_LOCKS, DOMAIN
 from .providers import INTEGRATIONS_CLASS_MAP, BaseLock
 
 _LOGGER = logging.getLogger(__name__)
 
 
-UI_CODE_SLOT_SCHEMA = vol.Schema(
-    {
-        vol.Optional(CONF_NAME): cv.string,
-        vol.Optional(CONF_PIN): cv.string,
-        vol.Required(CONF_ENABLED, default=True): cv.boolean,
-        vol.Optional(CONF_CALENDAR): sel.EntitySelector(
-            sel.EntitySelectorConfig(domain=CALENDAR_DOMAIN)
-        ),
-        vol.Optional(CONF_NUMBER_OF_USES): sel.TextSelector(
-            sel.TextSelectorConfig(type=sel.TextSelectorType.NUMBER)
-        ),
-    }
-)
-
-CODE_SLOT_SCHEMA = UI_CODE_SLOT_SCHEMA.extend(
-    {vol.Optional(CONF_NUMBER_OF_USES): vol.Coerce(int)}
-)
-
-
-def enabled_requires_pin(data: dict[str, Any]) -> dict[str, Any]:
-    """Validate that if enabled is True, pin is set."""
-    if any(val.get(CONF_ENABLED) and not val.get(CONF_PIN) for val in data.values()):
-        raise vol.Invalid("PIN must be set if enabled is True")
-    return data
-
-
-CODE_SLOTS_SCHEMA = vol.All(
-    vol.Schema({vol.Coerce(int): CODE_SLOT_SCHEMA}), enabled_requires_pin
-)
-
-
 @callback
 def async_create_lock_instance(
     hass: HomeAssistant,