Python ast.NodeTransformer() Examples

def generic_visit(self, node):
        super(NodeTransformer, self).generic_visit(node)
        if hasattr(node, 'body') and type(node.body) is list:
            returns = [i for i, child in enumerate(node.body) if type(child) is ast.Return]
            if len(returns) > 0:
                for wait in self.get_waits():
                    node.body.insert(returns[0], wait)
            inserts = []
            for i, child in enumerate(node.body):
                if type(child) is ast.Expr and self.is_concurrent_call(child.value):
                    self.encounter_call(child.value)
                elif self.is_valid_assignment(child):
                    call = child.value
                    self.encounter_call(call)
                    name = child.targets[0].value
                    self.arguments.add(SchedulerRewriter.top_level_name(name))
                    index = child.targets[0].slice.value
                    call.func = ast.Attribute(call.func, 'assign', ast.Load())
                    call.args = [ast.Tuple([name, index], ast.Load())] + call.args
                    node.body[i] = ast.Expr(call)
                elif self.references_arg(child):
                    inserts.insert(0, i)
            for index in inserts:
                for wait in self.get_waits():
                    node.body.insert(index, wait) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def __init__(self, script, aux_functions, aliases, *args, **kwargs):
        ast.NodeTransformer.__init__(self, *args, **kwargs)
        self.aux_functions = aux_functions
        self.aliases = aliases
        self.script = script 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def filterstr_to_filterfunc(filter_str: str, item_type: type):
    """Takes an --post-filter=... or --storyitem-filter=... filter
     specification and makes a filter_func Callable out of it."""

    # The filter_str is parsed, then all names occurring in its AST are replaced by loads to post.<name>. A
    # function Post->bool is returned which evaluates the filter with the post as 'post' in its namespace.

    class TransformFilterAst(ast.NodeTransformer):
        def visit_Name(self, node: ast.Name):
            # pylint:disable=no-self-use
            if not isinstance(node.ctx, ast.Load):
                raise InvalidArgumentException("Invalid filter: Modifying variables ({}) not allowed.".format(node.id))
            if node.id == "datetime":
                return node
            if not hasattr(item_type, node.id):
                raise InvalidArgumentException("Invalid filter: {} not a {} attribute.".format(node.id,
                                                                                               item_type.__name__))
            new_node = ast.Attribute(ast.copy_location(ast.Name('item', ast.Load()), node), node.id,
                                     ast.copy_location(ast.Load(), node))
            return ast.copy_location(new_node, node)

    input_filename = '<command line filter parameter>'
    compiled_filter = compile(TransformFilterAst().visit(ast.parse(filter_str, filename=input_filename, mode='eval')),
                              filename=input_filename, mode='eval')

    def filterfunc(item) -> bool:
        # pylint:disable=eval-used
        return bool(eval(compiled_filter, {'item': item, 'datetime': datetime.datetime}))

    return filterfunc 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def _generate_translators(to):

    class Translator(ast.NodeTransformer):

        def _visit(self, node):
            if isinstance(node, list):
                return [self._visit(n) for n in node]
            elif isinstance(node, ast.AST):
                return self.visit(node)
            else:
                return node

        def generic_visit(self, node):
            cls = type(node).__name__
            # handle nodes that are not part of the AST
            if not hasattr(to, cls):
                return
            new_node = getattr(to, cls)()
            for field in node._fields:
                setattr(new_node, field, self._visit(getattr(node, field)))
            for attr in getattr(node, '_attributes'):
                if hasattr(node, attr):
                    setattr(new_node, attr, getattr(node, attr))
            return new_node

    return Translator 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def visit(self, node):
        """Ensure statement only contains allowed nodes."""
        if not isinstance(node, self.ALLOWED):
            raise SyntaxError('Not allowed in environment markers.\n%s\n%s' %
                               (self.statement,
                               (' ' * node.col_offset) + '^'))
        return ast.NodeTransformer.visit(self, node) 

