-
-
Notifications
You must be signed in to change notification settings - Fork 31.3k
/
hass_dict.pyi
176 lines (156 loc) · 6.25 KB
/
hass_dict.pyi
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
"""Stub file for hass_dict. Provide overload for type checking."""
# ruff: noqa: PYI021 # Allow docstrings
from typing import Any, Generic, TypeVar, assert_type, overload
__all__ = [
"HassDict",
"HassEntryKey",
"HassKey",
]
_T = TypeVar("_T")
_U = TypeVar("_U")
class _Key(Generic[_T]):
"""Base class for Hass key types. At runtime delegated to str."""
def __init__(self, value: str, /) -> None: ...
def __len__(self) -> int: ...
def __hash__(self) -> int: ...
def __eq__(self, other: object) -> bool: ...
def __getitem__(self, index: int) -> str: ...
class HassEntryKey(_Key[_T]):
"""Key type for integrations with config entries."""
class HassKey(_Key[_T]):
"""Generic Hass key type."""
class HassDict(dict[_Key[Any] | str, Any]):
"""Custom dict type to provide better value type hints for Hass key types."""
@overload # type: ignore[override]
def __getitem__(self, key: HassEntryKey[_T], /) -> dict[str, _T]: ...
@overload
def __getitem__(self, key: HassKey[_T], /) -> _T: ...
@overload
def __getitem__(self, key: str, /) -> Any: ...
# ------
@overload # type: ignore[override]
def __setitem__(self, key: HassEntryKey[_T], value: dict[str, _T], /) -> None: ...
@overload
def __setitem__(self, key: HassKey[_T], value: _T, /) -> None: ...
@overload
def __setitem__(self, key: str, value: Any, /) -> None: ...
# ------
@overload # type: ignore[override]
def setdefault(
self, key: HassEntryKey[_T], default: dict[str, _T], /
) -> dict[str, _T]: ...
@overload
def setdefault(self, key: HassKey[_T], default: _T, /) -> _T: ...
@overload
def setdefault(self, key: str, default: None = None, /) -> Any | None: ...
@overload
def setdefault(self, key: str, default: Any, /) -> Any: ...
# ------
@overload # type: ignore[override]
def get(self, key: HassEntryKey[_T], /) -> dict[str, _T] | None: ...
@overload
def get(self, key: HassEntryKey[_T], default: _U, /) -> dict[str, _T] | _U: ...
@overload
def get(self, key: HassKey[_T], /) -> _T | None: ...
@overload
def get(self, key: HassKey[_T], default: _U, /) -> _T | _U: ...
@overload
def get(self, key: str, /) -> Any | None: ...
@overload
def get(self, key: str, default: Any, /) -> Any: ...
# ------
@overload # type: ignore[override]
def pop(self, key: HassEntryKey[_T], /) -> dict[str, _T]: ...
@overload
def pop(
self, key: HassEntryKey[_T], default: dict[str, _T], /
) -> dict[str, _T]: ...
@overload
def pop(self, key: HassEntryKey[_T], default: _U, /) -> dict[str, _T] | _U: ...
@overload
def pop(self, key: HassKey[_T], /) -> _T: ...
@overload
def pop(self, key: HassKey[_T], default: _T, /) -> _T: ...
@overload
def pop(self, key: HassKey[_T], default: _U, /) -> _T | _U: ...
@overload
def pop(self, key: str, /) -> Any: ...
@overload
def pop(self, key: str, default: _U, /) -> Any | _U: ...
def _test_hass_dict_typing() -> None: # noqa: PYI048
"""Test HassDict overloads work as intended.
This is tested during the mypy run. Do not move it to 'tests'!
"""
d = HassDict()
entry_key = HassEntryKey[int]("entry_key")
key = HassKey[int]("key")
key2 = HassKey[dict[int, bool]]("key2")
key3 = HassKey[set[str]]("key3")
other_key = "domain"
# __getitem__
assert_type(d[entry_key], dict[str, int])
assert_type(d[entry_key]["entry_id"], int)
assert_type(d[key], int)
assert_type(d[key2], dict[int, bool])
# __setitem__
d[entry_key] = {}
d[entry_key] = 2 # type: ignore[call-overload]
d[entry_key]["entry_id"] = 2
d[entry_key]["entry_id"] = "Hello World" # type: ignore[assignment]
d[key] = 2
d[key] = "Hello World" # type: ignore[misc]
d[key] = {} # type: ignore[misc]
d[key2] = {}
d[key2] = 2 # type: ignore[misc]
d[key3] = set()
d[key3] = 2 # type: ignore[misc]
d[other_key] = 2
d[other_key] = "Hello World"
# get
assert_type(d.get(entry_key), dict[str, int] | None)
assert_type(d.get(entry_key, True), dict[str, int] | bool)
assert_type(d.get(key), int | None)
assert_type(d.get(key, True), int | bool)
assert_type(d.get(key2), dict[int, bool] | None)
assert_type(d.get(key2, {}), dict[int, bool])
assert_type(d.get(key3), set[str] | None)
assert_type(d.get(key3, set()), set[str])
assert_type(d.get(other_key), Any | None)
assert_type(d.get(other_key, True), Any)
assert_type(d.get(other_key, {})["id"], Any)
# setdefault
assert_type(d.setdefault(entry_key, {}), dict[str, int])
assert_type(d.setdefault(entry_key, {})["entry_id"], int)
assert_type(d.setdefault(key, 2), int)
assert_type(d.setdefault(key2, {}), dict[int, bool])
assert_type(d.setdefault(key2, {})[2], bool)
assert_type(d.setdefault(key3, set()), set[str])
assert_type(d.setdefault(other_key, 2), Any)
assert_type(d.setdefault(other_key), Any | None)
d.setdefault(entry_key, {})["entry_id"] = 2
d.setdefault(entry_key, {})["entry_id"] = "Hello World" # type: ignore[assignment]
d.setdefault(key, 2)
d.setdefault(key, "Error") # type: ignore[misc]
d.setdefault(key2, {})[2] = True
d.setdefault(key2, {})[2] = "Error" # type: ignore[assignment]
d.setdefault(key3, set()).add("Hello World")
d.setdefault(key3, set()).add(2) # type: ignore[arg-type]
d.setdefault(other_key, {})["id"] = 2
d.setdefault(other_key, {})["id"] = "Hello World"
d.setdefault(entry_key) # type: ignore[call-overload]
d.setdefault(key) # type: ignore[call-overload]
d.setdefault(key2) # type: ignore[call-overload]
# pop
assert_type(d.pop(entry_key), dict[str, int])
assert_type(d.pop(entry_key, {}), dict[str, int])
assert_type(d.pop(entry_key, 2), dict[str, int] | int)
assert_type(d.pop(key), int)
assert_type(d.pop(key, 2), int)
assert_type(d.pop(key, "Hello World"), int | str)
assert_type(d.pop(key2), dict[int, bool])
assert_type(d.pop(key2, {}), dict[int, bool])
assert_type(d.pop(key2, 2), dict[int, bool] | int)
assert_type(d.pop(key3), set[str])
assert_type(d.pop(key3, set()), set[str])
assert_type(d.pop(other_key), Any)
assert_type(d.pop(other_key, True), Any | bool)