inference_tip.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  2. # For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE
  3. # Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt
  4. """Transform utilities (filters and decorator)."""
  5. from __future__ import annotations
  6. from collections import OrderedDict
  7. from collections.abc import Generator
  8. from typing import Any, TypeVar
  9. from astroid.context import InferenceContext
  10. from astroid.exceptions import InferenceOverwriteError, UseInferenceDefault
  11. from astroid.nodes import NodeNG
  12. from astroid.typing import (
  13. InferenceResult,
  14. InferFn,
  15. TransformFn,
  16. )
  17. _cache: OrderedDict[
  18. tuple[InferFn[Any], NodeNG, InferenceContext | None], list[InferenceResult]
  19. ] = OrderedDict()
  20. _CURRENTLY_INFERRING: set[tuple[InferFn[Any], NodeNG]] = set()
  21. _NodesT = TypeVar("_NodesT", bound=NodeNG)
  22. def clear_inference_tip_cache() -> None:
  23. """Clear the inference tips cache."""
  24. _cache.clear()
  25. def _inference_tip_cached(func: InferFn[_NodesT]) -> InferFn[_NodesT]:
  26. """Cache decorator used for inference tips."""
  27. def inner(
  28. node: _NodesT,
  29. context: InferenceContext | None = None,
  30. **kwargs: Any,
  31. ) -> Generator[InferenceResult]:
  32. partial_cache_key = (func, node)
  33. if partial_cache_key in _CURRENTLY_INFERRING:
  34. # If through recursion we end up trying to infer the same
  35. # func + node we raise here.
  36. _CURRENTLY_INFERRING.remove(partial_cache_key)
  37. raise UseInferenceDefault
  38. if context is not None and context.is_empty():
  39. # Fresh, empty contexts will defeat the cache.
  40. context = None
  41. try:
  42. yield from _cache[func, node, context]
  43. return
  44. except KeyError:
  45. # Recursion guard with a partial cache key.
  46. # Using the full key causes a recursion error on PyPy.
  47. # It's a pragmatic compromise to avoid so much recursive inference
  48. # with slightly different contexts while still passing the simple
  49. # test cases included with this commit.
  50. _CURRENTLY_INFERRING.add(partial_cache_key)
  51. try:
  52. # May raise UseInferenceDefault
  53. result = _cache[func, node, context] = list(
  54. func(node, context, **kwargs)
  55. )
  56. except Exception as e:
  57. # Suppress the KeyError from the cache miss.
  58. raise e from None
  59. finally:
  60. # Remove recursion guard.
  61. try:
  62. _CURRENTLY_INFERRING.remove(partial_cache_key)
  63. except KeyError:
  64. pass # Recursion may beat us to the punch.
  65. if len(_cache) > 64:
  66. _cache.popitem(last=False)
  67. # https://github.com/pylint-dev/pylint/issues/8686
  68. yield from result # pylint: disable=used-before-assignment
  69. return inner
  70. def inference_tip(
  71. infer_function: InferFn[_NodesT], raise_on_overwrite: bool = False
  72. ) -> TransformFn[_NodesT]:
  73. """Given an instance specific inference function, return a function to be
  74. given to AstroidManager().register_transform to set this inference function.
  75. :param bool raise_on_overwrite: Raise an `InferenceOverwriteError`
  76. if the inference tip will overwrite another. Used for debugging
  77. Typical usage
  78. .. sourcecode:: python
  79. AstroidManager().register_transform(Call, inference_tip(infer_named_tuple),
  80. predicate)
  81. .. Note::
  82. Using an inference tip will override
  83. any previously set inference tip for the given
  84. node. Use a predicate in the transform to prevent
  85. excess overwrites.
  86. """
  87. def transform(
  88. node: _NodesT, infer_function: InferFn[_NodesT] = infer_function
  89. ) -> _NodesT:
  90. if (
  91. raise_on_overwrite
  92. and node._explicit_inference is not None
  93. and node._explicit_inference is not infer_function
  94. ):
  95. raise InferenceOverwriteError(
  96. "Inference already set to {existing_inference}. "
  97. "Trying to overwrite with {new_inference} for {node}".format(
  98. existing_inference=infer_function,
  99. new_inference=node._explicit_inference,
  100. node=node,
  101. )
  102. )
  103. node._explicit_inference = _inference_tip_cached(infer_function)
  104. return node
  105. return transform