| """Track current scope to easily calculate the corresponding fine-grained target. |
| |
| TODO: Use everywhere where we track targets, including in mypy.errors. |
| """ |
| |
| from __future__ import annotations |
| |
| from contextlib import contextmanager, nullcontext |
| from typing import Iterator, Optional, Tuple |
| from typing_extensions import TypeAlias as _TypeAlias |
| |
| from mypy.nodes import FuncBase, TypeInfo |
| |
| SavedScope: _TypeAlias = Tuple[str, Optional[TypeInfo], Optional[FuncBase]] |
| |
| |
| class Scope: |
| """Track which target we are processing at any given time.""" |
| |
| def __init__(self) -> None: |
| self.module: str | None = None |
| self.classes: list[TypeInfo] = [] |
| self.function: FuncBase | None = None |
| self.functions: list[FuncBase] = [] |
| # Number of nested scopes ignored (that don't get their own separate targets) |
| self.ignored = 0 |
| |
| def current_module_id(self) -> str: |
| assert self.module |
| return self.module |
| |
| def current_target(self) -> str: |
| """Return the current target (non-class; for a class return enclosing module).""" |
| assert self.module |
| if self.function: |
| fullname = self.function.fullname |
| return fullname or "" |
| return self.module |
| |
| def current_full_target(self) -> str: |
| """Return the current target (may be a class).""" |
| assert self.module |
| if self.function: |
| return self.function.fullname |
| if self.classes: |
| return self.classes[-1].fullname |
| return self.module |
| |
| def current_type_name(self) -> str | None: |
| """Return the current type's short name if it exists""" |
| return self.classes[-1].name if self.classes else None |
| |
| def current_function_name(self) -> str | None: |
| """Return the current function's short name if it exists""" |
| return self.function.name if self.function else None |
| |
| @contextmanager |
| def module_scope(self, prefix: str) -> Iterator[None]: |
| self.module = prefix |
| self.classes = [] |
| self.function = None |
| self.ignored = 0 |
| yield |
| assert self.module |
| self.module = None |
| |
| @contextmanager |
| def function_scope(self, fdef: FuncBase) -> Iterator[None]: |
| self.functions.append(fdef) |
| if not self.function: |
| self.function = fdef |
| else: |
| # Nested functions are part of the topmost function target. |
| self.ignored += 1 |
| yield |
| self.functions.pop() |
| if self.ignored: |
| # Leave a scope that's included in the enclosing target. |
| self.ignored -= 1 |
| else: |
| assert self.function |
| self.function = None |
| |
| def outer_functions(self) -> list[FuncBase]: |
| return self.functions[:-1] |
| |
| def enter_class(self, info: TypeInfo) -> None: |
| """Enter a class target scope.""" |
| if not self.function: |
| self.classes.append(info) |
| else: |
| # Classes within functions are part of the enclosing function target. |
| self.ignored += 1 |
| |
| def leave_class(self) -> None: |
| """Leave a class target scope.""" |
| if self.ignored: |
| # Leave a scope that's included in the enclosing target. |
| self.ignored -= 1 |
| else: |
| assert self.classes |
| # Leave the innermost class. |
| self.classes.pop() |
| |
| @contextmanager |
| def class_scope(self, info: TypeInfo) -> Iterator[None]: |
| self.enter_class(info) |
| yield |
| self.leave_class() |
| |
| def save(self) -> SavedScope: |
| """Produce a saved scope that can be entered with saved_scope()""" |
| assert self.module |
| # We only save the innermost class, which is sufficient since |
| # the rest are only needed for when classes are left. |
| cls = self.classes[-1] if self.classes else None |
| return self.module, cls, self.function |
| |
| @contextmanager |
| def saved_scope(self, saved: SavedScope) -> Iterator[None]: |
| module, info, function = saved |
| with self.module_scope(module): |
| with self.class_scope(info) if info else nullcontext(): |
| with self.function_scope(function) if function else nullcontext(): |
| yield |