File: helpers.py

package info (click to toggle)
python-mongomock 4.3.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 1,028 kB
  • sloc: python: 16,412; makefile: 24
file content (426 lines) | stat: -rw-r--r-- 13,872 bytes parent folder | download
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
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
from collections import abc
from collections import OrderedDict
from datetime import datetime, timedelta, tzinfo
from mongomock import InvalidURI
from packaging import version
import re
import time
from urllib.parse import unquote_plus
import warnings


# Get ObjectId from bson if available or import a crafted one. This is not used
# in this module but is made available for callers of this module.
try:
    from bson import ObjectId  # pylint: disable=unused-import
    from bson import Timestamp
    from pymongo import version as pymongo_version
    PYMONGO_VERSION = version.parse(pymongo_version)
    HAVE_PYMONGO = True
except ImportError:
    from mongomock.object_id import ObjectId  # noqa
    Timestamp = None
    # Default Pymongo version if not present.
    PYMONGO_VERSION = version.parse('4.0')
    HAVE_PYMONGO = False

# Cache the RegExp pattern type.
RE_TYPE = type(re.compile(''))
_HOST_MATCH = re.compile(r'^([^@]+@)?([^:]+|\[[^\]]+\])(:([^:]+))?$')
_SIMPLE_HOST_MATCH = re.compile(r'^([^:]+|\[[^\]]+\])(:([^:]+))?$')

try:
    from bson.tz_util import utc
except ImportError:
    class _FixedOffset(tzinfo):

        def __init__(self, offset, name):
            self.__offset = timedelta(minutes=offset)
            self.__name = name

        def __getinitargs__(self):
            return self.__offset, self.__name

        def utcoffset(self, dt):
            return self.__offset

        def tzname(self, dt):
            return self.__name

        def dst(self, dt):
            return timedelta(0)
    utc = _FixedOffset(0, 'UTC')


ASCENDING = 1
DESCENDING = -1


def utcnow():
    """Simple wrapper for datetime.utcnow

    This provides a centralized definition of "now" in the mongomock realm,
    allowing users to transform the value of "now" to the future or the past,
    based on their testing needs. For example:

    ```python
    def test_x(self):
        with mock.patch("mongomock.utcnow") as mm_utc:
            mm_utc = datetime.utcnow() + timedelta(hours=100)
            # Test some things "100 hours" in the future
    ```
    """
    return datetime.utcnow()


def print_deprecation_warning(old_param_name, new_param_name):
    warnings.warn(
        "'%s' has been deprecated to be in line with pymongo implementation, a new parameter '%s' "
        'should be used instead. the old parameter will be kept for backward compatibility '
        'purposes.' % (old_param_name, new_param_name), DeprecationWarning)


def create_index_list(key_or_list, direction=None):
    """Helper to generate a list of (key, direction) pairs.

       It takes such a list, or a single key, or a single key and direction.
    """
    if isinstance(key_or_list, str):
        return [(key_or_list, direction or ASCENDING)]
    if not isinstance(key_or_list, (list, tuple, abc.Iterable)):
        raise TypeError('if no direction is specified, '
                        'key_or_list must be an instance of list')
    return key_or_list


def gen_index_name(index_list):
    """Generate an index name based on the list of keys with directions."""

    return '_'.join(['%s_%s' % item for item in index_list])


class hashdict(dict):
    """hashable dict implementation, suitable for use as a key into other dicts.

    >>> h1 = hashdict({'apples': 1, 'bananas':2})
    >>> h2 = hashdict({'bananas': 3, 'mangoes': 5})
    >>> h1+h2
    hashdict(apples=1, bananas=3, mangoes=5)
    >>> d1 = {}
    >>> d1[h1] = 'salad'
    >>> d1[h1]
    'salad'
    >>> d1[h2]
    Traceback (most recent call last):
    ...
    KeyError: hashdict(bananas=3, mangoes=5)

    based on answers from
    http://stackoverflow.com/questions/1151658/python-hashable-dicts
    """
    def __key(self):
        return frozenset((k,
                          hashdict(v) if isinstance(v, dict) else
                          tuple(v) if isinstance(v, list) else
                          v)
                         for k, v in self.items())

    def __repr__(self):
        return '{}({})'.format(
            self.__class__.__name__,
            ', '.join(f'{str(i[0])}={repr(i[1])}' for i in sorted(self.__key())))

    def __hash__(self):
        return hash(self.__key())

    def __setitem__(self, key, value):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def __delitem__(self, key):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def clear(self):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def pop(self, *args, **kwargs):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def popitem(self, *args, **kwargs):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def setdefault(self, *args, **kwargs):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def update(self, *args, **kwargs):
        raise TypeError('{} does not support item assignment'
                        .format(self.__class__.__name__))

    def __add__(self, right):
        result = hashdict(self)
        dict.update(result, right)
        return result


