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_typing.py 2.5KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. # Copyright (c) 2016 David Euresti <david@dropbox.com>
  2. """Astroid hooks for typing.py support."""
  3. import textwrap
  4. from astroid import (
  5. MANAGER, UseInferenceDefault, extract_node, inference_tip,
  6. nodes, InferenceError)
  7. from astroid.nodes import List, Tuple
  8. TYPING_NAMEDTUPLE_BASENAMES = {
  9. 'NamedTuple',
  10. 'typing.NamedTuple'
  11. }
  12. def infer_typing_namedtuple(node, context=None):
  13. """Infer a typing.NamedTuple(...) call."""
  14. # This is essentially a namedtuple with different arguments
  15. # so we extract the args and infer a named tuple.
  16. try:
  17. func = next(node.func.infer())
  18. except InferenceError:
  19. raise UseInferenceDefault
  20. if func.qname() != 'typing.NamedTuple':
  21. raise UseInferenceDefault
  22. if len(node.args) != 2:
  23. raise UseInferenceDefault
  24. if not isinstance(node.args[1], (List, Tuple)):
  25. raise UseInferenceDefault
  26. names = []
  27. for elt in node.args[1].elts:
  28. if not isinstance(elt, (List, Tuple)):
  29. raise UseInferenceDefault
  30. if len(elt.elts) != 2:
  31. raise UseInferenceDefault
  32. names.append(elt.elts[0].as_string())
  33. typename = node.args[0].as_string()
  34. node = extract_node('namedtuple(%(typename)s, (%(fields)s,)) ' %
  35. {'typename': typename, 'fields': ",".join(names)})
  36. return node.infer(context=context)
  37. def infer_typing_namedtuple_class(node, context=None):
  38. """Infer a subclass of typing.NamedTuple"""
  39. # Check if it has the corresponding bases
  40. annassigns_fields = [
  41. annassign.target.name for annassign in node.body
  42. if isinstance(annassign, nodes.AnnAssign)
  43. ]
  44. code = textwrap.dedent('''
  45. from collections import namedtuple
  46. namedtuple({typename!r}, {fields!r})
  47. ''').format(
  48. typename=node.name,
  49. fields=",".join(annassigns_fields)
  50. )
  51. node = extract_node(code)
  52. return node.infer(context=context)
  53. def has_namedtuple_base(node):
  54. """Predicate for class inference tip
  55. :type node: ClassDef
  56. :rtype: bool
  57. """
  58. return set(node.basenames) & TYPING_NAMEDTUPLE_BASENAMES
  59. def looks_like_typing_namedtuple(node):
  60. func = node.func
  61. if isinstance(func, nodes.Attribute):
  62. return func.attrname == 'NamedTuple'
  63. if isinstance(func, nodes.Name):
  64. return func.name == 'NamedTuple'
  65. return False
  66. MANAGER.register_transform(
  67. nodes.Call,
  68. inference_tip(infer_typing_namedtuple),
  69. looks_like_typing_namedtuple
  70. )
  71. MANAGER.register_transform(
  72. nodes.ClassDef,
  73. inference_tip(infer_typing_namedtuple_class),
  74. has_namedtuple_base
  75. )