120 lines
4.0 KiB
Python
120 lines
4.0 KiB
Python
|
"""Track current scope to easily calculate the corresponding fine-grained target.
|
||
|
|
|
||
|
|
TODO: Use everywhere where we track targets, including in mypy.errors.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from typing import List, Optional, Iterator, Tuple
|
||
|
|
from typing_extensions import TypeAlias as _TypeAlias
|
||
|
|
|
||
|
|
from mypy.backports import nullcontext
|
||
|
|
from mypy.nodes import TypeInfo, FuncBase
|
||
|
|
|
||
|
|
|
||
|
|
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: Optional[str] = None
|
||
|
|
self.classes: List[TypeInfo] = []
|
||
|
|
self.function: Optional[FuncBase] = None
|
||
|
|
# 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) -> Optional[str]:
|
||
|
|
"""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) -> Optional[str]:
|
||
|
|
"""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]:
|
||
|
|
if not self.function:
|
||
|
|
self.function = fdef
|
||
|
|
else:
|
||
|
|
# Nested functions are part of the topmost function target.
|
||
|
|
self.ignored += 1
|
||
|
|
yield
|
||
|
|
if self.ignored:
|
||
|
|
# Leave a scope that's included in the enclosing target.
|
||
|
|
self.ignored -= 1
|
||
|
|
else:
|
||
|
|
assert self.function
|
||
|
|
self.function = None
|
||
|
|
|
||
|
|
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
|