def fields_list_to_dict(fields):
    """Takes a list of field names and returns a matching dictionary.

    ['a', 'b'] becomes {'a': 1, 'b': 1}

    and

    ['a.b.c', 'd', 'a.c'] becomes {'a.b.c': 1, 'd': 1, 'a.c': 1}
    """
    as_dict = {}
    for field in fields:
        if not isinstance(field, str):
            raise TypeError('fields must be a list of key names, each an instance of str')
        as_dict[field] = 1
    return as_dict


def parse_uri(uri, default_port=27017, warn=False):
    """A simplified version of pymongo.uri_parser.parse_uri.

    Returns a dict with:
     - nodelist, a tuple of (host, port)
     - database the name of the database or None if no database is provided in the URI.

    An invalid MongoDB connection URI may raise an InvalidURI exception,
    however, the URI is not fully parsed and some invalid URIs may not result
    in an exception.

    'mongodb://host1/database' becomes 'host1', 27017, 'database'

    and

    'mongodb://host1' becomes 'host1', 27017, None
    """
    SCHEME = 'mongodb://'

    if not uri.startswith(SCHEME):
        raise InvalidURI('Invalid URI scheme: URI '
                         "must begin with '%s'" % (SCHEME,))

    scheme_free = uri[len(SCHEME):]

    if not scheme_free:
        raise InvalidURI('Must provide at least one hostname or IP.')

    dbase = None

    # Check for unix domain sockets in the uri
    if '.sock' in scheme_free:
        host_part, _, path_part = scheme_free.rpartition('/')
        if not host_part:
            host_part = path_part
            path_part = ''
        if '/' in host_part:
            raise InvalidURI("Any '/' in a unix domain socket must be"
                             ' URL encoded: %s' % host_part)
        path_part = unquote_plus(path_part)
    else:
        host_part, _, path_part = scheme_free.partition('/')

    if not path_part and '?' in host_part:
        raise InvalidURI("A '/' is required between "
                         'the host list and any options.')

    nodelist = []
    if ',' in host_part:
        hosts = host_part.split(',')
    else:
        hosts = [host_part]
    for host in hosts:
        match = _HOST_MATCH.match(host)
        if not match:
            raise ValueError(
                "Reserved characters such as ':' must be escaped according RFC "
                "2396. An IPv6 address literal must be enclosed in '[' and ']' "
                'according to RFC 2732.')
        host = match.group(2)
        if host.startswith('[') and host.endswith(']'):
            host = host[1:-1]

        port = match.group(4)
        if port:
            try:
                port = int(port)
                if port < 0 or port > 65535:
                    raise ValueError()
            except ValueError as err:
                raise ValueError('Port must be an integer between 0 and 65535:', port) from err
        else:
            port = default_port

        nodelist.append((host, port))

    if path_part and path_part[0] != '?':
        dbase, _, _ = path_part.partition('?')
        if '.' in dbase:
            dbase, _ = dbase.split('.', 1)

    if dbase is not None:
        dbase = unquote_plus(dbase)

    return {'nodelist': tuple(nodelist), 'database': dbase}


def split_hosts(hosts, default_port=27017):
    """Split the entity into a list of tuples of host and port."""

    nodelist = []
    for entity in hosts.split(','):
        port = default_port
        if entity.endswith('.sock'):
            port = None

        match = _SIMPLE_HOST_MATCH.match(entity)
        if not match:
            raise ValueError(
                "Reserved characters such as ':' must be escaped according RFC "
                "2396. An IPv6 address literal must be enclosed in '[' and ']' "
                'according to RFC 2732.')
        host = match.group(1)
        if host.startswith('[') and host.endswith(']'):
            host = host[1:-1]

        if match.group(3):
            try:
                port = int(match.group(3))
                if port < 0 or port > 65535:
                    raise ValueError()
            except ValueError as err:
                raise ValueError('Port must be an integer between 0 and 65535:', port) from err

        nodelist.append((host, port))

    return nodelist


