brain_re.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  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. from typing import Optional
  4. from astroid import context, inference_tip, nodes
  5. from astroid.brain.helpers import register_module_extender
  6. from astroid.builder import extract_node, parse
  7. from astroid.const import PY37_PLUS, PY39_PLUS
  8. from astroid.manager import AstroidManager
  9. def _re_transform():
  10. # Since Python 3.6 there is the RegexFlag enum
  11. # where every entry will be exposed via updating globals()
  12. return parse(
  13. """
  14. import sre_compile
  15. ASCII = sre_compile.SRE_FLAG_ASCII
  16. IGNORECASE = sre_compile.SRE_FLAG_IGNORECASE
  17. LOCALE = sre_compile.SRE_FLAG_LOCALE
  18. UNICODE = sre_compile.SRE_FLAG_UNICODE
  19. MULTILINE = sre_compile.SRE_FLAG_MULTILINE
  20. DOTALL = sre_compile.SRE_FLAG_DOTALL
  21. VERBOSE = sre_compile.SRE_FLAG_VERBOSE
  22. A = ASCII
  23. I = IGNORECASE
  24. L = LOCALE
  25. U = UNICODE
  26. M = MULTILINE
  27. S = DOTALL
  28. X = VERBOSE
  29. TEMPLATE = sre_compile.SRE_FLAG_TEMPLATE
  30. T = TEMPLATE
  31. DEBUG = sre_compile.SRE_FLAG_DEBUG
  32. """
  33. )
  34. register_module_extender(AstroidManager(), "re", _re_transform)
  35. CLASS_GETITEM_TEMPLATE = """
  36. @classmethod
  37. def __class_getitem__(cls, item):
  38. return cls
  39. """
  40. def _looks_like_pattern_or_match(node: nodes.Call) -> bool:
  41. """Check for re.Pattern or re.Match call in stdlib.
  42. Match these patterns from stdlib/re.py
  43. ```py
  44. Pattern = type(...)
  45. Match = type(...)
  46. ```
  47. """
  48. return (
  49. node.root().name == "re"
  50. and isinstance(node.func, nodes.Name)
  51. and node.func.name == "type"
  52. and isinstance(node.parent, nodes.Assign)
  53. and len(node.parent.targets) == 1
  54. and isinstance(node.parent.targets[0], nodes.AssignName)
  55. and node.parent.targets[0].name in {"Pattern", "Match"}
  56. )
  57. def infer_pattern_match(
  58. node: nodes.Call, ctx: Optional[context.InferenceContext] = None
  59. ):
  60. """Infer re.Pattern and re.Match as classes. For PY39+ add `__class_getitem__`."""
  61. class_def = nodes.ClassDef(
  62. name=node.parent.targets[0].name,
  63. lineno=node.lineno,
  64. col_offset=node.col_offset,
  65. parent=node.parent,
  66. )
  67. if PY39_PLUS:
  68. func_to_add = extract_node(CLASS_GETITEM_TEMPLATE)
  69. class_def.locals["__class_getitem__"] = [func_to_add]
  70. return iter([class_def])
  71. if PY37_PLUS:
  72. AstroidManager().register_transform(
  73. nodes.Call, inference_tip(infer_pattern_match), _looks_like_pattern_or_match
  74. )