-
Notifications
You must be signed in to change notification settings - Fork 41
/
records.py
112 lines (89 loc) · 3.36 KB
/
records.py
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
from collections.abc import Mapping
from typing import Any, Callable, Dict, Iterator, List, Tuple, Union
# Optional cython extension:
try:
from aiochclient._types import empty_convertor, what_py_converter
except ImportError:
from aiochclient.types import empty_convertor, what_py_converter
__all__ = ["RecordsFabric", "Record", "FromJsonFabric"]
class Record(Mapping):
"""Lightweight, memory efficient objects with full mapping interface, where
you can get fields by names or by indexes.
Usage:
.. code-block:: python
row = await client.fetchrow("SELECT a, b FROM t WHERE a=1")
assert row["a"] == 1
assert row[0] == 1
assert row[:] == (1, (dt.date(2018, 9, 8), 3.14))
assert list(row.keys()) == ["a", "b"]
assert list(row.values()) == [1, (dt.date(2018, 9, 8), 3.14)]
"""
__slots__ = ("_converters", "_decoded", "_names", "_row")
def __init__(self, row: bytes, names: Dict[str, Any], converters: List[Callable]):
self._row: Union[bytes, Tuple[Any]] = row
if not self._row:
# in case of empty row
self._decoded = True
self._converters = []
self._names = {}
else:
self._decoded = False
self._converters = converters
self._names = names
def __getitem__(self, key: Union[str, int, slice]) -> Any:
self._decode()
return self._getitem(key)
def _getitem(self, key: Union[str, int, slice]) -> Any:
if type(key) == str:
try:
return self._row[self._names[key]]
except KeyError:
if not self._row:
raise KeyError(
"Empty row. May be it is result of 'WITH TOTALS' query."
)
raise KeyError(f"No fields with name '{key}'")
try:
return self._row[key]
except IndexError:
if not self._row:
raise IndexError(
"Empty row. May be it is result of 'WITH TOTALS' query."
)
raise IndexError(f"No fields with index '{key}'")
def __iter__(self) -> Iterator:
return iter(self._names)
def __len__(self) -> int:
return len(self._names)
def _decode(self):
if self._decoded:
return None
self._row = tuple(
converter(val)
for converter, val in zip(self._converters, self._row.split(b"\t"))
)
self._decoded = True
class RecordsFabric:
__slots__ = ("converters", "names")
def __init__(self, tps: bytes, names: bytes, convert: bool = True):
names = names.decode().strip().split("\t")
self.names = {key: index for (index, key) in enumerate(names)}
if convert:
self.converters = [
what_py_converter(tp) for tp in tps.decode().strip().split("\t")
]
else:
self.converters = [
empty_convertor for _ in tps.decode().strip().split("\t")
]
def new(self, row: bytes) -> Record:
return Record(
row=row[:-1], # because of delimiter
names=self.names,
converters=self.converters,
)
class FromJsonFabric:
def __init__(self, loads):
self.loads = loads
def new(self, row: bytes) -> Any:
return self.loads(row)