From 7137e309b3d11556b235a0356a5f3992ee5f6361 Mon Sep 17 00:00:00 2001
From: denpamusic <denpa@denpa.pro>
Date: Sat, 7 Oct 2023 01:07:34 +0300
Subject: [PATCH] Allow extra states in schedules.

This allows the use of "day" and "night" states, thus improving code readability.
---
 pyplumio/helpers/schedule.py   | 21 +++++++++++++++++----
 tests/helpers/test_schedule.py |  4 ++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/pyplumio/helpers/schedule.py b/pyplumio/helpers/schedule.py
index d5d460ab..0f1d4133 100644
--- a/pyplumio/helpers/schedule.py
+++ b/pyplumio/helpers/schedule.py
@@ -16,8 +16,18 @@
 START_OF_DAY: Final = "00:00"
 END_OF_DAY: Final = "00:00"
 
-STATE_NIGHT: Final = STATE_OFF
-STATE_DAY: Final = STATE_ON
+STATE_NIGHT: Final = "night"
+STATE_DAY: Final = "day"
+
+ENABLED: Final[list[str]] = [
+    STATE_ON,
+    STATE_DAY,
+]
+
+DISABLED: Final[list[str]] = [
+    STATE_OFF,
+    STATE_NIGHT,
+]
 
 
 def _parse_interval(start: str, end: str) -> tuple[int, int]:
@@ -77,14 +87,17 @@ def append(self, item) -> None:
 
     def set_state(
         self,
-        state: Literal["off", "on"],
+        state: Literal["off", "on", "day", "night"],
         start: str = START_OF_DAY,
         end: str = END_OF_DAY,
     ) -> None:
         """Set state for interval."""
+        if state not in [*ENABLED, *DISABLED]:
+            raise ValueError(f'state "{state}" is not allowed')
+
         index, stop_index = _parse_interval(start, end)
         while index < stop_index:
-            self._intervals[index] = state == STATE_ON
+            self._intervals[index] = state in ENABLED
             index += 1
 
     def set_on(self, start: str = START_OF_DAY, end: str = END_OF_DAY) -> None:
diff --git a/tests/helpers/test_schedule.py b/tests/helpers/test_schedule.py
index 55645286..e36d9c1b 100644
--- a/tests/helpers/test_schedule.py
+++ b/tests/helpers/test_schedule.py
@@ -55,6 +55,10 @@ def test_schedule_day(schedule_day: ScheduleDay) -> None:
         with pytest.raises(ValueError):
             schedule_day.set_state(STATE_ON, start, end)
 
+    # Test with incorrect state.
+    with pytest.raises(ValueError):
+        schedule_day.set_state("invalid_state", "00:00", "01:00")  # type: ignore [arg-type]
+
     # set whole day schedule.
     schedule_day.set_on()
     assert schedule_day.intervals == [True for _ in range(48)]