8000 Add emeter support for strip sockets (#203) · python-kasa/python-kasa@94e5a90 · GitHub
[go: up one dir, main page]
Skip to content

Commit 94e5a90

Browse files
bdracobrendanburnsrytilahti
authored
Add emeter support for strip sockets (#203)
* Add support for plugs with emeters. * Tweaks for emeter * black * tweaks * tweaks * more tweaks * dry * flake8 * flake8 * legacy typing * Update kasa/smartstrip.py Co-authored-by: Teemu R. <tpr@iki.fi> * reduce * remove useless delegation * tweaks * tweaks * dry * tweak * tweak * tweak * tweak * update tests * wrap * preen * prune * prune * prune * guard * adjust * robust * prune * prune * reduce dict lookups by 1 * Update kasa/smartstrip.py Co-authored-by: Teemu R. <tpr@iki.fi> * delete utils * isort Co-authored-by: Brendan Burns <brendan.d.burns@gmail.com> Co-authored-by: Teemu R. <tpr@iki.fi>
1 parent d720288 commit 94e5a90

File tree

4 files changed

+112
-101
lines changed

4 files changed

+112
-101
lines changed

‎kasa/smartdevice.py

Lines changed: 37 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,14 @@
1111
You may obtain a copy of the license at
1212
http://www.apache.org/licenses/LICENSE-2.0
1313
"""
14+
import collections.abc
1415
import functools
1516
import inspect
1617
import logging
1718
from dataclasses import dataclass
1819
from datetime import datetime, timedelta
1920
from enum import Enum, auto
20-
from typing import Any, Dict, List, Optional
21+
from typing import Any, Dict, List, Optional, Set
2122

2223
from .emeterstatus import EmeterStatus
2324
from .exceptions import SmartDeviceException
@@ -51,6 +52,16 @@ class WifiNetwork:
5152
rssi: Optional[int] = None
5253

5354

55+
def merge(d, u):
56+
"""Update dict recursively."""
57+
for k, v in u.items():
58+
if isinstance(v, collections.abc.Mapping):
59+
d[k] = merge(d.get(k, {}), v)
60+
else:
61+
d[k] = v
62+
return d
63+
64+
5465
def requires_update(f):
5566
"""Indicate that `update` should be called before accessing this method.""" # noqa: D202
5667
if inspect.iscoroutinefunction(f):
@@ -204,6 +215,11 @@ def _create_request(
204215

205216
return request
206217

218+
def _verify_emeter(self) -> None:
219+
"""Raise an exception if there is no emeter."""
220+
if not self.has_emeter:
221+
raise SmartDeviceException("Device has no emeter")
222+
207223
async def _query_helper(
208224
self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
209225
) -> Any:
@@ -240,13 +256,17 @@ async def _query_helper(
240256

241257
return result
242258

259+
@property # type: ignore
260+
@requires_update
261+
def features(self) -> Set[str]:
262+
"""Return a set of features that the device supports."""
263+
return set(self.sys_info["feature"< B41A /span>].split(":"))
264+
243265
@property # type: ignore
244266
@requires_update
245267
def has_emeter(self) -> bool:
246268
"""Return True if device has an energy meter."""
247-
sys_info = self.sys_info
248-
features = sys_info["feature"].split(":")
249-
return "ENE" in features
269+
return "ENE" in self.features
250270

251271
async def get_sys_info(self) -> Dict[str, Any]:
252272
"""Retrieve system information."""
@@ -374,10 +394,8 @@ def location(self) -> Dict:
374394
@requires_update
375395
def rssi(self) -> Optional[int]:
376396
"""Return WiFi signal strenth (rssi)."""
377-
sys_info = self.sys_info
378-
if "rssi" in sys_info:
379-
return int(sys_info["rssi"])
380-
return None
397+
rssi = self.sys_info.get("rssi")
398+
return None if rssi is None else int(rssi)
381399

382400
@property # type: ignore
383401
@requires_update
@@ -410,16 +428,12 @@ async def set_mac(self, mac):
410428
@requires_update
411429
def emeter_realtime(self) -> EmeterStatus:
412430
"""Return current energy readings."""
413-
if not self.has_emeter:
414-
raise SmartDeviceException("Device has no emeter")
415-
431+
self._verify_emeter()
416432
return EmeterStatus(self._last_update[self.emeter_type]["get_realtime"])
417433

418434
async def get_emeter_realtime(self) -> EmeterStatus:
419435
"""Retrieve current energy readings."""
420-
if not self.has_emeter:
421-
raise SmartDeviceException("Device has no emeter")
422-
436+
self._verify_emeter()
423437
return EmeterStatus(await self._query_helper(self.emeter_type, "get_realtime"))
424438

425439
def _create_emeter_request(self, year: int = None, month: int = None):
@@ -429,23 +443,12 @@ def _create_emeter_request(self, year: int = None, month: int = None):
429443
if month is None:
430444
month = datetime.now().month
431445

432-
import collections.abc
433-
434-
def update(d, u):
435-
"""Update dict recursively."""
436-
for k, v in u.items():
437-
if isinstance(v, collections.abc.Mapping):
438-
d[k] = update(d.get(k, {}), v)
439-
else:
440-
d[k] = v
441-
return d
442-
443446
req: Dict[str, Any] = {}
444-
update(req, self._create_request(self.emeter_type, "get_realtime"))
445-
update(
447+
merge(req, self._create_request(self.emeter_type, "get_realtime"))
448+
merge(
446449
req, self._create_request(self.emeter_type, "get_monthstat", {"year": year})
447450
)
448-
update(
451+
merge(
449452
req,
450453
self._create_request(
451454
self.emeter_type, "get_daystat", {"month": month, "year": year}
@@ -458,9 +461,7 @@ def update(d, u):
458461
@requires_update
459462
def emeter_today(self) -> Optional[float]:
460463
"""Return today's energy consumption in kWh."""
461-
if not self.has_emeter:
462-
raise SmartDeviceException("Device has no emeter")
463-
464+
self._verify_emeter()
464465
raw_data = self._last_update[self.emeter_type]["get_daystat"]["day_list"]
465466
data = self._emeter_convert_emeter_data(raw_data)
466467
today = datetime.now().day
@@ -474,9 +475,7 @@ def emeter_today(self) -> Optional[float]:
474475
@requires_update
475476
def emeter_this_month(self) -> Optional[float]:
476477
"""Return this month's energy consumption in kWh."""
477-
if not self.has_emeter:
478-
raise SmartDeviceException("Device has no emeter")
479-
478+
self._verify_emeter()
480479
raw_data = self._last_update[self.emeter_type]["get_monthstat"]["month_list"]
481480
data = self._emeter_convert_emeter_data(raw_data)
482481
current_month = datetime.now().month
@@ -516,9 +515,7 @@ async def get_emeter_daily(
516515
:param kwh: return usage in kWh (default: True)
517516
:return: mapping of day of month to value
518517
"""
519-
if not self.has_emeter:
520-
raise SmartDeviceException("Device has no emeter")
521-
518+
self._verify_emeter()
522519
if year is None:
523520
year = datetime.now().year
524521
if month is None:
@@ -538,9 +535,7 @@ async def get_emeter_monthly(self, year: int = None, kwh: bool = True) -> Dict:
538535
:param kwh: return usage in kWh (default: True)
539536
:return: dict: mapping of month to value
540537
"""
541-
if not self.has_emeter:
542-
raise SmartDeviceException("Device has no emeter")
543-
538+
self._verify_emeter()
544539
if year is None:
545540
year = datetime.now().year
546541

@@ -553,17 +548,13 @@ async def get_emeter_monthly(self, year: int = None, kwh: bool = True) -> Dict:
553548
@requires_update
554549
async def erase_emeter_stats(self) -> Dict:
555550
"""Erase energy meter statistics."""
556-
if not self.has_emeter:
557-
raise SmartDeviceException("Device has no emeter")
558-
551+
self._verify_emeter()
559552
return await self._query_helper(self.emeter_type, "erase_emeter_stat", None)
560553

561554
@requires_update
562555
async def current_consumption(self) -> float:
563556
"""Get the current power consumption in Watt."""
564-
if not self.has_emeter:
565-
raise SmartDeviceException("Device has no emeter")
566-
557+
self._verify_emeter()
567558
response = EmeterStatus(await self.get_emeter_realtime())
568559
return float(response["power"])
569560

‎kasa/smartstrip.py

Lines changed: 73 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from kasa.smartdevice import (
88
DeviceType,
9+
EmeterStatus,
910
SmartDevice,
1011
SmartDeviceException,
1112
requires_update,
@@ -15,6 +16,15 @@
1516
_LOGGER = logging.getLogger(__name__)
1617

1718

19+
def merge_sums(dicts):
20+
"""Merge the sum of dicts."""
21+
total_dict: DefaultDict[int, float] = defaultdict(lambda: 0.0)
22+
for sum_dict in dicts:
23+
for day, value in sum_dict.items():
24+
total_dict[day] += value
25+
return total_dict
26+
27+
1828
class SmartStrip(SmartDevice):
1929
"""Representation of a TP-Link Smart Power Strip.
2030
@@ -75,11 +85,7 @@ def __init__(self, host: str) -> None:
7585
@requires_update
7686
def is_on(self) -> bool:
7787
"""Return if any of the outlets are on."""
78-
for plug in self.children:
79-
is_on = plug.is_on
80-
if is_on:
81-
return True
82-
return False
88+
return any(plug.is_on for plug in self.children)
8389

8490
async def update(self):
8591
"""Update some of the attributes.
@@ -97,6 +103,10 @@ async def update(self):
97103
SmartStripPlug(self.host, parent=self, child_id=child["id"])
98104
)
99105

106+
if self.has_emeter:
107+
for plug in self.children:
108+
await plug.update()
109+
100110
async def turn_on(self, **kwargs):
101111
"""Turn the strip on."""
102112
await self._query_helper("system", "set_relay_state", {"state": 1})
@@ -140,16 +150,16 @@ def state_information(self) -> Dict[str, Any]:
140150

141151
async def current_consumption(self) -> float:
142152
"""Get the current power consumption in watts."""
143-
consumption = sum(await plug.current_consumption() for plug in self.children)
144-
145-
return consumption
153+
return sum([await plug.current_consumption() for plug in self.children])
146154

147-
async def set_alias(self, alias: str) -> None:
148-
"""Set the alias for the strip.
149-
150-
:param alias: new alias
151-
"""
152-
return await super().set_alias(alias)
155+
@requires_update
156+
async def get_emeter_realtime(self) -> EmeterStatus:
157+
"""Retrieve current energy readings."""
158+
emeter_rt = await self._async_get_emeter_sum("get_emeter_realtime", {})
159+
# Voltage is averaged since each read will result
160+
# in a slightly different voltage since they are not atomic
161+
emeter_rt["voltage_mv"] = int(emeter_rt["voltage_mv"] / len(self.children))
162+
return EmeterStatus(emeter_rt)
153163

154164
@requires_update
155165
async def get_emeter_daily(
@@ -163,14 +173,9 @@ async def get_emeter_daily(
163173
:param kwh: return usage in kWh (default: True)
164174
:return: mapping of day of month to value
165175
"""
166-
emeter_daily: DefaultDict[int, float] = defaultdict(lambda: 0.0)
167-
for plug in self.children:
168-
plug_emeter_daily = await plug.get_emeter_daily(
169-
year=year, month=month, kwh=kwh
170-
)
171-
for day, value in plug_emeter_daily.items():
172-
emeter_daily[day] += value
173-
return emeter_daily
176+
return await self._async_get_emeter_sum(
177+
"get_emeter_daily", {"year": year, "month": month, "kwh": kwh}
178+
)
174179

175180
@requires_update
176181
async def get_emeter_monthly(self, year: int = None, kwh: bool = True) -> Dict:
@@ -179,20 +184,45 @@ async def get_emeter_monthly(self, year: int = None, kwh: bool = True) -> Dict:
179184
:param year: year for which to retrieve statistics (default: this year)
180185
:param kwh: return usage in kWh (default: True)
181186
"""
182-
emeter_monthly: DefaultDict[int, float] = defaultdict(lambda: 0.0)
183-
for plug in self.children:
184-
plug_emeter_monthly = await plug.get_emeter_monthly(year=year, kwh=kwh)
185-
for month, value in plug_emeter_monthly:
186-
emeter_monthly[month] += value
187+
return await self._async_get_emeter_sum(
188+
"get_emeter_monthly", {"year": year, "kwh": kwh}
189+
)
187190

188-
return emeter_monthly
191+
async def _async_get_emeter_sum(self, func: str, kwargs: Dict[str, Any]) -> Dict:
192+
"""Retreive emeter stats for a time period from children."""
193+
self._verify_emeter()
194+
return merge_sums(
195+
[await getattr(plug, func)(**kwargs) for plug in self.children]
196+
)
189197

190198
@requires_update
191199
async def erase_emeter_stats(self):
192200
"""Erase energy meter statistics for all plugs."""
193201
for plug in self.children:
194202
await plug.erase_emeter_stats()
195203

204+
@property # type: ignore
205+
@requires_update
206+
def emeter_this_month(self) -> Optional[float]:
207+
"""Return this month's energy consumption in kWh."""
208+
return sum([plug.emeter_this_month for plug in self.children])
209+
210+
@property # type: ignore
211+
@requires_update
212+
def emeter_today(self) -> Optional[float]:
213+
"""Return this month's energy consumption in kWh."""
214+
return sum([plug.emeter_today for plug in self.children])
215+
216+
@property # type: ignore
217+
@requires_update
218+
def emeter_realtime(self) -> EmeterStatus:
219+
"""Return current energy readings."""
220+
emeter = merge_sums([plug.emeter_realtime for plug in self.children])
221+
# Voltage is averaged since each read will result
222+
# in a slightly different voltage since they are not atomic
223+
emeter["voltage_mv"] = int(emeter["voltage_mv"] / len(self.children))
224+
return EmeterStatus(emeter)
225+
196226

197227
class SmartStripPlug(SmartPlug):
198228
"""Representation of a single socket in a power strip.
@@ -214,12 +244,22 @@ def __init__(self, host: str, parent: "SmartStrip", child_id: str) -> None:
214244
self._device_type = DeviceType.StripSocket
215245

216246
async def update(self):
217-
"""Override the update to no-op and inform the user."""
218-
_LOGGER.warning(
219-
"You called update() on a child device, which has no effect."
220-
"Call update() on the parent device instead."
247+
"""Query the device to update the data.
248+
249+
Needed for properties that are decorated with `requires_update`.
250+
"""
251+
self._last_update = await self.parent.protocol.query(
252+
self.host, self._create_emeter_request()
221253
)
222-
return
254+
255+
def _create_request(
256+
self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
257+
):
258+
request: Dict[str, Any] = {
259+
"context": {"child_ids": [self.child_id]},
260+
target: {cmd: arg},
261+
}
262+
return request
223263

224264
async def _query_helper(
225265
self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
@@ -245,12 +285,6 @@ def led(self) -> bool:
245285
"""
246286
return False
247287

248-
@property # type: ignore
249-
@requires_update
250-
def has_emeter(self) -> bool:
251-
"""Children have no emeter to my knowledge."""
252-
return False
253-
254288
@property # type: ignore
255289
@requires_update
256290
def device_id(self) -> str:

0 commit comments

Comments
 (0)
0