brain_functools.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. # Copyright (c) 2016, 2018-2020 Claudiu Popa <pcmanticore@gmail.com>
  2. # Copyright (c) 2018 hippo91 <guillaume.peillex@gmail.com>
  3. # Copyright (c) 2018 Bryce Guinta <bryce.paul.guinta@gmail.com>
  4. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  5. # Copyright (c) 2021 Alphadelta14 <alpha@alphaservcomputing.solutions>
  6. """Astroid hooks for understanding functools library module."""
  7. from functools import partial
  8. from itertools import chain
  9. from astroid import BoundMethod, arguments, extract_node, helpers, objects
  10. from astroid.exceptions import InferenceError, UseInferenceDefault
  11. from astroid.inference_tip import inference_tip
  12. from astroid.interpreter import objectmodel
  13. from astroid.manager import AstroidManager
  14. from astroid.nodes.node_classes import AssignName, Attribute, Call, Name
  15. from astroid.nodes.scoped_nodes import FunctionDef
  16. from astroid.util import Uninferable
  17. LRU_CACHE = "functools.lru_cache"
  18. class LruWrappedModel(objectmodel.FunctionModel):
  19. """Special attribute model for functions decorated with functools.lru_cache.
  20. The said decorators patches at decoration time some functions onto
  21. the decorated function.
  22. """
  23. @property
  24. def attr___wrapped__(self):
  25. return self._instance
  26. @property
  27. def attr_cache_info(self):
  28. cache_info = extract_node(
  29. """
  30. from functools import _CacheInfo
  31. _CacheInfo(0, 0, 0, 0)
  32. """
  33. )
  34. class CacheInfoBoundMethod(BoundMethod):
  35. def infer_call_result(self, caller, context=None):
  36. yield helpers.safe_infer(cache_info)
  37. return CacheInfoBoundMethod(proxy=self._instance, bound=self._instance)
  38. @property
  39. def attr_cache_clear(self):
  40. node = extract_node("""def cache_clear(self): pass""")
  41. return BoundMethod(proxy=node, bound=self._instance.parent.scope())
  42. def _transform_lru_cache(node, context=None) -> None:
  43. # TODO: this is not ideal, since the node should be immutable,
  44. # but due to https://github.com/PyCQA/astroid/issues/354,
  45. # there's not much we can do now.
  46. # Replacing the node would work partially, because,
  47. # in pylint, the old node would still be available, leading
  48. # to spurious false positives.
  49. node.special_attributes = LruWrappedModel()(node)
  50. def _functools_partial_inference(node, context=None):
  51. call = arguments.CallSite.from_call(node, context=context)
  52. number_of_positional = len(call.positional_arguments)
  53. if number_of_positional < 1:
  54. raise UseInferenceDefault("functools.partial takes at least one argument")
  55. if number_of_positional == 1 and not call.keyword_arguments:
  56. raise UseInferenceDefault(
  57. "functools.partial needs at least to have some filled arguments"
  58. )
  59. partial_function = call.positional_arguments[0]
  60. try:
  61. inferred_wrapped_function = next(partial_function.infer(context=context))
  62. except (InferenceError, StopIteration) as exc:
  63. raise UseInferenceDefault from exc
  64. if inferred_wrapped_function is Uninferable:
  65. raise UseInferenceDefault("Cannot infer the wrapped function")
  66. if not isinstance(inferred_wrapped_function, FunctionDef):
  67. raise UseInferenceDefault("The wrapped function is not a function")
  68. # Determine if the passed keywords into the callsite are supported
  69. # by the wrapped function.
  70. if not inferred_wrapped_function.args:
  71. function_parameters = []
  72. else:
  73. function_parameters = chain(
  74. inferred_wrapped_function.args.args or (),
  75. inferred_wrapped_function.args.posonlyargs or (),
  76. inferred_wrapped_function.args.kwonlyargs or (),
  77. )
  78. parameter_names = {
  79. param.name for param in function_parameters if isinstance(param, AssignName)
  80. }
  81. if set(call.keyword_arguments) - parameter_names:
  82. raise UseInferenceDefault("wrapped function received unknown parameters")
  83. partial_function = objects.PartialFunction(
  84. call,
  85. name=inferred_wrapped_function.name,
  86. doc=inferred_wrapped_function.doc,
  87. lineno=inferred_wrapped_function.lineno,
  88. col_offset=inferred_wrapped_function.col_offset,
  89. parent=node.parent,
  90. )
  91. partial_function.postinit(
  92. args=inferred_wrapped_function.args,
  93. body=inferred_wrapped_function.body,
  94. decorators=inferred_wrapped_function.decorators,
  95. returns=inferred_wrapped_function.returns,
  96. type_comment_returns=inferred_wrapped_function.type_comment_returns,
  97. type_comment_args=inferred_wrapped_function.type_comment_args,
  98. )
  99. return iter((partial_function,))
  100. def _looks_like_lru_cache(node):
  101. """Check if the given function node is decorated with lru_cache."""
  102. if not node.decorators:
  103. return False
  104. for decorator in node.decorators.nodes:
  105. if not isinstance(decorator, Call):
  106. continue
  107. if _looks_like_functools_member(decorator, "lru_cache"):
  108. return True
  109. return False
  110. def _looks_like_functools_member(node, member) -> bool:
  111. """Check if the given Call node is a functools.partial call"""
  112. if isinstance(node.func, Name):
  113. return node.func.name == member
  114. if isinstance(node.func, Attribute):
  115. return (
  116. node.func.attrname == member
  117. and isinstance(node.func.expr, Name)
  118. and node.func.expr.name == "functools"
  119. )
  120. return False
  121. _looks_like_partial = partial(_looks_like_functools_member, member="partial")
  122. AstroidManager().register_transform(
  123. FunctionDef, _transform_lru_cache, _looks_like_lru_cache
  124. )
  125. AstroidManager().register_transform(
  126. Call,
  127. inference_tip(_functools_partial_inference),
  128. _looks_like_partial,
  129. )