helpers.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  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. try:
  55. if isinstance(response[i + 1], list):
  56. res["Child iterators"].append(parse_list_to_dict(response[i + 1]))
  57. except IndexError:
  58. pass
  59. elif isinstance(response[i + 1], list):
  60. res["Child iterators"] = [parse_list_to_dict(response[i + 1])]
  61. else:
  62. try:
  63. res[response[i]] = float(response[i + 1])
  64. except (TypeError, ValueError):
  65. res[response[i]] = response[i + 1]
  66. return res
  67. def parse_to_dict(response):
  68. if response is None:
  69. return {}
  70. res = {}
  71. for det in response:
  72. if not isinstance(det, list) or not det:
  73. continue
  74. if len(det) == 1:
  75. res[det[0]] = True
  76. elif isinstance(det[1], list):
  77. res[det[0]] = parse_list_to_dict(det[1])
  78. else:
  79. try: # try to set the attribute. may be provided without value
  80. try: # try to convert the value to float
  81. res[det[0]] = float(det[1])
  82. except (TypeError, ValueError):
  83. res[det[0]] = det[1]
  84. except IndexError:
  85. pass
  86. return res
  87. def random_string(length=10):
  88. """
  89. Returns a random N character long string.
  90. """
  91. return "".join( # nosec
  92. random.choice(string.ascii_lowercase) for x in range(length)
  93. )
  94. def quote_string(v):
  95. """
  96. RedisGraph strings must be quoted,
  97. quote_string wraps given v with quotes incase
  98. v is a string.
  99. """
  100. if isinstance(v, bytes):
  101. v = v.decode()
  102. elif not isinstance(v, str):
  103. return v
  104. if len(v) == 0:
  105. return '""'
  106. v = v.replace("\\", "\\\\")
  107. v = v.replace('"', '\\"')
  108. return f'"{v}"'
  109. def decode_dict_keys(obj):
  110. """Decode the keys of the given dictionary with utf-8."""
  111. newobj = copy.copy(obj)
  112. for k in obj.keys():
  113. if isinstance(k, bytes):
  114. newobj[k.decode("utf-8")] = newobj[k]
  115. newobj.pop(k)
  116. return newobj
  117. def stringify_param_value(value):
  118. """
  119. Turn a parameter value into a string suitable for the params header of
  120. a Cypher command.
  121. You may pass any value that would be accepted by `json.dumps()`.
  122. Ways in which output differs from that of `str()`:
  123. * Strings are quoted.
  124. * None --> "null".
  125. * In dictionaries, keys are _not_ quoted.
  126. :param value: The parameter value to be turned into a string.
  127. :return: string
  128. """
  129. if isinstance(value, str):
  130. return quote_string(value)
  131. elif value is None:
  132. return "null"
  133. elif isinstance(value, (list, tuple)):
  134. return f'[{",".join(map(stringify_param_value, value))}]'
  135. elif isinstance(value, dict):
  136. return f'{{{",".join(f"{k}:{stringify_param_value(v)}" for k, v in value.items())}}}' # noqa
  137. else:
  138. return str(value)
  139. def get_protocol_version(client):
  140. if isinstance(client, redis.Redis) or isinstance(client, redis.asyncio.Redis):
  141. return client.connection_pool.connection_kwargs.get("protocol")
  142. elif isinstance(client, redis.cluster.AbstractRedisCluster):
  143. return client.nodes_manager.connection_kwargs.get("protocol")