helpers.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  1. import copy
  2. import random
  3. import string
  4. from typing import List, Tuple
  5. import redis
  6. from redis.typing import KeysT, KeyT
  7. def list_or_args(keys: KeysT, args: Tuple[KeyT, ...]) -> List[KeyT]:
  8. # returns a single new list combining keys and args
  9. try:
  10. iter(keys)
  11. # a string or bytes instance can be iterated, but indicates
  12. # keys wasn't passed as a list
  13. if isinstance(keys, (bytes, str)):
  14. keys = [keys]
  15. else:
  16. keys = list(keys)
  17. except TypeError:
  18. keys = [keys]
  19. if args:
  20. keys.extend(args)
  21. return keys
  22. def nativestr(x):
  23. """Return the decoded binary string, or a string, depending on type."""
  24. r = x.decode("utf-8", "replace") if isinstance(x, bytes) else x
  25. if r == "null":
  26. return
  27. return r
  28. def delist(x):
  29. """Given a list of binaries, return the stringified version."""
  30. if x is None:
  31. return x
  32. return [nativestr(obj) for obj in x]
  33. def parse_to_list(response):
  34. """Optimistically parse the response to a list."""
  35. res = []
  36. if response is None:
  37. return res
  38. for item in response:
  39. try:
  40. res.append(int(item))
  41. except ValueError:
  42. try:
  43. res.append(float(item))
  44. except ValueError:
  45. res.append(nativestr(item))
  46. except TypeError:
  47. res.append(None)
  48. return res
  49. def parse_list_to_dict(response):
  50. res = {}
  51. for i in range(0, len(response), 2):
  52. if isinstance(response[i], list):
  53. res["Child iterators"].append(parse_list_to_dict(response[i]))
  54. elif isinstance(response[i + 1], list):
  55. res["Child iterators"] = [parse_list_to_dict(response[i + 1])]
  56. else:
  57. try:
  58. res[response[i]] = float(response[i + 1])
  59. except (TypeError, ValueError):
  60. res[response[i]] = response[i + 1]
  61. return res
  62. def parse_to_dict(response):
  63. if response is None:
  64. return {}
  65. res = {}
  66. for det in response:
  67. if isinstance(det[1], list):
  68. res[det[0]] = parse_list_to_dict(det[1])
  69. else:
  70. try: # try to set the attribute. may be provided without value
  71. try: # try to convert the value to float
  72. res[det[0]] = float(det[1])
  73. except (TypeError, ValueError):
  74. res[det[0]] = det[1]
  75. except IndexError:
  76. pass
  77. return res
  78. def random_string(length=10):
  79. """
  80. Returns a random N character long string.
  81. """
  82. return "".join( # nosec
  83. random.choice(string.ascii_lowercase) for x in range(length)
  84. )
  85. def quote_string(v):
  86. """
  87. RedisGraph strings must be quoted,
  88. quote_string wraps given v with quotes incase
  89. v is a string.
  90. """
  91. if isinstance(v, bytes):
  92. v = v.decode()
  93. elif not isinstance(v, str):
  94. return v
  95. if len(v) == 0:
  96. return '""'
  97. v = v.replace("\\", "\\\\")
  98. v = v.replace('"', '\\"')
  99. return f'"{v}"'
  100. def decode_dict_keys(obj):
  101. """Decode the keys of the given dictionary with utf-8."""
  102. newobj = copy.copy(obj)
  103. for k in obj.keys():
  104. if isinstance(k, bytes):
  105. newobj[k.decode("utf-8")] = newobj[k]
  106. newobj.pop(k)
  107. return newobj
  108. def stringify_param_value(value):
  109. """
  110. Turn a parameter value into a string suitable for the params header of
  111. a Cypher command.
  112. You may pass any value that would be accepted by `json.dumps()`.
  113. Ways in which output differs from that of `str()`:
  114. * Strings are quoted.
  115. * None --> "null".
  116. * In dictionaries, keys are _not_ quoted.
  117. :param value: The parameter value to be turned into a string.
  118. :return: string
  119. """
  120. if isinstance(value, str):
  121. return quote_string(value)
  122. elif value is None:
  123. return "null"
  124. elif isinstance(value, (list, tuple)):
  125. return f'[{",".join(map(stringify_param_value, value))}]'
  126. elif isinstance(value, dict):
  127. return f'{{{",".join(f"{k}:{stringify_param_value(v)}" for k, v in value.items())}}}' # noqa
  128. else:
  129. return str(value)
  130. def get_protocol_version(client):
  131. if isinstance(client, redis.Redis) or isinstance(client, redis.asyncio.Redis):
  132. return client.connection_pool.connection_kwargs.get("protocol")
  133. elif isinstance(client, redis.cluster.AbstractRedisCluster):
  134. return client.nodes_manager.connection_kwargs.get("protocol")