frappe_docker/frappe-bench/env/lib/python2.7/site-packages/rq/compat/connections.py
2017-07-31 15:51:51 +05:30

57 lines
2 KiB
Python

# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from functools import partial
from redis import Redis, StrictRedis
def fix_return_type(func):
# deliberately no functools.wraps() call here, since the function being
# wrapped is a partial, which has no module
def _inner(*args, **kwargs):
value = func(*args, **kwargs)
if value is None:
value = -1
return value
return _inner
PATCHED_METHODS = ['_setex', '_lrem', '_zadd', '_pipeline', '_ttl']
def _hset(self, key, field_name, value, pipeline=None):
connection = pipeline if pipeline is not None else self
connection.hset(key, field_name, value)
def patch_connection(connection):
# Don't patch already patches objects
if all([hasattr(connection, attr) for attr in PATCHED_METHODS]):
return connection
connection._hset = partial(_hset, connection)
if isinstance(connection, Redis):
connection._setex = partial(StrictRedis.setex, connection)
connection._lrem = partial(StrictRedis.lrem, connection)
connection._zadd = partial(StrictRedis.zadd, connection)
connection._pipeline = partial(StrictRedis.pipeline, connection)
connection._ttl = fix_return_type(partial(StrictRedis.ttl, connection))
if hasattr(connection, 'pttl'):
connection._pttl = fix_return_type(partial(StrictRedis.pttl, connection))
# add support for mock redis objects
elif hasattr(connection, 'setex'):
connection._setex = connection.setex
connection._lrem = connection.lrem
connection._zadd = connection.zadd
connection._pipeline = connection.pipeline
connection._ttl = connection.ttl
if hasattr(connection, 'pttl'):
connection._pttl = connection.pttl
else:
raise ValueError('Unanticipated connection type: {}. Please report this.'.format(type(connection)))
return connection