_LAST_TIMESTAMP_INC = []


def get_current_timestamp():
    """Get the current timestamp as a bson Timestamp object."""
    if not Timestamp:
        raise NotImplementedError('timestamp is not supported. Import pymongo to use it.')
    now = int(time.time())
    if _LAST_TIMESTAMP_INC and _LAST_TIMESTAMP_INC[0] == now:
        _LAST_TIMESTAMP_INC[1] += 1
    else:
        del _LAST_TIMESTAMP_INC[:]
        _LAST_TIMESTAMP_INC.extend([now, 1])
    return Timestamp(now, _LAST_TIMESTAMP_INC[1])


def patch_datetime_awareness_in_document(value):
    # MongoDB is supposed to stock everything as timezone naive utc date
    # Hence we have to convert incoming datetimes to avoid errors while
    # mixing tz aware and naive.
    # On top of that, MongoDB date precision is up to millisecond, where Python
    # datetime use microsecond, so we must lower the precision to mimic mongo.
    for best_type in (OrderedDict, dict):
        if isinstance(value, best_type):
            return best_type((k, patch_datetime_awareness_in_document(v)) for k, v in value.items())
    if isinstance(value, (tuple, list)):
        return [patch_datetime_awareness_in_document(item) for item in value]
    if isinstance(value, datetime):
        mongo_us = (value.microsecond // 1000) * 1000
        if value.tzinfo:
            return (value - value.utcoffset()).replace(tzinfo=None, microsecond=mongo_us)
        return value.replace(microsecond=mongo_us)
    if Timestamp and isinstance(value, Timestamp) and not value.time and not value.inc:
        return get_current_timestamp()
    return value


def make_datetime_timezone_aware_in_document(value):
    # MongoClient support tz_aware=True parameter to return timezone-aware
    # datetime objects. Given the date is stored internally without timezone
    # information, all returned datetime have utc as timezone.
    if isinstance(value, dict):
        return {k: make_datetime_timezone_aware_in_document(v) for k, v in value.items()}
    if isinstance(value, (tuple, list)):
        return [make_datetime_timezone_aware_in_document(item) for item in value]
    if isinstance(value, datetime):
        return value.replace(tzinfo=utc)
    return value


def get_value_by_dot(doc, key, can_generate_array=False):
    """Get dictionary value using dotted key"""
    result = doc
    key_items = key.split('.')
    for key_index, key_item in enumerate(key_items):
        if isinstance(result, dict):
            result = result[key_item]

        elif isinstance(result, (list, tuple)):
            try:
                int_key = int(key_item)
            except ValueError as err:
                if not can_generate_array:
                    raise KeyError(key_index) from err
                remaining_key = '.'.join(key_items[key_index:])
                return [get_value_by_dot(subdoc, remaining_key) for subdoc in result]

            try:
                result = result[int_key]
            except (ValueError, IndexError) as err:
                raise KeyError(key_index) from err

        else:
            raise KeyError(key_index)

    return result


def set_value_by_dot(doc, key, value):
    """Set dictionary value using dotted key"""
    try:
        parent_key, child_key = key.rsplit('.', 1)
        parent = get_value_by_dot(doc, parent_key)
    except ValueError:
        child_key = key
        parent = doc

    if isinstance(parent, dict):
        parent[child_key] = value
    elif isinstance(parent, (list, tuple)):
        try:
            parent[int(child_key)] = value
        except (ValueError, IndexError) as err:
            raise KeyError() from err
    else:
        raise KeyError()

    return doc


def delete_value_by_dot(doc, key):
    """Delete dictionary value using dotted key.

    This function assumes that the value exists.
    """
    try:
        parent_key, child_key = key.rsplit('.', 1)
        parent = get_value_by_dot(doc, parent_key)
    except ValueError:
        child_key = key
        parent = doc

    del parent[child_key]

    return doc


def mongodb_to_bool(value):
    """Converts any value to bool the way MongoDB does it"""

    return value not in [False, None, 0]