expanded pydantic model

This commit is contained in:
John Lancaster
2024-03-10 18:17:14 -05:00
parent 541475d68b
commit 9845368159
2 changed files with 71 additions and 32 deletions

View File

@@ -1,38 +1,77 @@
from typing import Dict from datetime import datetime, timedelta
from pathlib import Path
from typing import Annotated, Dict, List, Self
from pydantic import ( import yaml
BaseModel, from astral import SunDirection
ValidationError, from pydantic import BaseModel, BeforeValidator, ValidationError, conint, root_validator
field_validator,
)
def validate_int(v): def str_to_timedelta(input_str: str) -> timedelta:
if not len(bytes(v)) == 1: try:
raise ValidationError() hours, minutes, seconds = map(int, input_str.split(':'))
return timedelta(hours=hours, minutes=minutes, seconds=seconds)
except Exception:
return timedelta()
def str_to_direction(input_str: str) -> SunDirection:
if input_str.lower() == 'setting':
return SunDirection.SETTING
elif input_str == 'rising':
return SunDirection.RISING
else:
raise ValidationError(f'Invalid sun direction: {input_str}')
OffDuration = Annotated[timedelta, BeforeValidator(str_to_timedelta)]
class State(BaseModel): class State(BaseModel):
state: bool = True state: bool = True
brightness: int = None brightness: conint(ge=1, le=255) = None
color_temp: int = None color_temp: conint(ge=200, le=650) = None
@field_validator('brightness')
@classmethod
def validate_brightness(cls, v: int) -> int:
assert 0 <= v <= 255
return v
@field_validator('color_temp')
@classmethod
def validate_color_temp(cls, v: int) -> int:
assert 200 <= v <= 600
return v
# Scene = RootModel[Dict[str, State]] class ApplyKwargs(BaseModel):
"""Arguments to call with the 'scene/apply' service"""
class ApplyScene(BaseModel):
entities: Dict[str, State] entities: Dict[str, State]
transition: int = None transition: int = None
class ControllerStateConfig(BaseModel):
time: str | datetime = None
elevation: float = None
direction: Annotated[SunDirection, BeforeValidator(str_to_direction)] = None
off_duration: OffDuration = None
scene: Dict[str, State]
@root_validator(pre=True)
def check_args(cls, values):
time, elevation = values.get('time'), values.get('elevation')
if time is None and elevation is None:
raise ValueError('Either time or elevation must be set.')
elif time is not None and elevation is not None:
raise ValueError('Only one of time or elevation can be set.')
elif elevation is not None:
assert 'direction' in values
return values
def to_apply_kwargs(self, transition: int = 0):
return ApplyKwargs(entities=self.scene, transition=transition).model_dump(exclude_none=True)
class RoomConfig(BaseModel):
states: List[ControllerStateConfig]
off_duration: OffDuration = None
@classmethod
def from_yaml(cls: Self, yaml_path: Path):
yaml_path = Path(yaml_path)
with yaml_path.open('r') as f:
for appname, app_cfg in yaml.load(f, Loader=yaml.SafeLoader).items():
if app_cfg['class'] == 'RoomController':
break
print(app_cfg)
return cls(**app_cfg)

View File

@@ -11,7 +11,7 @@ from appdaemon.plugins.hass.hassapi import Hass
from appdaemon.plugins.mqtt.mqttapi import Mqtt from appdaemon.plugins.mqtt.mqttapi import Mqtt
from astral import SunDirection from astral import SunDirection
from console import console, setup_handler from console import console, setup_handler
from model import ApplyScene from model import ApplyKwargs
from rich.console import Console, ConsoleOptions, RenderResult from rich.console import Console, ConsoleOptions, RenderResult
from rich.table import Column, Table from rich.table import Column, Table
@@ -74,11 +74,11 @@ class RoomState:
collapse_padding=True, collapse_padding=True,
) )
for name, state in self.scene.items(): for name, state in self.scene.items():
table.add_row(name, ApplyScene(entites=self.scene).model_dump(exclude_none=True)) table.add_row(name, ApplyKwargs(entites=self.scene).model_dump(exclude_none=True))
yield table yield table
def scene_model(self) -> ApplyScene: def scene_model(self) -> ApplyKwargs:
return ApplyScene( return ApplyKwargs(
entities=self.scene, entities=self.scene,
transition=0 transition=0
).model_dump(exclude_none=True) ).model_dump(exclude_none=True)
@@ -119,7 +119,7 @@ class RoomConfig:
collapse_padding=True, collapse_padding=True,
) )
for state in self.states: for state in self.states:
scene_json = ApplyScene(entities=state.scene).model_dump(exclude_none=True) scene_json = ApplyKwargs(entities=state.scene).model_dump(exclude_none=True)
lines = [ lines = [
f'{name:20}{state["state"]} Brightness: {state["brightness"]:<4} Temp: {state["color_temp"]}' f'{name:20}{state["state"]} Brightness: {state["brightness"]:<4} Temp: {state["color_temp"]}'
for name, state in scene_json['entities'].items() for name, state in scene_json['entities'].items()
@@ -363,7 +363,7 @@ class RoomController(Hass, Mqtt):
self.log(f'Turned on scene: {scene}') self.log(f'Turned on scene: {scene}')
elif isinstance(scene, dict): elif isinstance(scene, dict):
kwargs = ApplyScene(entities=scene, transition=0).model_dump(exclude_none=True) kwargs = ApplyKwargs(entities=scene, transition=0).model_dump(exclude_none=True)
self.log('Validated scene JSON', level='DEBUG') self.log('Validated scene JSON', level='DEBUG')
self.call_service('scene/apply', **kwargs) self.call_service('scene/apply', **kwargs)
if self.logger.isEnabledFor(logging.INFO): if self.logger.isEnabledFor(logging.INFO):