_ast.py 4.0 KB

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