_ast.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. import ast
  2. import sys
  3. import types
  4. from collections import namedtuple
  5. from functools import partial
  6. from typing import Dict, Optional
  7. from astroid.const import PY38_PLUS, Context
  8. if sys.version_info >= (3, 8):
  9. # On Python 3.8, typed_ast was merged back into `ast`
  10. _ast_py3: Optional[types.ModuleType] = ast
  11. else:
  12. try:
  13. import typed_ast.ast3 as _ast_py3
  14. except ImportError:
  15. _ast_py3 = None
  16. FunctionType = namedtuple("FunctionType", ["argtypes", "returns"])
  17. class ParserModule(
  18. namedtuple(
  19. "ParserModule",
  20. [
  21. "module",
  22. "unary_op_classes",
  23. "cmp_op_classes",
  24. "bool_op_classes",
  25. "bin_op_classes",
  26. "context_classes",
  27. ],
  28. )
  29. ):
  30. def parse(self, string: str, type_comments=True):
  31. if self.module is _ast_py3:
  32. if PY38_PLUS:
  33. parse_func = partial(self.module.parse, type_comments=type_comments)
  34. else:
  35. parse_func = partial(
  36. self.module.parse, feature_version=sys.version_info.minor
  37. )
  38. else:
  39. parse_func = self.module.parse
  40. return parse_func(string)
  41. def parse_function_type_comment(type_comment: str) -> Optional[FunctionType]:
  42. """Given a correct type comment, obtain a FunctionType object"""
  43. if _ast_py3 is None:
  44. return None
  45. func_type = _ast_py3.parse(type_comment, "<type_comment>", "func_type") # type: ignore[attr-defined]
  46. return FunctionType(argtypes=func_type.argtypes, returns=func_type.returns)
  47. def get_parser_module(type_comments=True) -> ParserModule:
  48. parser_module = ast
  49. if type_comments and _ast_py3:
  50. parser_module = _ast_py3
  51. unary_op_classes = _unary_operators_from_module(parser_module)
  52. cmp_op_classes = _compare_operators_from_module(parser_module)
  53. bool_op_classes = _bool_operators_from_module(parser_module)
  54. bin_op_classes = _binary_operators_from_module(parser_module)
  55. context_classes = _contexts_from_module(parser_module)
  56. return ParserModule(
  57. parser_module,
  58. unary_op_classes,
  59. cmp_op_classes,
  60. bool_op_classes,
  61. bin_op_classes,
  62. context_classes,
  63. )
  64. def _unary_operators_from_module(module):
  65. return {module.UAdd: "+", module.USub: "-", module.Not: "not", module.Invert: "~"}
  66. def _binary_operators_from_module(module):
  67. binary_operators = {
  68. module.Add: "+",
  69. module.BitAnd: "&",
  70. module.BitOr: "|",
  71. module.BitXor: "^",
  72. module.Div: "/",
  73. module.FloorDiv: "//",
  74. module.MatMult: "@",
  75. module.Mod: "%",
  76. module.Mult: "*",
  77. module.Pow: "**",
  78. module.Sub: "-",
  79. module.LShift: "<<",
  80. module.RShift: ">>",
  81. }
  82. return binary_operators
  83. def _bool_operators_from_module(module):
  84. return {module.And: "and", module.Or: "or"}
  85. def _compare_operators_from_module(module):
  86. return {
  87. module.Eq: "==",
  88. module.Gt: ">",
  89. module.GtE: ">=",
  90. module.In: "in",
  91. module.Is: "is",
  92. module.IsNot: "is not",
  93. module.Lt: "<",
  94. module.LtE: "<=",
  95. module.NotEq: "!=",
  96. module.NotIn: "not in",
  97. }
  98. def _contexts_from_module(module) -> Dict[ast.expr_context, Context]:
  99. return {
  100. module.Load: Context.Load,
  101. module.Store: Context.Store,
  102. module.Del: Context.Del,
  103. module.Param: Context.Store,
  104. }