You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

356 lines
12 KiB

4 years ago
  1. import math
  2. import json
  3. from jmespath import exceptions
  4. from jmespath.compat import string_type as STRING_TYPE
  5. from jmespath.compat import get_methods, with_metaclass
  6. # python types -> jmespath types
  7. TYPES_MAP = {
  8. 'bool': 'boolean',
  9. 'list': 'array',
  10. 'dict': 'object',
  11. 'NoneType': 'null',
  12. 'unicode': 'string',
  13. 'str': 'string',
  14. 'float': 'number',
  15. 'int': 'number',
  16. 'long': 'number',
  17. 'OrderedDict': 'object',
  18. '_Projection': 'array',
  19. '_Expression': 'expref',
  20. }
  21. # jmespath types -> python types
  22. REVERSE_TYPES_MAP = {
  23. 'boolean': ('bool',),
  24. 'array': ('list', '_Projection'),
  25. 'object': ('dict', 'OrderedDict',),
  26. 'null': ('None',),
  27. 'string': ('unicode', 'str'),
  28. 'number': ('float', 'int', 'long'),
  29. 'expref': ('_Expression',),
  30. }
  31. def signature(*arguments):
  32. def _record_signature(func):
  33. func.signature = arguments
  34. return func
  35. return _record_signature
  36. class FunctionRegistry(type):
  37. def __init__(cls, name, bases, attrs):
  38. cls._populate_function_table()
  39. super(FunctionRegistry, cls).__init__(name, bases, attrs)
  40. def _populate_function_table(cls):
  41. function_table = {}
  42. # Any method with a @signature decorator that also
  43. # starts with "_func_" is registered as a function.
  44. # _func_max_by -> max_by function.
  45. for name, method in get_methods(cls):
  46. if not name.startswith('_func_'):
  47. continue
  48. signature = getattr(method, 'signature', None)
  49. if signature is not None:
  50. function_table[name[6:]] = {
  51. 'function': method,
  52. 'signature': signature,
  53. }
  54. cls.FUNCTION_TABLE = function_table
  55. class Functions(with_metaclass(FunctionRegistry, object)):
  56. FUNCTION_TABLE = {
  57. }
  58. def call_function(self, function_name, resolved_args):
  59. try:
  60. spec = self.FUNCTION_TABLE[function_name]
  61. except KeyError:
  62. raise exceptions.UnknownFunctionError(
  63. "Unknown function: %s()" % function_name)
  64. function = spec['function']
  65. signature = spec['signature']
  66. self._validate_arguments(resolved_args, signature, function_name)
  67. return function(self, *resolved_args)
  68. def _validate_arguments(self, args, signature, function_name):
  69. if signature and signature[-1].get('variadic'):
  70. if len(args) < len(signature):
  71. raise exceptions.VariadictArityError(
  72. len(signature), len(args), function_name)
  73. elif len(args) != len(signature):
  74. raise exceptions.ArityError(
  75. len(signature), len(args), function_name)
  76. return self._type_check(args, signature, function_name)
  77. def _type_check(self, actual, signature, function_name):
  78. for i in range(len(signature)):
  79. allowed_types = signature[i]['types']
  80. if allowed_types:
  81. self._type_check_single(actual[i], allowed_types,
  82. function_name)
  83. def _type_check_single(self, current, types, function_name):
  84. # Type checking involves checking the top level type,
  85. # and in the case of arrays, potentially checking the types
  86. # of each element.
  87. allowed_types, allowed_subtypes = self._get_allowed_pytypes(types)
  88. # We're not using isinstance() on purpose.
  89. # The type model for jmespath does not map
  90. # 1-1 with python types (booleans are considered
  91. # integers in python for example).
  92. actual_typename = type(current).__name__
  93. if actual_typename not in allowed_types:
  94. raise exceptions.JMESPathTypeError(
  95. function_name, current,
  96. self._convert_to_jmespath_type(actual_typename), types)
  97. # If we're dealing with a list type, we can have
  98. # additional restrictions on the type of the list
  99. # elements (for example a function can require a
  100. # list of numbers or a list of strings).
  101. # Arrays are the only types that can have subtypes.
  102. if allowed_subtypes:
  103. self._subtype_check(current, allowed_subtypes,
  104. types, function_name)
  105. def _get_allowed_pytypes(self, types):
  106. allowed_types = []
  107. allowed_subtypes = []
  108. for t in types:
  109. type_ = t.split('-', 1)
  110. if len(type_) == 2:
  111. type_, subtype = type_
  112. allowed_subtypes.append(REVERSE_TYPES_MAP[subtype])
  113. else:
  114. type_ = type_[0]
  115. allowed_types.extend(REVERSE_TYPES_MAP[type_])
  116. return allowed_types, allowed_subtypes
  117. def _subtype_check(self, current, allowed_subtypes, types, function_name):
  118. if len(allowed_subtypes) == 1:
  119. # The easy case, we know up front what type
  120. # we need to validate.
  121. allowed_subtypes = allowed_subtypes[0]
  122. for element in current:
  123. actual_typename = type(element).__name__
  124. if actual_typename not in allowed_subtypes:
  125. raise exceptions.JMESPathTypeError(
  126. function_name, element, actual_typename, types)
  127. elif len(allowed_subtypes) > 1 and current:
  128. # Dynamic type validation. Based on the first
  129. # type we see, we validate that the remaining types
  130. # match.
  131. first = type(current[0]).__name__
  132. for subtypes in allowed_subtypes:
  133. if first in subtypes:
  134. allowed = subtypes
  135. break
  136. else:
  137. raise exceptions.JMESPathTypeError(
  138. function_name, current[0], first, types)
  139. for element in current:
  140. actual_typename = type(element).__name__
  141. if actual_typename not in allowed:
  142. raise exceptions.JMESPathTypeError(
  143. function_name, element, actual_typename, types)
  144. @signature({'types': ['number']})
  145. def _func_abs(self, arg):
  146. return abs(arg)
  147. @signature({'types': ['array-number']})
  148. def _func_avg(self, arg):
  149. if arg:
  150. return sum(arg) / float(len(arg))
  151. else:
  152. return None
  153. @signature({'types': [], 'variadic': True})
  154. def _func_not_null(self, *arguments):
  155. for argument in arguments:
  156. if argument is not None:
  157. return argument
  158. @signature({'types': []})
  159. def _func_to_array(self, arg):
  160. if isinstance(arg, list):
  161. return arg
  162. else:
  163. return [arg]
  164. @signature({'types': []})
  165. def _func_to_string(self, arg):
  166. if isinstance(arg, STRING_TYPE):
  167. return arg
  168. else:
  169. return json.dumps(arg, separators=(',', ':'),
  170. default=str)
  171. @signature({'types': []})
  172. def _func_to_number(self, arg):
  173. if isinstance(arg, (list, dict, bool)):
  174. return None
  175. elif arg is None:
  176. return None
  177. elif isinstance(arg, (int, float)):
  178. return arg
  179. else:
  180. try:
  181. return int(arg)
  182. except ValueError:
  183. try:
  184. return float(arg)
  185. except ValueError:
  186. return None
  187. @signature({'types': ['array', 'string']}, {'types': []})
  188. def _func_contains(self, subject, search):
  189. return search in subject
  190. @signature({'types': ['string', 'array', 'object']})
  191. def _func_length(self, arg):
  192. return len(arg)
  193. @signature({'types': ['string']}, {'types': ['string']})
  194. def _func_ends_with(self, search, suffix):
  195. return search.endswith(suffix)
  196. @signature({'types': ['string']}, {'types': ['string']})
  197. def _func_starts_with(self, search, suffix):
  198. return search.startswith(suffix)
  199. @signature({'types': ['array', 'string']})
  200. def _func_reverse(self, arg):
  201. if isinstance(arg, STRING_TYPE):
  202. return arg[::-1]
  203. else:
  204. return list(reversed(arg))
  205. @signature({"types": ['number']})
  206. def _func_ceil(self, arg):
  207. return math.ceil(arg)
  208. @signature({"types": ['number']})
  209. def _func_floor(self, arg):
  210. return math.floor(arg)
  211. @signature({"types": ['string']}, {"types": ['array-string']})
  212. def _func_join(self, separator, array):
  213. return separator.join(array)
  214. @signature({'types': ['expref']}, {'types': ['array']})
  215. def _func_map(self, expref, arg):
  216. result = []
  217. for element in arg:
  218. result.append(expref.visit(expref.expression, element))
  219. return result
  220. @signature({"types": ['array-number', 'array-string']})
  221. def _func_max(self, arg):
  222. if arg:
  223. return max(arg)
  224. else:
  225. return None
  226. @signature({"types": ["object"], "variadic": True})
  227. def _func_merge(self, *arguments):
  228. merged = {}
  229. for arg in arguments:
  230. merged.update(arg)
  231. return merged
  232. @signature({"types": ['array-number', 'array-string']})
  233. def _func_min(self, arg):
  234. if arg:
  235. return min(arg)
  236. else:
  237. return None
  238. @signature({"types": ['array-string', 'array-number']})
  239. def _func_sort(self, arg):
  240. return list(sorted(arg))
  241. @signature({"types": ['array-number']})
  242. def _func_sum(self, arg):
  243. return sum(arg)
  244. @signature({"types": ['object']})
  245. def _func_keys(self, arg):
  246. # To be consistent with .values()
  247. # should we also return the indices of a list?
  248. return list(arg.keys())
  249. @signature({"types": ['object']})
  250. def _func_values(self, arg):
  251. return list(arg.values())
  252. @signature({'types': []})
  253. def _func_type(self, arg):
  254. if isinstance(arg, STRING_TYPE):
  255. return "string"
  256. elif isinstance(arg, bool):
  257. return "boolean"
  258. elif isinstance(arg, list):
  259. return "array"
  260. elif isinstance(arg, dict):
  261. return "object"
  262. elif isinstance(arg, (float, int)):
  263. return "number"
  264. elif arg is None:
  265. return "null"
  266. @signature({'types': ['array']}, {'types': ['expref']})
  267. def _func_sort_by(self, array, expref):
  268. if not array:
  269. return array
  270. # sort_by allows for the expref to be either a number of
  271. # a string, so we have some special logic to handle this.
  272. # We evaluate the first array element and verify that it's
  273. # either a string of a number. We then create a key function
  274. # that validates that type, which requires that remaining array
  275. # elements resolve to the same type as the first element.
  276. required_type = self._convert_to_jmespath_type(
  277. type(expref.visit(expref.expression, array[0])).__name__)
  278. if required_type not in ['number', 'string']:
  279. raise exceptions.JMESPathTypeError(
  280. 'sort_by', array[0], required_type, ['string', 'number'])
  281. keyfunc = self._create_key_func(expref,
  282. [required_type],
  283. 'sort_by')
  284. return list(sorted(array, key=keyfunc))
  285. @signature({'types': ['array']}, {'types': ['expref']})
  286. def _func_min_by(self, array, expref):
  287. keyfunc = self._create_key_func(expref,
  288. ['number', 'string'],
  289. 'min_by')
  290. return min(array, key=keyfunc)
  291. @signature({'types': ['array']}, {'types': ['expref']})
  292. def _func_max_by(self, array, expref):
  293. keyfunc = self._create_key_func(expref,
  294. ['number', 'string'],
  295. 'min_by')
  296. return max(array, key=keyfunc)
  297. def _create_key_func(self, expref, allowed_types, function_name):
  298. def keyfunc(x):
  299. result = expref.visit(expref.expression, x)
  300. actual_typename = type(result).__name__
  301. jmespath_type = self._convert_to_jmespath_type(actual_typename)
  302. # allowed_types is in term of jmespath types, not python types.
  303. if jmespath_type not in allowed_types:
  304. raise exceptions.JMESPathTypeError(
  305. function_name, result, jmespath_type, allowed_types)
  306. return result
  307. return keyfunc
  308. def _convert_to_jmespath_type(self, pyobject):
  309. return TYPES_MAP.get(pyobject, 'unknown')