From 043402ad2f02eeb5e48c88ca9f5f0437050d14e2 Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Thu, 25 Jul 2024 00:24:07 -0500
Subject: [PATCH] changed datetime import

---
 src/room_control/model.py | 30 +++++++++++++++---------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/room_control/model.py b/src/room_control/model.py
index ef3b1fd..c70c348 100644
--- a/src/room_control/model.py
+++ b/src/room_control/model.py
@@ -1,21 +1,21 @@
-from datetime import datetime, time, timedelta
+import datetime
 from pathlib import Path
 from typing import Annotated, Dict, List, Optional, Self
 
 import yaml
 from astral import SunDirection
-from pydantic import BaseModel, BeforeValidator, Field, root_validator
+from pydantic import BaseModel, BeforeValidator, Field, model_validator, root_validator
 from pydantic_core import PydanticCustomError
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.table import Column, Table
 
 
-def str_to_timedelta(input_str: str) -> timedelta:
+def str_to_timedelta(input_str: str) -> datetime.timedelta:
     try:
         hours, minutes, seconds = map(int, input_str.split(':'))
-        return timedelta(hours=hours, minutes=minutes, seconds=seconds)
+        return datetime.timedelta(hours=hours, minutes=minutes, seconds=seconds)
     except Exception:
-        return timedelta()
+        return datetime.timedelta()
 
 
 def str_to_direction(input_str: str) -> SunDirection:
@@ -27,7 +27,7 @@ def str_to_direction(input_str: str) -> SunDirection:
         )
 
 
-OffDuration = Annotated[timedelta, BeforeValidator(str_to_timedelta)]
+OffDuration = Annotated[datetime.timedelta, BeforeValidator(str_to_timedelta)]
 
 
 class State(BaseModel):
@@ -45,13 +45,13 @@ class ApplyKwargs(BaseModel):
 
 
 class ControllerStateConfig(BaseModel):
-    time: Optional[str | datetime] = None
+    time: Optional[str | datetime.time | datetime.datetime] = None
     elevation: Optional[float] = None
     direction: Optional[Annotated[SunDirection, BeforeValidator(str_to_direction)]] = None
     off_duration: Optional[OffDuration] = None
-    scene: dict[str, State] | str
+    scene: dict[str, State] | str = Field(default_factory=dict)
 
-    @root_validator(pre=True)
+    @model_validator(mode='before')
     def check_args(cls, values):
         time, elevation = values.get('time'), values.get('elevation')
         if time is not None and elevation is not None:
@@ -60,8 +60,8 @@ class ControllerStateConfig(BaseModel):
             raise PydanticCustomError('no_sun_dir', 'Needs sun direction with elevation')
         return values
 
-    def to_apply_kwargs(self, **kwargs):
-        return ApplyKwargs(entities=self.scene, **kwargs).model_dump(exclude_none=True)
+    def to_apply_kwargs(self, transition: int = None):
+        return ApplyKwargs(entities=self.scene, transition=transition).model_dump(exclude_none=True)
 
 
 class RoomControllerConfig(BaseModel):
@@ -102,11 +102,11 @@ class RoomControllerConfig(BaseModel):
     def sort_states(self):
         """Should only be called after all the times have been resolved"""
         assert all(
-            isinstance(state.time, time) for state in self.states
+            isinstance(state.time, datetime.time) for state in self.states
         ), 'Times have not all been resolved yet'
         self.states = sorted(self.states, key=lambda s: s.time, reverse=True)
 
-    def current_state(self, now: time) -> ControllerStateConfig:
+    def current_state(self, now: datetime.time) -> ControllerStateConfig:
         self.sort_states()
         for state in self.states:
             if state.time <= now:
@@ -114,11 +114,11 @@ class RoomControllerConfig(BaseModel):
         else:
             return self.states[0]
 
-    def current_scene(self, now: time) -> Dict:
+    def current_scene(self, now: datetime.time) -> Dict:
         state = self.current_state(now)
         return state.scene
 
-    def current_off_duration(self, now: time) -> timedelta:
+    def current_off_duration(self, now: datetime.time) -> datetime.timedelta:
         state = self.current_state(now)
         if state.off_duration is None:
             if self.off_duration is None: