File: usertype.py

package info (click to toggle)
python-aioinflux 0.9.0-5
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, sid
  • size: 320 kB
  • sloc: python: 1,402; makefile: 41
file content (206 lines) | stat: -rw-r--r-- 7,610 bytes parent folder | download | duplicates (2)
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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
import enum
import ciso8601
import time
import decimal
from collections import Counter
from typing import TypeVar, Optional, Mapping, Union
from datetime import datetime

# noinspection PyUnresolvedReferences
from .common import *  # noqa
from ..compat import pd

__all__ = [
    'lineprotocol', 'SchemaError',
    'MEASUREMENT', 'TIMEINT', 'TIMESTR', 'TIMEDT',
    'TAG', 'TAGENUM',
    'BOOL', 'INT', 'DECIMAL', 'FLOAT', 'STR', 'ENUM',
]

MEASUREMENT = TypeVar('MEASUREMENT', bound=str)
TIMEINT = TypeVar('TIMEINT', bound=int)
TIMESTR = TypeVar('TIMESTR', bound=str)
TIMEDT = TypeVar('TIMEDT', bound=datetime)
TAG = TypeVar('TAG', bound=str)
TAGENUM = TypeVar('TAGENUM', bound=enum.Enum)
BOOL = TypeVar('BOOL', bound=bool)
INT = TypeVar('INT', bound=int)
DECIMAL = TypeVar('DECIMAL', bound=decimal.Decimal)
FLOAT = TypeVar('FLOAT', bound=float)
STR = TypeVar('STR', bound=str)
ENUM = TypeVar('ENUM', bound=enum.Enum)

time_types = [TIMEINT, TIMEDT, TIMESTR]
tag_types = [TAG, TAGENUM]
field_types = [BOOL, INT, DECIMAL, FLOAT, STR, ENUM]
optional_field_types = [Optional[f] for f in field_types]


class SchemaError(TypeError):
    """Raised when invalid schema is passed to :func:`lineprotocol`"""


def str_to_dt(s):
    dt = ciso8601.parse_datetime(s)
    if dt:
        return dt
    raise ValueError(f'Invalid datetime string: {dt!r}')


def dt_to_int(dt):
    if not dt.tzinfo:
        # Assume tz-naive input to be in UTC, not local time
        return int(dt.timestamp() - time.timezone) * 10 ** 9 + dt.microsecond * 1000
    return int(dt.timestamp()) * 10 ** 9 + dt.microsecond * 1000


def _validate_schema(schema, placeholder):
    c = Counter(schema.values())
    if not c:
        raise SchemaError("Schema/type annotations missing")
    if c[MEASUREMENT] > 1:
        raise SchemaError("Class can't have more than one 'MEASUREMENT' attribute")
    if sum(c[e] for e in time_types) > 1:
        raise SchemaError(f"Can't have more than one timestamp-type attribute {time_types}")
    if sum(c[e] for e in field_types + optional_field_types) < 1 and not placeholder:
        raise SchemaError(f"Must have one or more non-empty "
                          f"field-type attributes {field_types}")


def is_optional(t, base_type):
    """Checks if type hint is Optional[base_type]"""
    # NOTE: The 'typing' module is still "provisional" and documentation sub-optimal,
    #  which requires these kinds instrospection into undocumented implementation details
    # NOTE: May break in Python 3.8
    # TODO: Check if works on Python 3.6
    try:
        cond1 = getattr(t, '__origin__') is Union
        cond2 = {type(None), base_type} == set(getattr(t, '__args__', []))
        if cond1 and cond2:
            return True
    except AttributeError:
        return False
    return False


