cinderx 2026.1.16.2__cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

cinderx/compiler/strict/common.py

@@ -0,0 +1,507 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# pyre-strict
+
+from __future__ import annotations
+
+import ast
+import os.path
+import symtable
+import sys
+import typing
+from ast import (
+    alias,
+    AST,
+    AsyncFunctionDef,
+    ClassDef,
+    comprehension,
+    copy_location,
+    DictComp,
+    FunctionDef,
+    GeneratorExp,
+    iter_fields,
+    Lambda,
+    ListComp,
+    NodeVisitor,
+    SetComp,
+    Try,
+)
+from collections import deque
+from collections.abc import Callable, Mapping, MutableMapping
+from contextlib import nullcontext
+from dataclasses import dataclass
+from symtable import Class, SymbolTable
+from typing import final, Generic, TypeVar
+
+from .runtime import freeze_type, mutable
+
+
+# Increment this whenever we change the output of the strict modules
+# interpreter. It must stay below 32768 (15 bits), because we use the high bit
+# to encode strictness of the module.
+MAGIC_NUMBER = 55
+
+
+DEFAULT_STUB_PATH: str = os.path.dirname(__file__) + "/stubs"
+
+
+def make_fixed_modules() -> Mapping[str, Mapping[str, object]]:
+    typing_members = {}
+    for name in typing.__all__:
+        typing_members[name] = getattr(typing, name)
+    strict_mod_members = {
+        "freeze_type": freeze_type,
+        "mutable": mutable,
+    }
+
+    return {
+        "typing": typing_members,
+        "strict_modules": dict(strict_mod_members),
+        "__strict__": strict_mod_members,
+    }
+
+
+FIXED_MODULES: Mapping[str, Mapping[str, object]] = make_fixed_modules()
+
+
+TVar = TypeVar("TScope")
+TScopeData = TypeVar("TData", covariant=True)
+
+SymbolMap = dict[AST, SymbolTable]
+
+
+@dataclass
+class StrictModuleError(Exception):
+    msg: str
+    filename: str
+    lineno: int
+    col: int
+    metadata: str = ""
+
+
+def _is_scoped_generator_node(node: AST) -> bool:
+    return sys.version_info < (3, 12) or not isinstance(
+        node,
+        (
+            ast.ListComp,
+            ast.DictComp,
+            ast.SetComp,
+        ),
+    )
+
+
+# SymbolTableType is new in 3.13
+if sys.version_info >= (3, 14):
+    ANNOTATION_TYPE: str = symtable.SymbolTableType.ANNOTATION
+
+    def is_annotation(x: symtable.SymbolTable) -> bool:
+        return x.get_type() == ANNOTATION_TYPE
+else:
+
+    def is_annotation(x: symtable.SymbolTable) -> bool:
+        return False
+
+
+@final
+class SymbolMapBuilder(ast.NodeVisitor):
+    def __init__(self, symbols: SymbolTable) -> None:
+        self.symbol_stack: deque[SymbolTable] = deque([symbols])
+        children = self.symbol_stack.popleft().get_children()
+        self.symbol_stack.extendleft(
+            x for x in reversed(children) if not is_annotation(x)
+        )
+        self.mapping: SymbolMap = {}
+
+    def _process_scope_node(self, node: AST) -> None:
+        current_symbol = self.mapping[node] = self.symbol_stack.popleft()
+        children = current_symbol.get_children()
+        self.symbol_stack.extendleft(
+            x for x in reversed(children) if not is_annotation(x)
+        )
+
+    def visit_ClassDef(self, node: ClassDef) -> None:
+        for child in node.bases:
+            self.visit(child)
+        for child in node.keywords:
+            self.visit(child)
+        for child in node.decorator_list:
+            self.visit(child)
+        self._process_scope_node(node)
+        for child in node.body:
+            self.visit(child)
+
+    def visit_FunctionDef(self, node: FunctionDef) -> None:
+        self._visit_function_like(node)
+
+    def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> None:
+        self._visit_function_like(node)
+
+    def visit_Lambda(self, node: Lambda) -> None:
+        if node.args:
+            self.visit(node.args)
+        self._process_scope_node(node)
+        self.visit(node.body)
+
+    def _visit_function_like(self, node: FunctionDef | AsyncFunctionDef) -> None:
+        # args -> returns -> decorator_list -> body
+        if node.args:
+            self.visit(node.args)
+        returns = node.returns
+        if returns:
+            self.visit(returns)
+        for child in node.decorator_list:
+            self.visit(child)
+        self._process_scope_node(node)
+        for child in node.body:
+            self.visit(child)
+
+    def visit_ListComp(self, node: ListComp) -> None:
+        self._visit_generator_like(node, [node.elt], node.generators)
+
+    def visit_SetComp(self, node: SetComp) -> None:
+        self._visit_generator_like(node, [node.elt], node.generators)
+
+    def visit_DictComp(self, node: DictComp) -> None:
+        self._visit_generator_like(node, [node.key, node.value], node.generators)
+
+    def visit_GeneratorExp(self, node: GeneratorExp) -> None:
+        self._visit_generator_like(node, [node.elt], node.generators)
+
+    def _visit_generator_like(
+        self, node: AST, elements: list[AST], comprehensions: list[comprehension]
+    ) -> None:
+        # first iter is visited in the outer scope
+        self.visit(comprehensions[0].iter)
+        # everything else is in the inner scope
+        if _is_scoped_generator_node(node):
+            # In 3.12 list comprehensions are inlined
+            self._process_scope_node(node)
+        # process first comprehension, without iter
+        for child in comprehensions[0].ifs:
+            self.visit(child)
+        for comp in comprehensions[1:]:
+            self.visit(comp)
+        # process elements
+        for element in elements:
+            self.visit(element)
+
+    def visit_Try(self, node: Try) -> None:
+        # Need to match the order the symbol visitor constructs these in, which
+        # walks orelse before handlers.
+        for val in node.body:
+            self.visit(val)
+        for val in node.orelse:
+            self.visit(val)
+        for val in node.handlers:
+            self.visit(val)
+        for val in node.finalbody:
+            self.visit(val)
+
+
+def get_symbol_map(node: ast.AST, symtable: SymbolTable) -> SymbolMap:
+    visitor = SymbolMapBuilder(symtable)
+    visitor.visit(node)
+    return visitor.mapping
+
+
+@final
+class SymbolScope(Generic[TVar, TScopeData]):
+    def __init__(
+        self,
+        symbols: SymbolTable,
+        scope_data: TScopeData,
+        vars: MutableMapping[str, TVar] | None = None,
+        invisible: bool = False,
+    ) -> None:
+        self.symbols = symbols
+        self.vars = vars
+        self.scope_data = scope_data
+        self.invisible = invisible
+
+    def __getitem__(self, name: str) -> TVar:
+        v = self.vars
+        if v is None:
+            raise KeyError(name)
+
+        return v[name]
+
+    def __setitem__(self, name: str, value: TVar) -> None:
+        v = self.vars
+        if v is None:
+            v = self.vars = {}
+
+        v[name] = value
+
+    def __delitem__(self, name: str) -> None:
+        v = self.vars
+        if v is None:
+            raise KeyError(name)
+
+        del v[name]
+
+    def __contains__(self, name: str) -> bool:
+        v = self.vars
+        if v is None:
+            return False
+
+        return name in v
+
+
+def mangle_priv_name(name: str, scopes: list[SymbolScope[TVar, TScopeData]]) -> str:
+    if name.startswith("__") and not name.endswith("__"):
+        # symtable has name mangled private names.  Walk the scope list
+        # backwards and apply the mangled class name
+        for scope in reversed(scopes):
+            if isinstance(scope.symbols, symtable.Class) and not scope.invisible:
+                return "_" + scope.symbols.get_name().lstrip("_") + name
+
+    return name
+
+
+def imported_name(name: alias) -> str:
+    return name.asname or name.name.partition(".")[0]
+
+
+@final
+class ScopeContextManager(Generic[TVar, TScopeData]):
+    def __init__(
+        self, parent: ScopeStack[TVar, TScopeData], scope: SymbolScope[TVar, TScopeData]
+    ) -> None:
+        self.parent = parent
+        self.scope = scope
+
+    def __enter__(self) -> SymbolScope[TVar, TScopeData]:
+        self.parent.push(self.scope)
+        return self.scope
+
+    def __exit__(
+        self, exc_type: type[Exception], exc_val: Exception, exc_tb: object
+    ) -> None:
+        self.parent.pop()
+
+
+@final
+class ScopeStack(Generic[TVar, TScopeData]):
+    def __init__(
+        self,
+        *scopes: SymbolScope[TVar, TScopeData],
+        symbol_map: SymbolMap,
+        scope_factory: Callable[
+            [SymbolTable, AST, MutableMapping[str, TVar] | None],
+            SymbolScope[TVar, TScopeData],
+        ] = lambda symtable, node, vars: SymbolScope(symtable, None),
+    ) -> None:
+        self.scopes: list[SymbolScope[TVar, TScopeData]] = list(scopes)
+        self.symbol_map = symbol_map
+        self.scope_factory: Callable[
+            [SymbolTable, AST, MutableMapping[str, TVar] | None],
+            SymbolScope[TVar, TScopeData],
+        ] = scope_factory
+
+    def push(self, scope: SymbolScope[TVar, TScopeData]) -> None:
+        self.scopes.append(scope)
+
+    def pop(self) -> None:
+        self.scopes.pop()
+
+    @property
+    def current_symbols(self) -> SymbolTable:
+        return self.scopes[-1].symbols
+
+    @property
+    def current(self) -> SymbolScope[TVar, TScopeData]:
+        return self.scopes[-1]
+
+    @property
+    def in_global_scope(self) -> bool:
+        return len(self.scopes) == 1
+
+    @property
+    def in_class_scope(self) -> bool:
+        return isinstance(self.scopes[-1].symbols, Class)
+
+    def with_node_scope(
+        self, node: AST, vars: MutableMapping[str, TVar] | None = None
+    ) -> ScopeContextManager[TVar, TScopeData] | nullcontext[None]:
+        if not _is_scoped_generator_node(node):
+            assert isinstance(node, (ast.ListComp, ast.DictComp, ast.SetComp))
+            # In 3.12 list/dict/set comprehensions are inlined
+            return nullcontext()
+
+        next_symtable = self.symbol_map[node]
+        return ScopeContextManager(self, self.scope_factory(next_symtable, node, vars))
+
+    def is_global(self, name: str) -> bool:
+        if self.in_global_scope:
+            return True
+
+        return self.current_symbols.lookup(
+            mangle_priv_name(name, self.scopes)
+        ).is_global()
+
+    def is_nonlocal(self, name: str) -> bool:
+        return (
+            not self.is_global(name)
+            and self.current_symbols.lookup(
+                mangle_priv_name(name, self.scopes)
+            ).is_free()
+        )
+
+    def is_local(self, name: str) -> bool:
+        return self.current_symbols.lookup(
+            mangle_priv_name(name, self.scopes)
+        ).is_local()
+
+    def scope_for(self, name: str) -> SymbolScope[TVar, TScopeData]:
+        if self.is_global(name):
+            return self.scopes[0]
+        elif self.is_nonlocal(name):
+            for scope in reversed(self.scopes):
+                lookup = scope.symbols.lookup(name)
+                if lookup.is_global() or lookup.is_free():
+                    return scope
+
+        return self.scopes[-1]
+
+    def __getitem__(self, name: str) -> TVar:
+        is_leaf = True
+        for scope in reversed(self.scopes):
+            if (is_leaf or not isinstance(scope.symbols, Class)) and name in scope:
+                return scope[name]
+            is_leaf = False
+        raise KeyError(f"{name} not found in scope")
+
+    def __contains__(self, name: str) -> bool:
+        is_leaf = True
+        for scope in reversed(self.scopes):
+            if (is_leaf or not isinstance(scope.symbols, Class)) and name in scope:
+                return True
+            is_leaf = False
+        return False
+
+    def __setitem__(self, name: str, value: TVar) -> None:
+        if self.is_global(name):
+            self.scopes[0][name] = value
+        elif self.is_nonlocal(name):
+            for scope in reversed(self.scopes[:-1]):
+                if name in scope:
+                    scope[name] = value
+                    break
+        else:
+            current_scope = self.scopes[-1]
+            current_scope[name] = value
+
+    def __delitem__(self, name: str) -> None:
+        if self.is_global(name):
+            if name in self.scopes[0]:
+                del self.scopes[0][name]
+        elif self.is_nonlocal(name):
+            for scope in reversed(self.scopes[:-1]):
+                if name in scope:
+                    del scope[name]
+                    break
+        else:
+            current_scope = self.scopes[-1]
+            if name in current_scope:
+                del current_scope[name]
+
+    def local_contains(self, name: str) -> bool:
+        return name in self.scopes[-1]
+
+    def shallow_copy(self) -> ScopeStack[TVar, TScopeData]:
+        # same underlying stack but a different dict
+        return ScopeStack(
+            *(d for d in self.scopes),
+            symbol_map=self.symbol_map,
+            scope_factory=self.scope_factory,
+        )
+
+    def with_scope(
+        self, scope: SymbolScope[TVar, TScopeData]
+    ) -> ScopeContextManager[TVar, TScopeData]:
+        return ScopeContextManager(self, scope)
+
+
+TAst = TypeVar("TAst", bound=AST)
+
+
+class AstRewriter(NodeVisitor):
+    """performs rewrites on the AST, but only produces new nodes, rather than
+    modifying nodes in place like ast.NodeTransformer."""
+
+    @staticmethod
+    def update_node(node: TAst, **replacement: object) -> TAst:
+        res = node
+        for name, val in replacement.items():
+            existing = getattr(res, name)
+            if existing == val:
+                continue
+
+            if node is res:
+                res = AstRewriter.clone_node(node)
+
+            setattr(res, name, val)
+        return res
+
+    @staticmethod
+    def clone_node(node: TAst) -> TAst:
+        attrs = []
+        for name in node._fields:
+            attr = getattr(node, name, None)
+            if isinstance(attr, list):
+                attr = list(attr)
+            attrs.append(attr)
+
+        new = type(node)(*attrs)
+        return copy_location(new, node)
+
+    def walk_list(self, old_values: list[TAst]) -> list[TAst]:
+        new_values = []
+        changed = False
+        for old_value in old_values:
+            if isinstance(old_value, AST):
+                new_value = self.visit(old_value)
+                changed |= new_value is not old_value
+                if new_value is None:
+                    continue
+                elif not isinstance(new_value, AST):
+                    new_values.extend(new_value)
+                    continue
+
+                new_values.append(new_value)
+            else:
+                new_values.append(old_value)
+        return new_values if changed else old_values
+
+    def generic_visit(self, node: TAst) -> TAst:
+        ret_node = node
+        for field, old_value in iter_fields(node):
+            if isinstance(old_value, list):
+                new_values = self.walk_list(old_value)
+                if new_values != old_value:
+                    if ret_node is node:
+                        ret_node = self.clone_node(node)
+                    setattr(ret_node, field, new_values)
+            elif isinstance(old_value, AST):
+                new_node = self.visit(old_value)
+                assert new_node is not None, (
+                    "can't remove AST nodes that aren't part of a list"
+                )
+                if new_node is not old_value:
+                    if ret_node is node:
+                        ret_node = self.clone_node(node)
+
+                    setattr(ret_node, field, new_node)
+
+        return ret_node
+
+
+def lineinfo(node: TAst, target: AST | None = None) -> TAst:
+    if not target:
+        # set lineno to -1 to indicate non-user code
+        node.lineno = -1
+        node.col_offset = -1
+        node.end_lineno = -1
+        node.end_col_offset = -1
+    else:
+        copy_location(node, target)
+    return node