diff --git a/python_snoo/baby.py b/python_snoo/baby.py index 6a3895b..4ced743 100644 --- a/python_snoo/baby.py +++ b/python_snoo/baby.py @@ -1,8 +1,8 @@ from datetime import datetime -from python_snoo.containers import Activity, BabyData, BreastfeedingActivity, DiaperActivity, DiaperTypes -from python_snoo.exceptions import SnooBabyError -from python_snoo.snoo import Snoo +from .containers import Activity, BabyData, BreastfeedingActivity, DiaperActivity, DiaperTypes +from .exceptions import SnooBabyError +from .snoo import Snoo class Baby: diff --git a/python_snoo/containers.py b/python_snoo/containers.py index 53972a5..e6f8283 100644 --- a/python_snoo/containers.py +++ b/python_snoo/containers.py @@ -1,6 +1,6 @@ import dataclasses import datetime -from enum import StrEnum +from enum import IntEnum, StrEnum from typing import Any, Union from mashumaro.mixins.json import DataClassJSONMixin @@ -16,6 +16,45 @@ class SnooLevels(StrEnum): stop = "ONLINE" +class SnooNoiseTimeoutLevels(IntEnum): + _5_minutes = 5 + _10_minutes = 10 + _15_minutes = 15 + _20_minutes = 20 + _25_minutes = 25 + _30_minutes = 30 + _35_minutes = 35 + _40_minutes = 40 + _45_minutes = 45 + _50_minutes = 50 + _55_minutes = 55 + _60_minutes = 60 + _65_minutes = 65 + _70_minutes = 70 + _75_minutes = 75 + _80_minutes = 80 + _85_minutes = 85 + _90_minutes = 90 + _95_minutes = 95 + _100_minutes = 100 + _105_minutes = 105 + _110_minutes = 110 + _115_minutes = 115 + _120_minutes = 120 + _125_minutes = 125 + _130_minutes = 130 + _135_minutes = 135 + _140_minutes = 140 + _145_minutes = 145 + _150_minutes = 150 + _155_minutes = 155 + _160_minutes = 160 + _165_minutes = 165 + _170_minutes = 170 + _175_minutes = 175 + _180_minutes = 180 + + class SnooStates(StrEnum): baseline = "BASELINE" level1 = "LEVEL1" diff --git a/python_snoo/snoo.py b/python_snoo/snoo.py index b78add9..f2895d5 100644 --- a/python_snoo/snoo.py +++ b/python_snoo/snoo.py @@ -13,13 +13,7 @@ from pubnub.pnconfiguration import PNConfiguration from pubnub.pubnub_asyncio import PubNubAsyncio -from .containers import ( - AuthorizationInfo, - BabyData, - SnooData, - SnooDevice, - SnooStates, -) +from .containers import AuthorizationInfo, BabyData, SnooData, SnooDevice, SnooNoiseTimeoutLevels, SnooStates from .exceptions import InvalidSnooAuth, SnooAuthException, SnooBabyError, SnooCommandException, SnooDeviceError from .pubnub_async import SnooPubNub @@ -233,11 +227,25 @@ async def set_level(self, device: SnooDevice, level: SnooStates, hold: bool = Fa await self.send_command("go_to_state", device, **{"state": level.value, "hold": hold}) - async def set_sticky_white_noise(self, device: SnooDevice, on: bool): + async def set_sticky_white_noise( + self, + device: SnooDevice, + on: bool, + timeout_value: SnooNoiseTimeoutLevels | int = SnooNoiseTimeoutLevels._15_minutes, + ): + """Enable or disable sticky white noise for a device. + + Args: + device: The SnooDevice to control. + on: True to turn sticky white noise on, False to turn it off. + timeout_value: How long the white noise should remain on before timing out. + Must be a member of the SnooNoiseTimeoutLevels enum. The default is + SnooNoiseTimeoutLevels._15_minutes. + """ await self.send_command( "set_sticky_white_noise", device, - **{"state": "on" if on else "off", "timeout_min": 15}, + **{"state": "on" if on else "off", "timeout_min": timeout_value}, ) async def get_status(self, device: SnooDevice):