utils.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310
  1. import datetime
  2. import logging
  3. import textwrap
  4. from contextlib import contextmanager
  5. from functools import wraps
  6. from typing import Any, Dict, List, Mapping, Optional, Union
  7. from redis.exceptions import DataError
  8. from redis.typing import AbsExpiryT, EncodableT, ExpiryT
  9. try:
  10. import hiredis # noqa
  11. # Only support Hiredis >= 3.0:
  12. hiredis_version = hiredis.__version__.split(".")
  13. HIREDIS_AVAILABLE = int(hiredis_version[0]) > 3 or (
  14. int(hiredis_version[0]) == 3 and int(hiredis_version[1]) >= 2
  15. )
  16. if not HIREDIS_AVAILABLE:
  17. raise ImportError("hiredis package should be >= 3.2.0")
  18. except ImportError:
  19. HIREDIS_AVAILABLE = False
  20. try:
  21. import ssl # noqa
  22. SSL_AVAILABLE = True
  23. except ImportError:
  24. SSL_AVAILABLE = False
  25. try:
  26. import cryptography # noqa
  27. CRYPTOGRAPHY_AVAILABLE = True
  28. except ImportError:
  29. CRYPTOGRAPHY_AVAILABLE = False
  30. from importlib import metadata
  31. def from_url(url, **kwargs):
  32. """
  33. Returns an active Redis client generated from the given database URL.
  34. Will attempt to extract the database id from the path url fragment, if
  35. none is provided.
  36. """
  37. from redis.client import Redis
  38. return Redis.from_url(url, **kwargs)
  39. @contextmanager
  40. def pipeline(redis_obj):
  41. p = redis_obj.pipeline()
  42. yield p
  43. p.execute()
  44. def str_if_bytes(value: Union[str, bytes]) -> str:
  45. return (
  46. value.decode("utf-8", errors="replace") if isinstance(value, bytes) else value
  47. )
  48. def safe_str(value):
  49. return str(str_if_bytes(value))
  50. def dict_merge(*dicts: Mapping[str, Any]) -> Dict[str, Any]:
  51. """
  52. Merge all provided dicts into 1 dict.
  53. *dicts : `dict`
  54. dictionaries to merge
  55. """
  56. merged = {}
  57. for d in dicts:
  58. merged.update(d)
  59. return merged
  60. def list_keys_to_dict(key_list, callback):
  61. return dict.fromkeys(key_list, callback)
  62. def merge_result(command, res):
  63. """
  64. Merge all items in `res` into a list.
  65. This command is used when sending a command to multiple nodes
  66. and the result from each node should be merged into a single list.
  67. res : 'dict'
  68. """
  69. result = set()
  70. for v in res.values():
  71. for value in v:
  72. result.add(value)
  73. return list(result)
  74. def warn_deprecated(name, reason="", version="", stacklevel=2):
  75. import warnings
  76. msg = f"Call to deprecated {name}."
  77. if reason:
  78. msg += f" ({reason})"
  79. if version:
  80. msg += f" -- Deprecated since version {version}."
  81. warnings.warn(msg, category=DeprecationWarning, stacklevel=stacklevel)
  82. def deprecated_function(reason="", version="", name=None):
  83. """
  84. Decorator to mark a function as deprecated.
  85. """
  86. def decorator(func):
  87. @wraps(func)
  88. def wrapper(*args, **kwargs):
  89. warn_deprecated(name or func.__name__, reason, version, stacklevel=3)
  90. return func(*args, **kwargs)
  91. return wrapper
  92. return decorator
  93. def warn_deprecated_arg_usage(
  94. arg_name: Union[list, str],
  95. function_name: str,
  96. reason: str = "",
  97. version: str = "",
  98. stacklevel: int = 2,
  99. ):
  100. import warnings
  101. msg = (
  102. f"Call to '{function_name}' function with deprecated"
  103. f" usage of input argument/s '{arg_name}'."
  104. )
  105. if reason:
  106. msg += f" ({reason})"
  107. if version:
  108. msg += f" -- Deprecated since version {version}."
  109. warnings.warn(msg, category=DeprecationWarning, stacklevel=stacklevel)
  110. def deprecated_args(
  111. args_to_warn: list = ["*"],
  112. allowed_args: list = [],
  113. reason: str = "",
  114. version: str = "",
  115. ):
  116. """
  117. Decorator to mark specified args of a function as deprecated.
  118. If '*' is in args_to_warn, all arguments will be marked as deprecated.
  119. """
  120. def decorator(func):
  121. @wraps(func)
  122. def wrapper(*args, **kwargs):
  123. # Get function argument names
  124. arg_names = func.__code__.co_varnames[: func.__code__.co_argcount]
  125. provided_args = dict(zip(arg_names, args))
  126. provided_args.update(kwargs)
  127. provided_args.pop("self", None)
  128. for allowed_arg in allowed_args:
  129. provided_args.pop(allowed_arg, None)
  130. for arg in args_to_warn:
  131. if arg == "*" and len(provided_args) > 0:
  132. warn_deprecated_arg_usage(
  133. list(provided_args.keys()),
  134. func.__name__,
  135. reason,
  136. version,
  137. stacklevel=3,
  138. )
  139. elif arg in provided_args:
  140. warn_deprecated_arg_usage(
  141. arg, func.__name__, reason, version, stacklevel=3
  142. )
  143. return func(*args, **kwargs)
  144. return wrapper
  145. return decorator
  146. def _set_info_logger():
  147. """
  148. Set up a logger that log info logs to stdout.
  149. (This is used by the default push response handler)
  150. """
  151. if "push_response" not in logging.root.manager.loggerDict.keys():
  152. logger = logging.getLogger("push_response")
  153. logger.setLevel(logging.INFO)
  154. handler = logging.StreamHandler()
  155. handler.setLevel(logging.INFO)
  156. logger.addHandler(handler)
  157. def get_lib_version():
  158. try:
  159. libver = metadata.version("redis")
  160. except metadata.PackageNotFoundError:
  161. libver = "99.99.99"
  162. return libver
  163. def format_error_message(host_error: str, exception: BaseException) -> str:
  164. if not exception.args:
  165. return f"Error connecting to {host_error}."
  166. elif len(exception.args) == 1:
  167. return f"Error {exception.args[0]} connecting to {host_error}."
  168. else:
  169. return (
  170. f"Error {exception.args[0]} connecting to {host_error}. "
  171. f"{exception.args[1]}."
  172. )
  173. def compare_versions(version1: str, version2: str) -> int:
  174. """
  175. Compare two versions.
  176. :return: -1 if version1 > version2
  177. 0 if both versions are equal
  178. 1 if version1 < version2
  179. """
  180. num_versions1 = list(map(int, version1.split(".")))
  181. num_versions2 = list(map(int, version2.split(".")))
  182. if len(num_versions1) > len(num_versions2):
  183. diff = len(num_versions1) - len(num_versions2)
  184. for _ in range(diff):
  185. num_versions2.append(0)
  186. elif len(num_versions1) < len(num_versions2):
  187. diff = len(num_versions2) - len(num_versions1)
  188. for _ in range(diff):
  189. num_versions1.append(0)
  190. for i, ver in enumerate(num_versions1):
  191. if num_versions1[i] > num_versions2[i]:
  192. return -1
  193. elif num_versions1[i] < num_versions2[i]:
  194. return 1
  195. return 0
  196. def ensure_string(key):
  197. if isinstance(key, bytes):
  198. return key.decode("utf-8")
  199. elif isinstance(key, str):
  200. return key
  201. else:
  202. raise TypeError("Key must be either a string or bytes")
  203. def extract_expire_flags(
  204. ex: Optional[ExpiryT] = None,
  205. px: Optional[ExpiryT] = None,
  206. exat: Optional[AbsExpiryT] = None,
  207. pxat: Optional[AbsExpiryT] = None,
  208. ) -> List[EncodableT]:
  209. exp_options: list[EncodableT] = []
  210. if ex is not None:
  211. exp_options.append("EX")
  212. if isinstance(ex, datetime.timedelta):
  213. exp_options.append(int(ex.total_seconds()))
  214. elif isinstance(ex, int):
  215. exp_options.append(ex)
  216. elif isinstance(ex, str) and ex.isdigit():
  217. exp_options.append(int(ex))
  218. else:
  219. raise DataError("ex must be datetime.timedelta or int")
  220. elif px is not None:
  221. exp_options.append("PX")
  222. if isinstance(px, datetime.timedelta):
  223. exp_options.append(int(px.total_seconds() * 1000))
  224. elif isinstance(px, int):
  225. exp_options.append(px)
  226. else:
  227. raise DataError("px must be datetime.timedelta or int")
  228. elif exat is not None:
  229. if isinstance(exat, datetime.datetime):
  230. exat = int(exat.timestamp())
  231. exp_options.extend(["EXAT", exat])
  232. elif pxat is not None:
  233. if isinstance(pxat, datetime.datetime):
  234. pxat = int(pxat.timestamp() * 1000)
  235. exp_options.extend(["PXAT", pxat])
  236. return exp_options
  237. def truncate_text(txt, max_length=100):
  238. return textwrap.shorten(
  239. text=txt, width=max_length, placeholder="...", break_long_words=True
  240. )