def _make_serializer(meas, schema, extra_tags, placeholder):  # noqa: C901
    """Factory of line protocol parsers"""
    _validate_schema(schema, placeholder)
    tags = []
    fields = []
    ts = None
    meas = meas
    for k, t in schema.items():
        if t is MEASUREMENT:
            meas = f"{{i.{k}}}"
        elif t is TIMEINT:
            ts = f"{{i.{k}}}"
        elif t is TIMESTR:
            if pd:
                ts = f"{{pd.Timestamp(i.{k} or 0).value}}"
            else:
                ts = f"{{dt_to_int(str_to_dt(i.{k}))}}"
        elif t is TIMEDT:
            if pd:
                ts = f"{{pd.Timestamp(i.{k} or 0).value}}"
            else:
                ts = f"{{dt_to_int(i.{k})}}"
        elif t is TAG or is_optional(t, TAG):
            tags.append(f"{k}={{str(i.{k}).translate(tag_escape)}}")
        elif t is TAGENUM or is_optional(t, TAGENUM):
            tags.append(f"{k}={{getattr(i.{k}, 'name', i.{k} or None)}}")
        elif t is FLOAT or is_optional(t, FLOAT):
            fields.append(f"{k}={{i.{k}}}")
        elif t is DECIMAL or is_optional(t, DECIMAL):
            fields.append(f"{k}={{i.{k}}}")
        elif t is BOOL or is_optional(t, BOOL):
            fields.append(f"{k}={{i.{k}}}")
        elif t is INT or is_optional(t, INT):
            fields.append(f"{k}={{i.{k}}}i")
        elif t is STR or is_optional(t, STR):
            fields.append(f"{k}=\\\"{{str(i.{k}).translate(str_escape)}}\\\"")
        elif t is ENUM or is_optional(t, ENUM):
            fields.append(f"{k}=\\\"{{getattr(i.{k}, 'name', i.{k} or None)}}\\\"")
        else:
            raise SchemaError(f"Invalid attribute type {k!r}: {t!r}")
    extra_tags = extra_tags or {}
    for k, v in extra_tags.items():
        tags.append(f"{k}={v}")
    if placeholder:
        fields.insert(0, f"_=true")

    sep = ',' if tags else ''
    ts = f' {ts}' if ts else ''
    fmt = f"{meas}{sep}{','.join(tags)} {','.join(fields)}{ts}"
    f = eval(f'lambda i: f"{fmt}".encode()')
    f.__doc__ = "Returns InfluxDB line protocol representation of user-defined class"
    return f


def lineprotocol(
        cls=None,
        *,
        schema: Optional[Mapping[str, type]] = None,
        rm_none: bool = False,
        extra_tags: Optional[Mapping[str, str]] = None,
        placeholder: bool = False
):
    """Adds ``to_lineprotocol`` method to arbitrary user-defined classes

    :param cls: Class to monkey-patch
    :param schema: Schema dictionary (attr/type pairs).
    :param rm_none: Whether apply a regex to remove ``None`` values.
        If ``False``, passing ``None`` values to boolean, integer or float or time fields
        will result in write errors. Setting to ``True`` is "safer" but impacts performance.
    :param extra_tags: Hard coded tags to be added to every point generated.
    :param placeholder: If no field attributes are present, add a placeholder attribute (``_``)
        which is always equal to ``True``. This is a workaround for creating field-less points
        (which is not supported natively by InfluxDB)
    """
    opts = dict(
        schema=schema,
        rm_none=rm_none,
        extra_tags=extra_tags or {},
        placeholder=placeholder,
    )

    def _lineprotocol(cls):
        _schema = schema or getattr(cls, '__annotations__', {})
        # TODO: Raise warning or exception if schema has optionals but rm_none is False
        # for t in _schema.values():
        #     for bt in field_types + tag_types:
        #         if is_optional(t, bt):
        #             warnings.warn("")
        f = _make_serializer(cls.__name__, _schema, extra_tags, placeholder)
        cls.to_lineprotocol = f
        cls.to_lineprotocol.opts = opts
        return cls

    def _rm_none_lineprotocol(cls):

        def _parser_selector(i):
            if not hasattr(i, '_asdict'):
                raise ValueError("'rm_none' can only be used with namedtuples")
            key = tuple([k for k, v in i._asdict().items() if v != '' and v is not None])
            if key not in parsers:
                _schema = schema or getattr(cls, '__annotations__', {})
                _schema = {k: v for k, v in _schema.items() if k in key}
                parsers[key] = _make_serializer(cls.__name__, _schema, extra_tags, placeholder)
            return parsers[key](i)

        parsers = {}
        cls.to_lineprotocol = _parser_selector
        cls.to_lineprotocol.opts = opts
        return cls

    if cls:
        if rm_none:
            # Using rm_none has substantial runtime impact.
            # Best avoided if performance is critical.
            return _rm_none_lineprotocol(cls)
        # No options
        return _lineprotocol(cls)
    else:
        if rm_none:
            return _rm_none_lineprotocol
        return _lineprotocol