Skip to content

Commit 79ac954

Browse files
sdb9696sdb9696
authored
Replace custom deviceconfig serialization with mashumaru (#1274)
1 parent bf23f73 commit 79ac954

File tree

6 files changed

+165
-96
lines changed

6 files changed

+165
-96
lines changed

‎kasa/deviceconfig.py‎

Lines changed: 55 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -17,25 +17,32 @@
1717
>>> config_dict = device.config.to_dict()
1818
>>> # DeviceConfig.to_dict() can be used to store for later
1919
>>> print(config_dict)
20-
{'host': '127.0.0.3', 'timeout': 5, 'credentials': Credentials(), 'connection_type'\
21-
: {'device_family': 'SMART.TAPOBULB', 'encryption_type': 'KLAP', 'https': False, \
22-
'login_version': 2}, 'uses_http': True}
20+
{'host': '127.0.0.3', 'timeout': 5, 'credentials': {'username': 'user@example.com', \
21+
'password': 'great_password'}, 'connection_type'\
22+
: {'device_family': 'SMART.TAPOBULB', 'encryption_type': 'KLAP', 'login_version': 2, \
23+
'https': False}, 'uses_http': True}
2324
2425
>>> later_device = await Device.connect(config=Device.Config.from_dict(config_dict))
2526
>>> print(later_device.alias) # Alias is available as connect() calls update()
2627
Living Room Bulb
2728
2829
"""
2930

30-
# Module cannot use from __future__ import annotations until migrated to mashumaru
31-
# as dataclass.fields() will not resolve the type.
31+
from __future__ import annotations
32+
3233
import logging
33-
from dataclasses import asdict, dataclass, field, fields, is_dataclass
34+
from dataclasses import dataclass, field, replace
3435
from enum import Enum
35-
from typing import TYPE_CHECKING, Any, Optional, TypedDict
36+
from typing import TYPE_CHECKING, Any, Self, TypedDict
37+
38+
from aiohttp import ClientSession
39+
from mashumaro import field_options
40+
from mashumaro.config import BaseConfig
41+
from mashumaro.types import SerializationStrategy
3642

3743
from .credentials import Credentials
3844
from .exceptions import KasaException
45+
from .json import DataClassJSONMixin
3946

4047
if TYPE_CHECKING:
4148
from aiohttp import ClientSession
@@ -73,45 +80,17 @@ class DeviceFamily(Enum):
7380
SmartIpCamera = "SMART.IPCAMERA"
7481

7582

76-
def _dataclass_from_dict(klass: Any, in_val: dict) -> Any:
77-
if is_dataclass(klass):
78-
fieldtypes = {f.name: f.type for f in fields(klass)}
79-
val = {}
80-
for dict_key in in_val:
81-
if dict_key in fieldtypes:
82-
if hasattr(fieldtypes[dict_key], "from_dict"):
83-
val[dict_key] = fieldtypes[dict_key].from_dict(in_val[dict_key]) # type: ignore[union-attr]
84-
else:
85-
val[dict_key] = _dataclass_from_dict(
86-
fieldtypes[dict_key], in_val[dict_key]
87-
)
88-
else:
89-
raise KasaException(
90-
f"Cannot create dataclass from dict, unknown key: {dict_key}"
91-
)
92-
return klass(**val) # type: ignore[operator]
93-
else:
94-
return in_val
95-
96-
97-
def _dataclass_to_dict(in_val: Any) -> dict:
98-
fieldtypes = {f.name: f.type for f in fields(in_val) if f.compare}
99-
out_val = {}
100-
for field_name in fieldtypes:
101-
val = getattr(in_val, field_name)
102-
if val is None:
103-
continue
104-
elif hasattr(val, "to_dict"):
105-
out_val[field_name] = val.to_dict()
106-
elif is_dataclass(fieldtypes[field_name]):
107-
out_val[field_name] = asdict(val)
108-
else:
109-
out_val[field_name] = val
110-
return out_val
83+
class _DeviceConfigBaseMixin(DataClassJSONMixin):
84+
"""Base class for serialization mixin."""
85+
86+
class Config(BaseConfig):
87+
"""Serialization config."""
88+
89+
omit_none = True
11190

11291

11392
@dataclass
114-
class DeviceConnectionParameters:
93+
class DeviceConnectionParameters(_DeviceConfigBaseMixin):
11594
"""Class to hold the the parameters determining connection type."""
11695

11796
device_family: DeviceFamily
@@ -125,7 +104,7 @@ def from_values(
125104
encryption_type: str,
126105
login_version: int | None = None,
127106
https: bool | None = None,
128-
) -> "DeviceConnectionParameters":
107+
) -> DeviceConnectionParameters:
129108
"""Return connection parameters from string values."""
130109
try:
131110
if https is None:
@@ -142,39 +121,17 @@ def from_values(
142121
+ f"{encryption_type}.{login_version}"
143122
) from ex
144123

145-
@staticmethod
146-
def from_dict(connection_type_dict: dict[str, Any]) -> "DeviceConnectionParameters":
147-
"""Return connection parameters from dict."""
148-
if (
149-
isinstance(connection_type_dict, dict)
150-
and (device_family := connection_type_dict.get("device_family"))
151-
and (encryption_type := connection_type_dict.get("encryption_type"))
152-
):
153-
if login_version := connection_type_dict.get("login_version"):
154-
login_version = int(login_version) # type: ignore[assignment]
155-
return DeviceConnectionParameters.from_values(
156-
device_family,
157-
encryption_type,
158-
login_version, # type: ignore[arg-type]
159-
connection_type_dict.get("https", False),
160-
)
161124

162-
raise KasaException(f"Invalid connection type data for {connection_type_dict}")
125+
class _DoNotSerialize(SerializationStrategy):
126+
def serialize(self, value: Any) -> None:
127+
return None # pragma: no cover
163128

164-
def to_dict(self) -> dict[str, str | int | bool]:
165-
"""Convert connection params to dict."""
166-
result: dict[str, str | int] = {
167-
"device_family": self.device_family.value,
168-
"encryption_type": self.encryption_type.value,
169-
"https": self.https,
170-
}
171-
if self.login_version:
172-
result["login_version"] = self.login_version
173-
return result
129+
def deserialize(self, value: Any) -> None:
130+
return None # pragma: no cover
174131

175132

176133
@dataclass
177-
class DeviceConfig:
134+
class DeviceConfig(_DeviceConfigBaseMixin):
178135
"""Class to represent paramaters that determine how to connect to devices."""
179136

180137
DEFAULT_TIMEOUT = 5
@@ -202,9 +159,12 @@ class DeviceConfig:
202159
#: in order to determine whether they should pass a custom http client if desired.
203160
uses_http: bool = False
204161

205-
# compare=False will be excluded from the serialization and object comparison.
206162
#: Set a custom http_client for the device to use.
207-
http_client: Optional["ClientSession"] = field(default=None, compare=False)
163+
http_client: ClientSession | None = field(
164+
default=None,
165+
compare=False,
166+
metadata=field_options(serialization_strategy=_DoNotSerialize()),
167+
)
208168

209169
aes_keys: KeyPairDict | None = None
210170

@@ -214,22 +174,30 @@ def __post_init__(self) -> None:
214174
DeviceFamily.IotSmartPlugSwitch, DeviceEncryptionType.Xor
215175
)
216176

217-
def to_dict(
177+
def __pre_serialize__(self) -> Self:
178+
return replace(self, http_client=None)
179+
180+
def to_dict_control_credentials(
218181
self,
219182
*,
220183
credentials_hash: str | None = None,
221184
exclude_credentials: bool = False,
222185
) -> dict[str, dict[str, str]]:
223-
"""Convert device config to dict."""
224-
if credentials_hash is not None or exclude_credentials:
225-
self.credentials = None
226-
if credentials_hash:
227-
self.credentials_hash = credentials_hash
228-
return _dataclass_to_dict(self)
186+
"""Convert deviceconfig to dict controlling how to serialize credentials.
187+
188+
If credentials_hash is provided credentials will be None.
189+
If credentials_hash is '' credentials_hash and credentials will be None.
190+
exclude credentials controls whether to include credentials.
191+
The defaults are the same as calling to_dict().
192+
"""
193+
if credentials_hash is None:
194+
if not exclude_credentials:
195+
return self.to_dict()
196+
else:
197+
return replace(self, credentials=None).to_dict()
229198

230-
@staticmethod
231-
def from_dict(config_dict: dict[str, dict[str, str]]) -> "DeviceConfig":
232-
"""Return device config from dict."""
233-
if isinstance(config_dict, dict):
234-
return _dataclass_from_dict(DeviceConfig, config_dict)
235-
raise KasaException(f"Invalid device config data: {config_dict}")
199+
return replace(
200+
self,
201+
credentials_hash=credentials_hash if credentials_hash else None,
202+
credentials=None,
203+
).to_dict()

‎tests/conftest.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import asyncio
44
import sys
55
import warnings
6+
from pathlib import Path
67
from unittest.mock import MagicMock, patch
78

89
import pytest
@@ -21,6 +22,13 @@
2122
turn_on = pytest.mark.parametrize("turn_on", [True, False])
2223

2324

25+
def load_fixture(foldername, filename):
26+
"""Load a fixture."""
27+
path = Path(Path(__file__).parent / "fixtures" / foldername / filename)
28+
with path.open() as fdp:
29+
return fdp.read()
30+
31+
2432
async def handle_turn_on(dev, turn_on):
2533
if turn_on:
2634
await dev.turn_on()

‎tests/fixtures/serialization/deviceconfig_camera-aes-https.json‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
{
2+
"host": "127.0.0.1",
3+
"timeout": 5,
4+
"connection_type": {
5+
"device_family": "SMART.IPCAMERA",
6+
"encryption_type": "AES",
7+
"https": true
8+
},
9+
"uses_http": false
10+
}

‎tests/fixtures/serialization/deviceconfig_plug-klap.json‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
{
2+
"host": "127.0.0.1",
3+
"timeout": 5,
4+
"connection_type": {
5+
"device_family": "SMART.TAPOPLUG",
6+
"encryption_type": "KLAP",
7+
"https": false,
8+
"login_version": 2
9+
},
10+
"uses_http": false
11+
}

‎tests/fixtures/serialization/deviceconfig_plug-xor.json‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
{
2+
"host": "127.0.0.1",
3+
"timeout": 5,
4+
"connection_type": {
5+
"device_family": "IOT.SMARTPLUGSWITCH",
6+
"encryption_type": "XOR",
7+
"https": false
8+
},
9+
"uses_http": false
10+
}

‎tests/test_deviceconfig.py‎

Lines changed: 71 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,97 @@
1+
import json
2+
from dataclasses import replace
13
from json import dumps as json_dumps
24
from json import loads as json_loads
35

46
import aiohttp
57
import pytest
8+
from mashumaro import MissingField
69

710
from kasa.credentials import Credentials
811
from kasa.deviceconfig import (
912
DeviceConfig,
13+
DeviceConnectionParameters,
14+
DeviceEncryptionType,
15+
DeviceFamily,
16+
)
17+
18+
from .conftest import load_fixture
19+
20+
PLUG_XOR_CONFIG = DeviceConfig(host="127.0.0.1")
21+
PLUG_KLAP_CONFIG = DeviceConfig(
22+
host="127.0.0.1",
23+
connection_type=DeviceConnectionParameters(
24+
DeviceFamily.SmartTapoPlug, DeviceEncryptionType.Klap, login_version=2
25+
),
26+
)
27+
CAMERA_AES_CONFIG = DeviceConfig(
28+
host="127.0.0.1",
29+
connection_type=DeviceConnectionParameters(
30+
DeviceFamily.SmartIpCamera, DeviceEncryptionType.Aes, https=True
31+
),
1032
)
11-
from kasa.exceptions import KasaException
1233

1334

1435
async def test_serialization():
36+
"""Test device config serialization."""
1537
config = DeviceConfig(host="Foo", http_client=aiohttp.ClientSession())
1638
config_dict = config.to_dict()
1739
config_json = json_dumps(config_dict)
1840
config2_dict = json_loads(config_json)
1941
config2 = DeviceConfig.from_dict(config2_dict)
2042
assert config == config2
43+
assert config.to_dict_control_credentials() == config.to_dict()
44+
45+
46+
@pytest.mark.parametrize(
47+
("fixture_name", "expected_value"),
48+
[
49+
("deviceconfig_plug-xor.json", PLUG_XOR_CONFIG),
50+
("deviceconfig_plug-klap.json", PLUG_KLAP_CONFIG),
51+
("deviceconfig_camera-aes-https.json", CAMERA_AES_CONFIG),
52+
],
53+
ids=lambda arg: arg.split("_")[-1] if isinstance(arg, str) else "",
54+
)
55+
async def test_deserialization(fixture_name: str, expected_value: DeviceConfig):
56+
"""Test device config deserialization."""
57+
dict_val = json.loads(load_fixture("serialization", fixture_name))
58+
config = DeviceConfig.from_dict(dict_val)
59+
assert config == expected_value
60+
assert expected_value.to_dict() == dict_val
61+
62+
63+
async def test_serialization_http_client():
64+
"""Test that the http client does not try to serialize."""
65+
dict_val = json.loads(load_fixture("serialization", "deviceconfig_plug-klap.json"))
66+
67+
config = replace(PLUG_KLAP_CONFIG, http_client=object())
68+
assert config.http_client
69+
70+
assert config.to_dict() == dict_val
71+
72+
73+
async def test_conn_param_no_https():
74+
"""Test no https in connection param defaults to False."""
75+
dict_val = {
76+
"device_family": "SMART.TAPOPLUG",
77+
"encryption_type": "KLAP",
78+
"login_version": 2,
79+
}
80+
param = DeviceConnectionParameters.from_dict(dict_val)
81+
assert param.https is False
82+
assert param.to_dict() == {**dict_val, "https": False}
2183

2284

2385
@pytest.mark.parametrize(
24-
("input_value", "expected_msg"),
86+
("input_value", "expected_error"),
2587
[
26-
({"Foo": "Bar"}, "Cannot create dataclass from dict, unknown key: Foo"),
27-
("foobar", "Invalid device config data: foobar"),
88+
({"Foo": "Bar"}, MissingField),
89+
("foobar", ValueError),
2890
],
2991
ids=["invalid-dict", "not-dict"],
3092
)
31-
def test_deserialization_errors(input_value, expected_msg):
32-
with pytest.raises(KasaException, match=expected_msg):
93+
def test_deserialization_errors(input_value, expected_error):
94+
with pytest.raises(expected_error):
3395
DeviceConfig.from_dict(input_value)
3496

3597

@@ -39,7 +101,7 @@ async def test_credentials_hash():
39101
http_client=aiohttp.ClientSession(),
40102
credentials=Credentials("foo", "bar"),
41103
)
42-
config_dict = config.to_dict(credentials_hash="credhash")
104+
config_dict = config.to_dict_control_credentials(credentials_hash="credhash")
43105
config_json = json_dumps(config_dict)
44106
config2_dict = json_loads(config_json)
45107
config2 = DeviceConfig.from_dict(config2_dict)
@@ -53,7 +115,7 @@ async def test_blank_credentials_hash():
53115
http_client=aiohttp.ClientSession(),
54116
credentials=Credentials("foo", "bar"),
55117
)
56-
config_dict = config.to_dict(credentials_hash="")
118+
config_dict = config.to_dict_control_credentials(credentials_hash="")
57119
config_json = json_dumps(config_dict)
58120
config2_dict = json_loads(config_json)
59121
config2 = DeviceConfig.from_dict(config2_dict)
@@ -67,7 +129,7 @@ async def test_exclude_credentials():
67129
http_client=aiohttp.ClientSession(),
68130
credentials=Credentials("foo", "bar"),
69131
)
70-
config_dict = config.to_dict(exclude_credentials=True)
132+
config_dict = config.to_dict_control_credentials(exclude_credentials=True)
71133
config_json = json_dumps(config_dict)
72134
config2_dict = json_loads(config_json)
73135
config2 = DeviceConfig.from_dict(config2_dict)

0 commit comments

Comments
 (0)