mypy/scope.py

"""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