Development of an internal social media platform with personalised dashboards for students
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.

brain_namedtuple_enum.py 10KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285
  1. # Copyright (c) 2012-2015 LOGILAB S.A. (Paris, FRANCE) <contact@logilab.fr>
  2. # Copyright (c) 2014-2016 Claudiu Popa <pcmanticore@gmail.com>
  3. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  4. # For details: https://github.com/PyCQA/astroid/blob/master/COPYING.LESSER
  5. """Astroid hooks for the Python standard library."""
  6. import functools
  7. import sys
  8. import keyword
  9. from textwrap import dedent
  10. from astroid import (
  11. MANAGER, UseInferenceDefault, inference_tip,
  12. InferenceError)
  13. from astroid import arguments
  14. from astroid import exceptions
  15. from astroid import nodes
  16. from astroid.builder import AstroidBuilder, extract_node
  17. from astroid import util
  18. def _infer_first(node, context):
  19. if node is util.Uninferable:
  20. raise UseInferenceDefault
  21. try:
  22. value = next(node.infer(context=context))
  23. if value is util.Uninferable:
  24. raise UseInferenceDefault()
  25. else:
  26. return value
  27. except StopIteration:
  28. raise InferenceError()
  29. def _find_func_form_arguments(node, context):
  30. def _extract_namedtuple_arg_or_keyword(position, key_name=None):
  31. if len(args) > position:
  32. return _infer_first(args[position], context)
  33. if key_name and key_name in found_keywords:
  34. return _infer_first(found_keywords[key_name], context)
  35. args = node.args
  36. keywords = node.keywords
  37. found_keywords = {
  38. keyword.arg: keyword.value for keyword in keywords
  39. } if keywords else {}
  40. name = _extract_namedtuple_arg_or_keyword(
  41. position=0,
  42. key_name='typename'
  43. )
  44. names = _extract_namedtuple_arg_or_keyword(
  45. position=1,
  46. key_name='field_names'
  47. )
  48. if name and names:
  49. return name.value, names
  50. raise UseInferenceDefault()
  51. def infer_func_form(node, base_type, context=None, enum=False):
  52. """Specific inference function for namedtuple or Python 3 enum. """
  53. # node is a Call node, class name as first argument and generated class
  54. # attributes as second argument
  55. # namedtuple or enums list of attributes can be a list of strings or a
  56. # whitespace-separate string
  57. try:
  58. name, names = _find_func_form_arguments(node, context)
  59. try:
  60. attributes = names.value.replace(',', ' ').split()
  61. except AttributeError:
  62. if not enum:
  63. attributes = [_infer_first(const, context).value
  64. for const in names.elts]
  65. else:
  66. # Enums supports either iterator of (name, value) pairs
  67. # or mappings.
  68. # TODO: support only list, tuples and mappings.
  69. if hasattr(names, 'items') and isinstance(names.items, list):
  70. attributes = [_infer_first(const[0], context).value
  71. for const in names.items
  72. if isinstance(const[0], nodes.Const)]
  73. elif hasattr(names, 'elts'):
  74. # Enums can support either ["a", "b", "c"]
  75. # or [("a", 1), ("b", 2), ...], but they can't
  76. # be mixed.
  77. if all(isinstance(const, nodes.Tuple)
  78. for const in names.elts):
  79. attributes = [_infer_first(const.elts[0], context).value
  80. for const in names.elts
  81. if isinstance(const, nodes.Tuple)]
  82. else:
  83. attributes = [_infer_first(const, context).value
  84. for const in names.elts]
  85. else:
  86. raise AttributeError
  87. if not attributes:
  88. raise AttributeError
  89. except (AttributeError, exceptions.InferenceError):
  90. raise UseInferenceDefault()
  91. # If we can't infer the name of the class, don't crash, up to this point
  92. # we know it is a namedtuple anyway.
  93. name = name or 'Uninferable'
  94. # we want to return a Class node instance with proper attributes set
  95. class_node = nodes.ClassDef(name, 'docstring')
  96. class_node.parent = node.parent
  97. # set base class=tuple
  98. class_node.bases.append(base_type)
  99. # XXX add __init__(*attributes) method
  100. for attr in attributes:
  101. fake_node = nodes.EmptyNode()
  102. fake_node.parent = class_node
  103. fake_node.attrname = attr
  104. class_node.instance_attrs[attr] = [fake_node]
  105. return class_node, name, attributes
  106. def _looks_like(node, name):
  107. func = node.func
  108. if isinstance(func, nodes.Attribute):
  109. return func.attrname == name
  110. if isinstance(func, nodes.Name):
  111. return func.name == name
  112. return False
  113. _looks_like_namedtuple = functools.partial(_looks_like, name='namedtuple')
  114. _looks_like_enum = functools.partial(_looks_like, name='Enum')
  115. def infer_named_tuple(node, context=None):
  116. """Specific inference function for namedtuple Call node"""
  117. class_node, name, attributes = infer_func_form(node, nodes.Tuple._proxied,
  118. context=context)
  119. call_site = arguments.CallSite.from_call(node)
  120. func = next(extract_node('import collections; collections.namedtuple').infer())
  121. try:
  122. rename = next(call_site.infer_argument(func, 'rename', context)).bool_value()
  123. except InferenceError:
  124. rename = False
  125. if rename:
  126. attributes = _get_renamed_namedtuple_atributes(attributes)
  127. replace_args = ', '.join(
  128. '{arg}=None'.format(arg=arg)
  129. for arg in attributes
  130. )
  131. field_def = (" {name} = property(lambda self: self[{index:d}], "
  132. "doc='Alias for field number {index:d}')")
  133. field_defs = '\n'.join(field_def.format(name=name, index=index)
  134. for index, name in enumerate(attributes))
  135. fake = AstroidBuilder(MANAGER).string_build('''
  136. class %(name)s(tuple):
  137. __slots__ = ()
  138. _fields = %(fields)r
  139. def _asdict(self):
  140. return self.__dict__
  141. @classmethod
  142. def _make(cls, iterable, new=tuple.__new__, len=len):
  143. return new(cls, iterable)
  144. def _replace(self, %(replace_args)s):
  145. return self
  146. def __getnewargs__(self):
  147. return tuple(self)
  148. %(field_defs)s
  149. ''' % {'name': name,
  150. 'fields': attributes,
  151. 'field_defs': field_defs,
  152. 'replace_args': replace_args})
  153. class_node.locals['_asdict'] = fake.body[0].locals['_asdict']
  154. class_node.locals['_make'] = fake.body[0].locals['_make']
  155. class_node.locals['_replace'] = fake.body[0].locals['_replace']
  156. class_node.locals['_fields'] = fake.body[0].locals['_fields']
  157. for attr in attributes:
  158. class_node.locals[attr] = fake.body[0].locals[attr]
  159. # we use UseInferenceDefault, we can't be a generator so return an iterator
  160. return iter([class_node])
  161. def _get_renamed_namedtuple_atributes(field_names):
  162. names = list(field_names)
  163. seen = set()
  164. for i, name in enumerate(field_names):
  165. if (not all(c.isalnum() or c == '_' for c in name) or keyword.iskeyword(name)
  166. or not name or name[0].isdigit() or name.startswith('_') or name in seen):
  167. names[i] = '_%d' % i
  168. seen.add(name)
  169. return tuple(names)
  170. def infer_enum(node, context=None):
  171. """ Specific inference function for enum Call node. """
  172. enum_meta = extract_node('''
  173. class EnumMeta(object):
  174. 'docstring'
  175. def __call__(self, node):
  176. class EnumAttribute(object):
  177. name = ''
  178. value = 0
  179. return EnumAttribute()
  180. def __iter__(self):
  181. class EnumAttribute(object):
  182. name = ''
  183. value = 0
  184. return [EnumAttribute()]
  185. def __next__(self):
  186. return next(iter(self))
  187. def __getitem__(self, attr):
  188. class Value(object):
  189. @property
  190. def name(self):
  191. return ''
  192. @property
  193. def value(self):
  194. return attr
  195. return Value()
  196. __members__ = ['']
  197. ''')
  198. class_node = infer_func_form(node, enum_meta,
  199. context=context, enum=True)[0]
  200. return iter([class_node.instantiate_class()])
  201. def infer_enum_class(node):
  202. """ Specific inference for enums. """
  203. names = set(('Enum', 'IntEnum', 'enum.Enum', 'enum.IntEnum'))
  204. for basename in node.basenames:
  205. # TODO: doesn't handle subclasses yet. This implementation
  206. # is a hack to support enums.
  207. if basename not in names:
  208. continue
  209. if node.root().name == 'enum':
  210. # Skip if the class is directly from enum module.
  211. break
  212. for local, values in node.locals.items():
  213. if any(not isinstance(value, nodes.AssignName)
  214. for value in values):
  215. continue
  216. stmt = values[0].statement()
  217. if isinstance(stmt, nodes.Assign):
  218. if isinstance(stmt.targets[0], nodes.Tuple):
  219. targets = stmt.targets[0].itered()
  220. else:
  221. targets = stmt.targets
  222. elif isinstance(stmt, nodes.AnnAssign):
  223. targets = [stmt.target]
  224. new_targets = []
  225. for target in targets:
  226. # Replace all the assignments with our mocked class.
  227. classdef = dedent('''
  228. class %(name)s(%(types)s):
  229. @property
  230. def value(self):
  231. # Not the best return.
  232. return None
  233. @property
  234. def name(self):
  235. return %(name)r
  236. ''' % {'name': target.name, 'types': ', '.join(node.basenames)})
  237. fake = AstroidBuilder(MANAGER).string_build(classdef)[target.name]
  238. fake.parent = target.parent
  239. for method in node.mymethods():
  240. fake.locals[method.name] = [method]
  241. new_targets.append(fake.instantiate_class())
  242. node.locals[local] = new_targets
  243. break
  244. return node
  245. MANAGER.register_transform(nodes.Call, inference_tip(infer_named_tuple),
  246. _looks_like_namedtuple)
  247. MANAGER.register_transform(nodes.Call, inference_tip(infer_enum),
  248. _looks_like_enum)
  249. MANAGER.register_transform(nodes.ClassDef, infer_enum_class)