scope.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. """Track current scope to easily calculate the corresponding fine-grained target.
  2. TODO: Use everywhere where we track targets, including in mypy.errors.
  3. """
  4. from __future__ import annotations
  5. from contextlib import contextmanager, nullcontext
  6. from typing import Iterator, Optional, Tuple
  7. from typing_extensions import TypeAlias as _TypeAlias
  8. from mypy.nodes import FuncBase, TypeInfo
  9. SavedScope: _TypeAlias = Tuple[str, Optional[TypeInfo], Optional[FuncBase]]
  10. class Scope:
  11. """Track which target we are processing at any given time."""
  12. def __init__(self) -> None:
  13. self.module: str | None = None
  14. self.classes: list[TypeInfo] = []
  15. self.function: FuncBase | None = None
  16. self.functions: list[FuncBase] = []
  17. # Number of nested scopes ignored (that don't get their own separate targets)
  18. self.ignored = 0
  19. def current_module_id(self) -> str:
  20. assert self.module
  21. return self.module
  22. def current_target(self) -> str:
  23. """Return the current target (non-class; for a class return enclosing module)."""
  24. assert self.module
  25. if self.function:
  26. fullname = self.function.fullname
  27. return fullname or ""
  28. return self.module
  29. def current_full_target(self) -> str:
  30. """Return the current target (may be a class)."""
  31. assert self.module
  32. if self.function:
  33. return self.function.fullname
  34. if self.classes:
  35. return self.classes[-1].fullname
  36. return self.module
  37. def current_type_name(self) -> str | None:
  38. """Return the current type's short name if it exists"""
  39. return self.classes[-1].name if self.classes else None
  40. def current_function_name(self) -> str | None:
  41. """Return the current function's short name if it exists"""
  42. return self.function.name if self.function else None
  43. @contextmanager
  44. def module_scope(self, prefix: str) -> Iterator[None]:
  45. self.module = prefix
  46. self.classes = []
  47. self.function = None
  48. self.ignored = 0
  49. yield
  50. assert self.module
  51. self.module = None
  52. @contextmanager
  53. def function_scope(self, fdef: FuncBase) -> Iterator[None]:
  54. self.functions.append(fdef)
  55. if not self.function:
  56. self.function = fdef
  57. else:
  58. # Nested functions are part of the topmost function target.
  59. self.ignored += 1
  60. yield
  61. self.functions.pop()
  62. if self.ignored:
  63. # Leave a scope that's included in the enclosing target.
  64. self.ignored -= 1
  65. else:
  66. assert self.function
  67. self.function = None
  68. def outer_functions(self) -> list[FuncBase]:
  69. return self.functions[:-1]
  70. def enter_class(self, info: TypeInfo) -> None:
  71. """Enter a class target scope."""
  72. if not self.function:
  73. self.classes.append(info)
  74. else:
  75. # Classes within functions are part of the enclosing function target.
  76. self.ignored += 1
  77. def leave_class(self) -> None:
  78. """Leave a class target scope."""
  79. if self.ignored:
  80. # Leave a scope that's included in the enclosing target.
  81. self.ignored -= 1
  82. else:
  83. assert self.classes
  84. # Leave the innermost class.
  85. self.classes.pop()
  86. @contextmanager
  87. def class_scope(self, info: TypeInfo) -> Iterator[None]:
  88. self.enter_class(info)
  89. yield
  90. self.leave_class()
  91. def save(self) -> SavedScope:
  92. """Produce a saved scope that can be entered with saved_scope()"""
  93. assert self.module
  94. # We only save the innermost class, which is sufficient since
  95. # the rest are only needed for when classes are left.
  96. cls = self.classes[-1] if self.classes else None
  97. return self.module, cls, self.function
  98. @contextmanager
  99. def saved_scope(self, saved: SavedScope) -> Iterator[None]:
  100. module, info, function = saved
  101. with self.module_scope(module):
  102. with self.class_scope(info) if info else nullcontext():
  103. with self.function_scope(function) if function else nullcontext():
  104. yield