mypy-dev 1.9.0.dev0__py3-none-any.whl

mypy/refinfo.py

@@ -0,0 +1,92 @@
+"""Find line-level reference information from a mypy AST (undocumented feature)"""
+
+from __future__ import annotations
+
+from mypy.nodes import (
+    LDEF,
+    Expression,
+    FuncDef,
+    MemberExpr,
+    MypyFile,
+    NameExpr,
+    RefExpr,
+    SymbolNode,
+    TypeInfo,
+)
+from mypy.traverser import TraverserVisitor
+from mypy.typeops import tuple_fallback
+from mypy.types import (
+    FunctionLike,
+    Instance,
+    TupleType,
+    Type,
+    TypeType,
+    TypeVarLikeType,
+    get_proper_type,
+)
+
+
+class RefInfoVisitor(TraverserVisitor):
+    def __init__(self, type_map: dict[Expression, Type]) -> None:
+        super().__init__()
+        self.type_map = type_map
+        self.data: list[dict[str, object]] = []
+
+    def visit_name_expr(self, expr: NameExpr) -> None:
+        super().visit_name_expr(expr)
+        self.record_ref_expr(expr)
+
+    def visit_member_expr(self, expr: MemberExpr) -> None:
+        super().visit_member_expr(expr)
+        self.record_ref_expr(expr)
+
+    def visit_func_def(self, func: FuncDef) -> None:
+        if func.expanded:
+            for item in func.expanded:
+                if isinstance(item, FuncDef):
+                    super().visit_func_def(item)
+        else:
+            super().visit_func_def(func)
+
+    def record_ref_expr(self, expr: RefExpr) -> None:
+        fullname = None
+        if expr.kind != LDEF and "." in expr.fullname:
+            fullname = expr.fullname
+        elif isinstance(expr, MemberExpr):
+            typ = self.type_map.get(expr.expr)
+            sym = None
+            if isinstance(expr.expr, RefExpr):
+                sym = expr.expr.node
+            if typ:
+                tfn = type_fullname(typ, sym)
+                if tfn:
+                    fullname = f"{tfn}.{expr.name}"
+            if not fullname:
+                fullname = f"*.{expr.name}"
+        if fullname is not None:
+            self.data.append({"line": expr.line, "column": expr.column, "target": fullname})
+
+
+def type_fullname(typ: Type, node: SymbolNode | None = None) -> str | None:
+    typ = get_proper_type(typ)
+    if isinstance(typ, Instance):
+        return typ.type.fullname
+    elif isinstance(typ, TypeType):
+        return type_fullname(typ.item)
+    elif isinstance(typ, FunctionLike) and typ.is_type_obj():
+        if isinstance(node, TypeInfo):
+            return node.fullname
+        return type_fullname(typ.fallback)
+    elif isinstance(typ, TupleType):
+        return type_fullname(tuple_fallback(typ))
+    elif isinstance(typ, TypeVarLikeType):
+        return type_fullname(typ.upper_bound)
+    return None
+
+
+def get_undocumented_ref_info_json(
+    tree: MypyFile, type_map: dict[Expression, Type]
+) -> list[dict[str, object]]:
+    visitor = RefInfoVisitor(type_map)
+    tree.accept(visitor)
+    return visitor.data

mypy/renaming.py