def _insert_expression_markers(self, node):
        """
        TODO: this docstring is outdated
        each expression e gets wrapped like this:
            _after(_before(_loc, _node_is_zoomable), e, _node_role, _parent_range)
        where
            _after is function that gives the resulting value
            _before is function that signals the beginning of evaluation of e
            _loc gives the code range of e
            _node_is_zoomable indicates whether this node has subexpressions
            _node_role is either 'last_call_arg', 'last_op_arg', 'first_or_arg', 
                                 'first_and_arg', 'function' or None
        """
        tracer = self

        class ExpressionVisitor(ast.NodeTransformer):
            def generic_visit(self, node):
                if isinstance(node, _ast.expr):
                    if isinstance(node, ast.Starred):
                        # keep this node as is, but instrument its children
                        return ast.NodeTransformer.generic_visit(self, node)
                    elif tracer._should_instrument_as_expression(node):
                        # before marker
                        before_marker = tracer._create_simple_marker_call(
                            node, BEFORE_EXPRESSION_MARKER
                        )
                        ast.copy_location(before_marker, node)

                        if "ignore_children" in node.tags:
                            transformed_node = node
                        else:
                            transformed_node = ast.NodeTransformer.generic_visit(self, node)

                        # after marker
                        after_marker = ast.Call(
                            func=ast.Name(id=AFTER_EXPRESSION_MARKER, ctx=ast.Load()),
                            args=[before_marker, transformed_node],
                            keywords=[],
                        )
                        ast.copy_location(after_marker, node)
                        ast.fix_missing_locations(after_marker)
                        # further transformations may query original node location from after marker
                        if hasattr(node, "end_lineno"):
                            after_marker.end_lineno = node.end_lineno
                            after_marker.end_col_offset = node.end_col_offset

                        return after_marker
                    else:
                        # This expression (and its children) should be ignored
                        return node
                else:
                    # Descend into statements
                    return ast.NodeTransformer.generic_visit(self, node)

        return ExpressionVisitor().visit(node) 

def run_script(self, *, file_path=None, stdout=None) -> None:
        import ast
        import contextlib
        import os
        with open(file_path) as f:
            script = f.read()
        script_name = os.path.basename(file_path)
        script_ast = ast.parse(script, script_name, 'exec')

        class AddCallFunctionNodeTransformer(ast.NodeTransformer):
            def __init__(self, func_id, arg_id):
                self.__func_id = func_id
                self.__arg_id = arg_id

            def visit_Module(self, node):
                name_expr = ast.Name(id=self.__func_id, ctx=ast.Load())
                arg_expr = ast.Name(id=self.__arg_id, ctx=ast.Load())
                call_expr = ast.Expr(value=ast.Call(func=name_expr, args=[arg_expr], keywords=[]))
                new_node = copy.deepcopy(node)
                new_node.body.append(call_expr)
                ast.fix_missing_locations(new_node)
                return new_node

        # if script_main exists, add a node to call it
        for node in script_ast.body:
            if getattr(node, "name", None) == "script_main":
                script_ast = AddCallFunctionNodeTransformer('script_main', 'api_broker').visit(script_ast)

        compiled = compile(script_ast, script_name, 'exec')

        if not stdout:
            print_fn = print
            class StdoutCatcher:
                def __init__(self):
                    pass
                def write(self, stuff):
                    print_fn(stuff.rstrip())
                def flush(self):
                    pass
            stdout = StdoutCatcher()

        class APIBroker:
            def get_api(self, version, ui_version=None):
                ui_version = ui_version if ui_version else "~1.0"
                return PlugInManager.api_broker_fn(version, ui_version)
            def get_ui(self, version):
                actual_version = "1.0.0"
                if Utility.compare_versions(version, actual_version) > 0:
                    raise NotImplementedError("API requested version %s is greater than %s." % (version, actual_version))
                return Declarative.DeclarativeUI()

        with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stdout):
            g = dict()
            g["api_broker"] = APIBroker()
            g["print"] = stdout.write
            exec(compiled, g)