codekingpro/portable-devtools

1import typing as t2 3from . import nodes4from .visitor import NodeVisitor5 6if t.TYPE_CHECKING:7    import typing_extensions as te8 9VAR_LOAD_PARAMETER = "param"10VAR_LOAD_RESOLVE = "resolve"11VAR_LOAD_ALIAS = "alias"12VAR_LOAD_UNDEFINED = "undefined"13 14 15def find_symbols(16    nodes: t.Iterable[nodes.Node], parent_symbols: t.Optional["Symbols"] = None17) -> "Symbols":18    sym = Symbols(parent=parent_symbols)19    visitor = FrameSymbolVisitor(sym)20    for node in nodes:21        visitor.visit(node)22    return sym23 24 25def symbols_for_node(26    node: nodes.Node, parent_symbols: t.Optional["Symbols"] = None27) -> "Symbols":28    sym = Symbols(parent=parent_symbols)29    sym.analyze_node(node)30    return sym31 32 33class Symbols:34    def __init__(35        self, parent: t.Optional["Symbols"] = None, level: t.Optional[int] = None36    ) -> None:37        if level is None:38            if parent is None:39                level = 040            else:41                level = parent.level + 142 43        self.level: int = level44        self.parent = parent45        self.refs: t.Dict[str, str] = {}46        self.loads: t.Dict[str, t.Any] = {}47        self.stores: t.Set[str] = set()48 49    def analyze_node(self, node: nodes.Node, **kwargs: t.Any) -> None:50        visitor = RootVisitor(self)51        visitor.visit(node, **kwargs)52 53    def _define_ref(54        self, name: str, load: t.Optional[t.Tuple[str, t.Optional[str]]] = None55    ) -> str:56        ident = f"l_{self.level}_{name}"57        self.refs[name] = ident58        if load is not None:59            self.loads[ident] = load60        return ident61 62    def find_load(self, target: str) -> t.Optional[t.Any]:63        if target in self.loads:64            return self.loads[target]65 66        if self.parent is not None:67            return self.parent.find_load(target)68 69        return None70 71    def find_ref(self, name: str) -> t.Optional[str]:72        if name in self.refs:73            return self.refs[name]74 75        if self.parent is not None:76            return self.parent.find_ref(name)77 78        return None79 80    def ref(self, name: str) -> str:81        rv = self.find_ref(name)82        if rv is None:83            raise AssertionError(84                "Tried to resolve a name to a reference that was"85                f" unknown to the frame ({name!r})"86            )87        return rv88 89    def copy(self) -> "te.Self":90        rv = object.__new__(self.__class__)91        rv.__dict__.update(self.__dict__)92        rv.refs = self.refs.copy()93        rv.loads = self.loads.copy()94        rv.stores = self.stores.copy()95        return rv96 97    def store(self, name: str) -> None:98        self.stores.add(name)99 100        # If we have not see the name referenced yet, we need to figure101        # out what to set it to.102        if name not in self.refs:103            # If there is a parent scope we check if the name has a104            # reference there.  If it does it means we might have to alias105            # to a variable there.106            if self.parent is not None:107                outer_ref = self.parent.find_ref(name)108                if outer_ref is not None:109                    self._define_ref(name, load=(VAR_LOAD_ALIAS, outer_ref))110                    return111 112            # Otherwise we can just set it to undefined.113            self._define_ref(name, load=(VAR_LOAD_UNDEFINED, None))114 115    def declare_parameter(self, name: str) -> str:116        self.stores.add(name)117        return self._define_ref(name, load=(VAR_LOAD_PARAMETER, None))118 119    def load(self, name: str) -> None:120        if self.find_ref(name) is None:121            self._define_ref(name, load=(VAR_LOAD_RESOLVE, name))122 123    def branch_update(self, branch_symbols: t.Sequence["Symbols"]) -> None:124        stores: t.Set[str] = set()125 126        for branch in branch_symbols:127            stores.update(branch.stores)128 129        stores.difference_update(self.stores)130 131        for sym in branch_symbols:132            self.refs.update(sym.refs)133            self.loads.update(sym.loads)134            self.stores.update(sym.stores)135 136        for name in stores:137            target = self.find_ref(name)138            assert target is not None, "should not happen"139 140            if self.parent is not None:141                outer_target = self.parent.find_ref(name)142                if outer_target is not None:143                    self.loads[target] = (VAR_LOAD_ALIAS, outer_target)144                    continue145            self.loads[target] = (VAR_LOAD_RESOLVE, name)146 147    def dump_stores(self) -> t.Dict[str, str]:148        rv: t.Dict[str, str] = {}149        node: t.Optional[Symbols] = self150 151        while node is not None:152            for name in sorted(node.stores):153                if name not in rv:154                    rv[name] = self.find_ref(name)  # type: ignore155 156            node = node.parent157 158        return rv159 160    def dump_param_targets(self) -> t.Set[str]:161        rv = set()162        node: t.Optional[Symbols] = self163 164        while node is not None:165            for target, (instr, _) in self.loads.items():166                if instr == VAR_LOAD_PARAMETER:167                    rv.add(target)168 169            node = node.parent170 171        return rv172 173 174class RootVisitor(NodeVisitor):175    def __init__(self, symbols: "Symbols") -> None:176        self.sym_visitor = FrameSymbolVisitor(symbols)177 178    def _simple_visit(self, node: nodes.Node, **kwargs: t.Any) -> None:179        for child in node.iter_child_nodes():180            self.sym_visitor.visit(child)181 182    visit_Template = _simple_visit183    visit_Block = _simple_visit184    visit_Macro = _simple_visit185    visit_FilterBlock = _simple_visit186    visit_Scope = _simple_visit187    visit_If = _simple_visit188    visit_ScopedEvalContextModifier = _simple_visit189 190    def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:191        for child in node.body:192            self.sym_visitor.visit(child)193 194    def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:195        for child in node.iter_child_nodes(exclude=("call",)):196            self.sym_visitor.visit(child)197 198    def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:199        for child in node.body:200            self.sym_visitor.visit(child)201 202    def visit_For(203        self, node: nodes.For, for_branch: str = "body", **kwargs: t.Any204    ) -> None:205        if for_branch == "body":206            self.sym_visitor.visit(node.target, store_as_param=True)207            branch = node.body208        elif for_branch == "else":209            branch = node.else_210        elif for_branch == "test":211            self.sym_visitor.visit(node.target, store_as_param=True)212            if node.test is not None:213                self.sym_visitor.visit(node.test)214            return215        else:216            raise RuntimeError("Unknown for branch")217 218        if branch:219            for item in branch:220                self.sym_visitor.visit(item)221 222    def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:223        for target in node.targets:224            self.sym_visitor.visit(target)225        for child in node.body:226            self.sym_visitor.visit(child)227 228    def generic_visit(self, node: nodes.Node, *args: t.Any, **kwargs: t.Any) -> None:229        raise NotImplementedError(f"Cannot find symbols for {type(node).__name__!r}")230 231 232class FrameSymbolVisitor(NodeVisitor):233    """A visitor for `Frame.inspect`."""234 235    def __init__(self, symbols: "Symbols") -> None:236        self.symbols = symbols237 238    def visit_Name(239        self, node: nodes.Name, store_as_param: bool = False, **kwargs: t.Any240    ) -> None:241        """All assignments to names go through this function."""242        if store_as_param or node.ctx == "param":243            self.symbols.declare_parameter(node.name)244        elif node.ctx == "store":245            self.symbols.store(node.name)246        elif node.ctx == "load":247            self.symbols.load(node.name)248 249    def visit_NSRef(self, node: nodes.NSRef, **kwargs: t.Any) -> None:250        self.symbols.load(node.name)251 252    def visit_If(self, node: nodes.If, **kwargs: t.Any) -> None:253        self.visit(node.test, **kwargs)254        original_symbols = self.symbols255 256        def inner_visit(nodes: t.Iterable[nodes.Node]) -> "Symbols":257            self.symbols = rv = original_symbols.copy()258 259            for subnode in nodes:260                self.visit(subnode, **kwargs)261 262            self.symbols = original_symbols263            return rv264 265        body_symbols = inner_visit(node.body)266        elif_symbols = inner_visit(node.elif_)267        else_symbols = inner_visit(node.else_ or ())268        self.symbols.branch_update([body_symbols, elif_symbols, else_symbols])269 270    def visit_Macro(self, node: nodes.Macro, **kwargs: t.Any) -> None:271        self.symbols.store(node.name)272 273    def visit_Import(self, node: nodes.Import, **kwargs: t.Any) -> None:274        self.generic_visit(node, **kwargs)275        self.symbols.store(node.target)276 277    def visit_FromImport(self, node: nodes.FromImport, **kwargs: t.Any) -> None:278        self.generic_visit(node, **kwargs)279 280        for name in node.names:281            if isinstance(name, tuple):282                self.symbols.store(name[1])283            else:284                self.symbols.store(name)285 286    def visit_Assign(self, node: nodes.Assign, **kwargs: t.Any) -> None:287        """Visit assignments in the correct order."""288        self.visit(node.node, **kwargs)289        self.visit(node.target, **kwargs)290 291    def visit_For(self, node: nodes.For, **kwargs: t.Any) -> None:292        """Visiting stops at for blocks.  However the block sequence293        is visited as part of the outer scope.294        """295        self.visit(node.iter, **kwargs)296 297    def visit_CallBlock(self, node: nodes.CallBlock, **kwargs: t.Any) -> None:298        self.visit(node.call, **kwargs)299 300    def visit_FilterBlock(self, node: nodes.FilterBlock, **kwargs: t.Any) -> None:301        self.visit(node.filter, **kwargs)302 303    def visit_With(self, node: nodes.With, **kwargs: t.Any) -> None:304        for target in node.values:305            self.visit(target)306 307    def visit_AssignBlock(self, node: nodes.AssignBlock, **kwargs: t.Any) -> None:308        """Stop visiting at block assigns."""309        self.visit(node.target, **kwargs)310 311    def visit_Scope(self, node: nodes.Scope, **kwargs: t.Any) -> None:312        """Stop visiting at scopes."""313 314    def visit_Block(self, node: nodes.Block, **kwargs: t.Any) -> None:315        """Stop visiting at blocks."""316 317    def visit_OverlayScope(self, node: nodes.OverlayScope, **kwargs: t.Any) -> None:318        """Do not visit into overlay scopes."""319