@@ -0,0 +1,568 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Final, Iterator
+
+from mypy.nodes import (
+    AssignmentStmt,
+    Block,
+    BreakStmt,
+    ClassDef,
+    ContinueStmt,
+    ForStmt,
+    FuncDef,
+    Import,
+    ImportAll,
+    ImportFrom,
+    IndexExpr,
+    ListExpr,
+    Lvalue,
+    MatchStmt,
+    MemberExpr,
+    MypyFile,
+    NameExpr,
+    StarExpr,
+    TryStmt,
+    TupleExpr,
+    WhileStmt,
+    WithStmt,
+)
+from mypy.patterns import AsPattern
+from mypy.traverser import TraverserVisitor
+
+# Scope kinds
+FILE: Final = 0
+FUNCTION: Final = 1
+CLASS: Final = 2
+
+
+class VariableRenameVisitor(TraverserVisitor):
+    """Rename variables to allow redefinition of variables.
+
+    For example, consider this code:
+
+      x = 0
+      f(x)
+
+      x = "a"
+      g(x)
+
+    It will be transformed like this:
+
+      x' = 0
+      f(x')
+
+      x = "a"
+      g(x)
+
+    There will be two independent variables (x' and x) that will have separate
+    inferred types. The publicly exposed variant will get the non-suffixed name.
+    This is the last definition at module top level and the first definition
+    (argument) within a function.
+
+    Renaming only happens for assignments within the same block. Renaming is
+    performed before semantic analysis, immediately after parsing.
+
+    The implementation performs a rudimentary static analysis. The analysis is
+    overly conservative to keep things simple.
+    """
+
+    def __init__(self) -> None:
+        # Counter for labeling new blocks
+        self.block_id = 0
+        # Number of surrounding try statements that disallow variable redefinition
+        self.disallow_redef_depth = 0
+        # Number of surrounding loop statements
+        self.loop_depth = 0
+        # Map block id to loop depth.
+        self.block_loop_depth: dict[int, int] = {}
+        # Stack of block ids being processed.
+        self.blocks: list[int] = []
+        # List of scopes; each scope maps short (unqualified) name to block id.
+        self.var_blocks: list[dict[str, int]] = []
+
+        # References to variables that we may need to rename. List of
+        # scopes; each scope is a mapping from name to list of collections
+        # of names that refer to the same logical variable.
+        self.refs: list[dict[str, list[list[NameExpr]]]] = []
+        # Number of reads of the most recent definition of a variable (per scope)
+        self.num_reads: list[dict[str, int]] = []
+        # Kinds of nested scopes (FILE, FUNCTION or CLASS)
+        self.scope_kinds: list[int] = []
+
+    def visit_mypy_file(self, file_node: MypyFile) -> None:
+        """Rename variables within a file.
+
+        This is the main entry point to this class.
+        """
+        self.clear()
+        with self.enter_scope(FILE), self.enter_block():
+            for d in file_node.defs:
+                d.accept(self)
+
+    def visit_func_def(self, fdef: FuncDef) -> None:
+        # Conservatively do not allow variable defined before a function to
+        # be redefined later, since function could refer to either definition.
+        self.reject_redefinition_of_vars_in_scope()
+
+        with self.enter_scope(FUNCTION), self.enter_block():
+            for arg in fdef.arguments:
+                name = arg.variable.name
+                # 'self' can't be redefined since it's special as it allows definition of
+                # attributes. 'cls' can't be used to define attributes so we can ignore it.
+                can_be_redefined = name != "self"  # TODO: Proper check
+                self.record_assignment(arg.variable.name, can_be_redefined)
+                self.handle_arg(name)
+
+            for stmt in fdef.body.body:
+                stmt.accept(self)
+
+    def visit_class_def(self, cdef: ClassDef) -> None:
+        self.reject_redefinition_of_vars_in_scope()
+        with self.enter_scope(CLASS):
+            super().visit_class_def(cdef)
+
+    def visit_block(self, block: Block) -> None:
+        with self.enter_block():
+            super().visit_block(block)
+
+    def visit_while_stmt(self, stmt: WhileStmt) -> None:
+        with self.enter_loop():
+            super().visit_while_stmt(stmt)
+
+    def visit_for_stmt(self, stmt: ForStmt) -> None:
+        stmt.expr.accept(self)
+        self.analyze_lvalue(stmt.index, True)
+        # Also analyze as non-lvalue so that every for loop index variable is assumed to be read.
+        stmt.index.accept(self)
+        with self.enter_loop():
+            stmt.body.accept(self)
+        if stmt.else_body:
+            stmt.else_body.accept(self)
+
+    def visit_break_stmt(self, stmt: BreakStmt) -> None:
+        self.reject_redefinition_of_vars_in_loop()
+
+    def visit_continue_stmt(self, stmt: ContinueStmt) -> None:
+        self.reject_redefinition_of_vars_in_loop()
+
+    def visit_try_stmt(self, stmt: TryStmt) -> None:
+        # Variables defined by a try statement get special treatment in the
+        # type checker which allows them to be always redefined, so no need to
+        # do renaming here.
+        with self.enter_try():
+            super().visit_try_stmt(stmt)
+
+    def visit_with_stmt(self, stmt: WithStmt) -> None:
+        for expr in stmt.expr:
+            expr.accept(self)
+        for target in stmt.target:
+            if target is not None:
+                self.analyze_lvalue(target)
+        # We allow redefinitions in the body of a with statement for
+        # convenience.  This is unsafe since with statements can affect control
+        # flow by catching exceptions, but this is rare except for
+        # assertRaises() and other similar functions, where the exception is
+        # raised by the last statement in the body, which usually isn't a
+        # problem.
+        stmt.body.accept(self)
+
+    def visit_import(self, imp: Import) -> None:
+        for id, as_id in imp.ids:
+            self.record_assignment(as_id or id, False)
+
+    def visit_import_from(self, imp: ImportFrom) -> None:
+        for id, as_id in imp.names:
+            self.record_assignment(as_id or id, False)
+
+    def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
+        s.rvalue.accept(self)
+        for lvalue in s.lvalues:
+            self.analyze_lvalue(lvalue)
+
+    def visit_match_stmt(self, s: MatchStmt) -> None:
+        s.subject.accept(self)
+        for i in range(len(s.patterns)):
+            with self.enter_block():
+                s.patterns[i].accept(self)
+                guard = s.guards[i]
+                if guard is not None:
+                    guard.accept(self)
+                # We already entered a block, so visit this block's statements directly
+                for stmt in s.bodies[i].body:
+                    stmt.accept(self)
+
+    def visit_capture_pattern(self, p: AsPattern) -> None:
+        if p.name is not None:
+            self.analyze_lvalue(p.name)
+
+    def analyze_lvalue(self, lvalue: Lvalue, is_nested: bool = False) -> None:
+        """Process assignment; in particular, keep track of (re)defined names.
+
+        Args:
+            is_nested: True for non-outermost Lvalue in a multiple assignment such as
+                "x, y = ..."
+        """
+        if isinstance(lvalue, NameExpr):
+            name = lvalue.name
+            is_new = self.record_assignment(name, True)
+            if is_new:
+                self.handle_def(lvalue)
+            else:
+                self.handle_refine(lvalue)
+            if is_nested:
+                # This allows these to be redefined freely even if never read. Multiple
+                # assignment like "x, _ _ = y" defines dummy variables that are never read.
+                self.handle_ref(lvalue)
+        elif isinstance(lvalue, (ListExpr, TupleExpr)):
+            for item in lvalue.items:
+                self.analyze_lvalue(item, is_nested=True)
+        elif isinstance(lvalue, MemberExpr):
+            lvalue.expr.accept(self)
+        elif isinstance(lvalue, IndexExpr):
+            lvalue.base.accept(self)
+            lvalue.index.accept(self)
+        elif isinstance(lvalue, StarExpr):
+            # Propagate is_nested since in a typical use case like "x, *rest = ..." 'rest' may
+            # be freely reused.
+            self.analyze_lvalue(lvalue.expr, is_nested=is_nested)
+
+    def visit_name_expr(self, expr: NameExpr) -> None:
+        self.handle_ref(expr)
+
+    # Helpers for renaming references
+
+    def handle_arg(self, name: str) -> None:
+        """Store function argument."""
+        self.refs[-1][name] = [[]]
+        self.num_reads[-1][name] = 0
+
+    def handle_def(self, expr: NameExpr) -> None:
+        """Store new name definition."""
+        name = expr.name
+        names = self.refs[-1].setdefault(name, [])
+        names.append([expr])
+        self.num_reads[-1][name] = 0
+
+    def handle_refine(self, expr: NameExpr) -> None:
+        """Store assignment to an existing name (that replaces previous value, if any)."""
+        name = expr.name
+        if name in self.refs[-1]:
+            names = self.refs[-1][name]
+            if not names:
+                names.append([])
+            names[-1].append(expr)
+
+    def handle_ref(self, expr: NameExpr) -> None:
+        """Store reference to defined name."""
+        name = expr.name
+        if name in self.refs[-1]:
+            names = self.refs[-1][name]
+            if not names:
+                names.append([])
+            names[-1].append(expr)
+        num_reads = self.num_reads[-1]
+        num_reads[name] = num_reads.get(name, 0) + 1
+
+    def flush_refs(self) -> None:
+        """Rename all references within the current scope.
+
+        This will be called at the end of a scope.
+        """
+        is_func = self.scope_kinds[-1] == FUNCTION
+        for refs in self.refs[-1].values():
+            if len(refs) == 1:
+                # Only one definition -- no renaming needed.
+                continue
+            if is_func:
+                # In a function, don't rename the first definition, as it
+                # may be an argument that must preserve the name.
+                to_rename = refs[1:]
+            else:
+                # At module top level, don't rename the final definition,
+                # as it will be publicly visible outside the module.
+                to_rename = refs[:-1]
+            for i, item in enumerate(to_rename):
+                rename_refs(item, i)
+        self.refs.pop()
+
+    # Helpers for determining which assignments define new variables
+
+    def clear(self) -> None:
+        self.blocks = []
+        self.var_blocks = []
+
+    @contextmanager
+    def enter_block(self) -> Iterator[None]:
+        self.block_id += 1
+        self.blocks.append(self.block_id)
+        self.block_loop_depth[self.block_id] = self.loop_depth
+        try:
+            yield
+        finally:
+            self.blocks.pop()
+
+    @contextmanager
+    def enter_try(self) -> Iterator[None]:
+        self.disallow_redef_depth += 1
+        try:
+            yield
+        finally:
+            self.disallow_redef_depth -= 1
+
+    @contextmanager
+    def enter_loop(self) -> Iterator[None]:
+        self.loop_depth += 1
+        try:
+            yield
+        finally:
+            self.loop_depth -= 1
+
+    def current_block(self) -> int:
+        return self.blocks[-1]
+
+    @contextmanager
+    def enter_scope(self, kind: int) -> Iterator[None]:
+        self.var_blocks.append({})
+        self.refs.append({})
+        self.num_reads.append({})
+        self.scope_kinds.append(kind)
+        try:
+            yield
+        finally:
+            self.flush_refs()
+            self.var_blocks.pop()
+            self.num_reads.pop()
+            self.scope_kinds.pop()
+
+    def is_nested(self) -> int:
+        return len(self.var_blocks) > 1
+
+    def reject_redefinition_of_vars_in_scope(self) -> None:
+        """Make it impossible to redefine defined variables in the current scope.
+
+        This is used if we encounter a function definition that
+        can make it ambiguous which definition is live. Example:
+
+          x = 0
+
+          def f() -> int:
+              return x
+
+          x = ''  # Error -- cannot redefine x across function definition
+        """
+        var_blocks = self.var_blocks[-1]
+        for key in var_blocks:
+            var_blocks[key] = -1
+
+    def reject_redefinition_of_vars_in_loop(self) -> None:
+        """Reject redefinition of variables in the innermost loop.
+
+        If there is an early exit from a loop, there may be ambiguity about which
+        value may escape the loop. Example where this matters:
+
+          while f():
+              x = 0
+              if g():
+                  break
+              x = ''  # Error -- not a redefinition
+          reveal_type(x)  # int
+
+        This method ensures that the second assignment to 'x' doesn't introduce a new
+        variable.
+        """
+        var_blocks = self.var_blocks[-1]
+        for key, block in var_blocks.items():
+            if self.block_loop_depth.get(block) == self.loop_depth:
+                var_blocks[key] = -1
+
+    def record_assignment(self, name: str, can_be_redefined: bool) -> bool:
+        """Record assignment to given name and return True if it defines a new variable.
+
+        Args:
+            can_be_redefined: If True, allows assignment in the same block to redefine
+                this name (if this is a new definition)
+        """
+        if self.num_reads[-1].get(name, -1) == 0:
+            # Only set, not read, so no reason to redefine
+            return False
+        if self.disallow_redef_depth > 0:
+            # Can't redefine within try/with a block.
+            can_be_redefined = False
+        block = self.current_block()
+        var_blocks = self.var_blocks[-1]
+        if name not in var_blocks:
+            # New definition in this scope.
+            if can_be_redefined:
+                # Store the block where this was defined to allow redefinition in
+                # the same block only.
+                var_blocks[name] = block
+            else:
+                # This doesn't support arbitrary redefinition.
+                var_blocks[name] = -1
+            return True
+        elif var_blocks[name] == block:
+            # Redefinition -- defines a new variable with the same name.
+            return True
+        else:
+            # Assigns to an existing variable.
+            return False
+
+
+class LimitedVariableRenameVisitor(TraverserVisitor):
+    """Perform some limited variable renaming in with statements.
+
+    This allows reusing a variable in multiple with statements with
+    different types. For example, the two instances of 'x' can have
+    incompatible types:
+
+       with C() as x:
+           f(x)
+       with D() as x:
+           g(x)
+
+    The above code gets renamed conceptually into this (not valid Python!):
+
+       with C() as x':
+           f(x')
+       with D() as x:
+           g(x)
+
+    If there's a reference to a variable defined in 'with' outside the
+    statement, or if there's any trickiness around variable visibility
+    (e.g. function definitions), we give up and won't perform renaming.
+
+    The main use case is to allow binding both readable and writable
+    binary files into the same variable. These have different types:
+
+        with open(fnam, 'rb') as f: ...
+        with open(fnam, 'wb') as f: ...
+    """
+
+    def __init__(self) -> None:
+        # Short names of variables bound in with statements using "as"
+        # in a surrounding scope
+        self.bound_vars: list[str] = []
+        # Stack of names that can't be safely renamed, per scope ('*' means that
+        # no names can be renamed)
+        self.skipped: list[set[str]] = []
+        # References to variables that we may need to rename. Stack of
+        # scopes; each scope is a mapping from name to list of collections
+        # of names that refer to the same logical variable.
+        self.refs: list[dict[str, list[list[NameExpr]]]] = []
+
+    def visit_mypy_file(self, file_node: MypyFile) -> None:
+        """Rename variables within a file.
+
+        This is the main entry point to this class.
+        """
+        with self.enter_scope():
+            for d in file_node.defs:
+                d.accept(self)
+
+    def visit_func_def(self, fdef: FuncDef) -> None:
+        self.reject_redefinition_of_vars_in_scope()
+        with self.enter_scope():
+            for arg in fdef.arguments:
+                self.record_skipped(arg.variable.name)
+            super().visit_func_def(fdef)
+
+    def visit_class_def(self, cdef: ClassDef) -> None:
+        self.reject_redefinition_of_vars_in_scope()
+        with self.enter_scope():
+            super().visit_class_def(cdef)
+
+    def visit_with_stmt(self, stmt: WithStmt) -> None:
+        for expr in stmt.expr:
+            expr.accept(self)
+        old_len = len(self.bound_vars)
+        for target in stmt.target:
+            if target is not None:
+                self.analyze_lvalue(target)
+        for target in stmt.target:
+            if target:
+                target.accept(self)
+        stmt.body.accept(self)
+
+        while len(self.bound_vars) > old_len:
+            self.bound_vars.pop()
+
+    def analyze_lvalue(self, lvalue: Lvalue) -> None:
+        if isinstance(lvalue, NameExpr):
+            name = lvalue.name
+            if name in self.bound_vars:
+                # Name bound in a surrounding with statement, so it can be renamed
+                self.visit_name_expr(lvalue)
+            else:
+                var_info = self.refs[-1]
+                if name not in var_info:
+                    var_info[name] = []
+                var_info[name].append([])
+                self.bound_vars.append(name)
+        elif isinstance(lvalue, (ListExpr, TupleExpr)):
+            for item in lvalue.items:
+                self.analyze_lvalue(item)
+        elif isinstance(lvalue, MemberExpr):
+            lvalue.expr.accept(self)
+        elif isinstance(lvalue, IndexExpr):
+            lvalue.base.accept(self)
+            lvalue.index.accept(self)
+        elif isinstance(lvalue, StarExpr):
+            self.analyze_lvalue(lvalue.expr)
+
+    def visit_import(self, imp: Import) -> None:
+        # We don't support renaming imports
+        for id, as_id in imp.ids:
+            self.record_skipped(as_id or id)
+
+    def visit_import_from(self, imp: ImportFrom) -> None:
+        # We don't support renaming imports
+        for id, as_id in imp.names:
+            self.record_skipped(as_id or id)
+
+    def visit_import_all(self, imp: ImportAll) -> None:
+        # Give up, since we don't know all imported names yet
+        self.reject_redefinition_of_vars_in_scope()
+
+    def visit_name_expr(self, expr: NameExpr) -> None:
+        name = expr.name
+        if name in self.bound_vars:
+            # Record reference so that it can be renamed later
+            for scope in reversed(self.refs):
+                if name in scope:
+                    scope[name][-1].append(expr)
+        else:
+            self.record_skipped(name)
+
+    @contextmanager
+    def enter_scope(self) -> Iterator[None]:
+        self.skipped.append(set())
+        self.refs.append({})
+        yield None
+        self.flush_refs()
+
+    def reject_redefinition_of_vars_in_scope(self) -> None:
+        self.record_skipped("*")
+
+    def record_skipped(self, name: str) -> None:
+        self.skipped[-1].add(name)
+
+    def flush_refs(self) -> None:
+        ref_dict = self.refs.pop()
+        skipped = self.skipped.pop()
+        if "*" not in skipped:
+            for name, refs in ref_dict.items():
+                if len(refs) <= 1 or name in skipped:
+                    continue
+                # At module top level we must not rename the final definition,
+                # as it may be publicly visible
+                to_rename = refs[:-1]
+                for i, item in enumerate(to_rename):
+                    rename_refs(item, i)
+
+
+def rename_refs(names: list[NameExpr], index: int) -> None:
+    name = names[0].name
+    new_name = name + "'" * (index + 1)
+    for expr in names:
+        expr.name = new_name