brain_collections.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. # Copyright (c) 2016, 2018, 2020 Claudiu Popa <pcmanticore@gmail.com>
  2. # Copyright (c) 2016-2017 Łukasz Rogalski <rogalski.91@gmail.com>
  3. # Copyright (c) 2017 Derek Gustafson <degustaf@gmail.com>
  4. # Copyright (c) 2018 Ioana Tagirta <ioana.tagirta@gmail.com>
  5. # Copyright (c) 2019 Hugo van Kemenade <hugovk@users.noreply.github.com>
  6. # Copyright (c) 2020-2021 hippo91 <guillaume.peillex@gmail.com>
  7. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  8. # Copyright (c) 2021 John Belmonte <john@neggie.net>
  9. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  10. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  11. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  12. from astroid.brain.helpers import register_module_extender
  13. from astroid.builder import extract_node, parse
  14. from astroid.const import PY39_PLUS
  15. from astroid.exceptions import AttributeInferenceError
  16. from astroid.manager import AstroidManager
  17. from astroid.nodes.scoped_nodes import ClassDef
  18. def _collections_transform():
  19. return parse(
  20. """
  21. class defaultdict(dict):
  22. default_factory = None
  23. def __missing__(self, key): pass
  24. def __getitem__(self, key): return default_factory
  25. """
  26. + _deque_mock()
  27. + _ordered_dict_mock()
  28. )
  29. def _deque_mock():
  30. base_deque_class = """
  31. class deque(object):
  32. maxlen = 0
  33. def __init__(self, iterable=None, maxlen=None):
  34. self.iterable = iterable or []
  35. def append(self, x): pass
  36. def appendleft(self, x): pass
  37. def clear(self): pass
  38. def count(self, x): return 0
  39. def extend(self, iterable): pass
  40. def extendleft(self, iterable): pass
  41. def pop(self): return self.iterable[0]
  42. def popleft(self): return self.iterable[0]
  43. def remove(self, value): pass
  44. def reverse(self): return reversed(self.iterable)
  45. def rotate(self, n=1): return self
  46. def __iter__(self): return self
  47. def __reversed__(self): return self.iterable[::-1]
  48. def __getitem__(self, index): return self.iterable[index]
  49. def __setitem__(self, index, value): pass
  50. def __delitem__(self, index): pass
  51. def __bool__(self): return bool(self.iterable)
  52. def __nonzero__(self): return bool(self.iterable)
  53. def __contains__(self, o): return o in self.iterable
  54. def __len__(self): return len(self.iterable)
  55. def __copy__(self): return deque(self.iterable)
  56. def copy(self): return deque(self.iterable)
  57. def index(self, x, start=0, end=0): return 0
  58. def insert(self, i, x): pass
  59. def __add__(self, other): pass
  60. def __iadd__(self, other): pass
  61. def __mul__(self, other): pass
  62. def __imul__(self, other): pass
  63. def __rmul__(self, other): pass"""
  64. if PY39_PLUS:
  65. base_deque_class += """
  66. @classmethod
  67. def __class_getitem__(self, item): return cls"""
  68. return base_deque_class
  69. def _ordered_dict_mock():
  70. base_ordered_dict_class = """
  71. class OrderedDict(dict):
  72. def __reversed__(self): return self[::-1]
  73. def move_to_end(self, key, last=False): pass"""
  74. if PY39_PLUS:
  75. base_ordered_dict_class += """
  76. @classmethod
  77. def __class_getitem__(cls, item): return cls"""
  78. return base_ordered_dict_class
  79. register_module_extender(AstroidManager(), "collections", _collections_transform)
  80. def _looks_like_subscriptable(node: ClassDef) -> bool:
  81. """
  82. Returns True if the node corresponds to a ClassDef of the Collections.abc module that
  83. supports subscripting
  84. :param node: ClassDef node
  85. """
  86. if node.qname().startswith("_collections") or node.qname().startswith(
  87. "collections"
  88. ):
  89. try:
  90. node.getattr("__class_getitem__")
  91. return True
  92. except AttributeInferenceError:
  93. pass
  94. return False
  95. CLASS_GET_ITEM_TEMPLATE = """
  96. @classmethod
  97. def __class_getitem__(cls, item):
  98. return cls
  99. """
  100. def easy_class_getitem_inference(node, context=None):
  101. # Here __class_getitem__ exists but is quite a mess to infer thus
  102. # put an easy inference tip
  103. func_to_add = extract_node(CLASS_GET_ITEM_TEMPLATE)
  104. node.locals["__class_getitem__"] = [func_to_add]
  105. if PY39_PLUS:
  106. # Starting with Python39 some objects of the collection module are subscriptable
  107. # thanks to the __class_getitem__ method but the way it is implemented in
  108. # _collection_abc makes it difficult to infer. (We would have to handle AssignName inference in the
  109. # getitem method of the ClassDef class) Instead we put here a mock of the __class_getitem__ method
  110. AstroidManager().register_transform(
  111. ClassDef, easy_class_getitem_inference, _looks_like_subscriptable
  112. )