Python astor.to_source() Examples

def check_compiles(dsl_expr, code):
    types = ENV.__types__

    expr, functions = to_expr(dsl_expr)
    env = fn_types(functions)
    env.update(types['__root__'].__field_types__)

    expr = check(expr, types, env)

    # test eval
    lambda_expr = ExpressionCompiler.compile_lambda_expr(expr)
    eval(compile(lambda_expr, '<expr>', 'eval'))

    # test compile
    py_expr = ExpressionCompiler.compile_expr(expr)

    first = astor.to_source(py_expr).strip()
    second = dedent(code).strip()
    if first != second:
        msg = ('Compiled code is not equal:\n\n{}'
               .format('\n'.join(difflib.ndiff(first.splitlines(),
                                               second.splitlines()))))
        raise AssertionError(msg) 

def transform_concat(code: str, *args, **kwargs) -> Tuple[str, bool]:
    tree = ast.parse(f"({code})")

    ft = ConcatTransformer()
    ft.visit(tree)
    il = FstrInliner()
    il.visit(tree)

    new_code = astor.to_source(tree)
    if new_code[-1] == "\n":
        new_code = new_code[:-1]

    new_code = new_code.replace("\n", "\\n")
    if new_code[:4] == 'f"""':
        new_code = set_quote_type(new_code, QuoteTypes.double)

    return new_code, ft.counter > 0 

def ast_formatted_value(
    val, fmt_str: str = None, conversion=None
) -> ast.FormattedValue:
    if isinstance(val, ast.FormattedValue):
        return val

    if astor.to_source(val)[0] == "{":
        raise FlyntException(
            "values starting with '{' are better left not transformed."
        )

    if fmt_str:
        format_spec = ast.JoinedStr([ast_string_node(fmt_str.replace(":", ""))])
    else:
        format_spec = None

    if conversion is None:
        conversion = -1
    else:
        conversion = ord(conversion.replace("!", ""))
    return ast.FormattedValue(value=val, conversion=conversion, format_spec=format_spec) 

def ast_formatted_value(
    val, fmt_str: str = None, conversion=None
) -> ast.FormattedValue:

    if astor.to_source(val)[0] == "{":
        raise FlyntException(
            "values starting with '{' are better left not transformed."
        )

    if fmt_str:
        format_spec = ast.JoinedStr([ast_string_node(fmt_str.replace(":", ""))])
    else:
        format_spec = None

    if conversion is None:
        conversion = -1
    else:
        conversion = ord(conversion.replace("!", ""))
    return ast.FormattedValue(value=val, conversion=conversion, format_spec=format_spec) 

def parse_functions(code):
    """Parse all the global functions present in the input code.

    Parse all the ast nodes ast.FunctionDef that are global functions in the
    source code. These also include function that are defined inside other
    Python statements, like `try`. ast.ClassDef nodes are skipped from the
    parsing so that class functions are ignored.

    Args:
        code (str): Multiline string representing Python code

    Returns (dict): A dictionary [fn_name] -> function_source
    """
    fns = dict()
    tree = ast.parse(code)
    for block in tree.body:
        for node in walk(block,
                         stop_at=(ast.FunctionDef,),
                         ignore=(ast.ClassDef,)):
            if isinstance(node, (ast.FunctionDef,)):
                fn_name = node.name
                fns[fn_name] = astor.to_source(node)
    return fns 

def generic_visit(self, node):
        """Catch-all for nodes that slip through

        Basically everything I haven't gotten around to writing a custom
        annotator for gets caught here and wrapped in an annotation. Currently
        the one that jumps to mind are context managers.

        This is necessary because some nodes we recursively call `self.visit()`
        on and we may run into expressions that we have not written a node
        tranformer for.

        """
        # try:
        #     c = astor.to_source(node)
        #     print(c.strip())
        #     print(getattr(node, 'lineno', -1))
        #     print()
        # except:
        #     pass
        # return super().generic_visit(node)
        if not self.target_node and (getattr(node, 'lineno', -1) == self.lineno):
            self.target_node = node
        return super().generic_visit(node) 

def make_annotation(node=None, buffer='outside', content=None, cell_type='code', lineno=None):
    """Return a ast.Expr that looks like

    ```
    __cell__('make-cell', [content, buffer, cell_type])
    ```

    """
    content = astor.to_source(node).strip() if node else content
    lineno = str(node.lineno) if hasattr(node, 'lineno') else str(-1) if not lineno else str(lineno)
    call = ast.Call(
        func=ast.Name(id='__cell__', ctx=ast.Load()),
        args=[
            ast.Str(s=content),
            ast.Str(s=f'{buffer}'),
            ast.Str(s=cell_type),
            ast.Str(s=lineno),
        ],
        keywords=[]
    )
    return ast.Expr(call) 

def generate(module_name, code):
    """Generate search space.
    Return a serializable search space object.
    module_name: name of the module (str)
    code: user code (str)
    """
    try:
        ast_tree = ast.parse(code)
    except Exception:
        raise RuntimeError('Bad Python code')

    visitor = SearchSpaceGenerator(module_name)
    try:
        visitor.visit(ast_tree)
    except AssertionError as exc:
        raise RuntimeError('%d: %s' % (visitor.last_line, exc.args[0]))
    return visitor.search_space, astor.to_source(ast_tree) 

def parse_nni_variable(code):
    """Parse `nni.variable` expression.
    Return the name argument and AST node of annotated expression.
    code: annotation string
    """
    name, call = parse_annotation_function(code, 'variable')

    assert len(call.args) == 1, 'nni.variable contains more than one arguments'
    arg = call.args[0]
    assert type(arg) is ast.Call, 'Value of nni.variable is not a function call'
    assert type(arg.func) is ast.Attribute, 'nni.variable value is not a NNI function'
    assert type(arg.func.value) is ast.Name, 'nni.variable value is not a NNI function'
    assert arg.func.value.id == 'nni', 'nni.variable value is not a NNI function'

    name_str = astor.to_source(name).strip()
    keyword_arg = ast.keyword(arg='name', value=ast.Str(s=name_str))
    arg.keywords.append(keyword_arg)
    if arg.func.attr == 'choice':
        convert_args_to_dict(arg)

    return name, arg 

def test_filter_class(self):
        code = """

class Foo:
    def bar():
        pass

class Bar:
    def biz():
        pass

"""
        tree = ast.parse(code)
        out = codebook.node_transformers.ClassFinder(class_name='Foo').visit(tree)
        c = astor.to_source(out)
        self.assertIn('Foo', c)
        self.assertNotIn('Bar', c) 

def convert_args_to_dict(call, with_lambda=False):
    """Convert all args to a dict such that every key and value in the dict is the same as the value of the arg.
    Return the AST Call node with only one arg that is the dictionary
    """
    keys, values = list(), list()
    for arg in call.args:
        if type(arg) in [ast.Str, ast.Num]:
            arg_value = arg
        else:
            # if arg is not a string or a number, we use its source code as the key
            arg_value = astor.to_source(arg).strip('\n"')
            arg_value = ast.Str(str(arg_value))
        arg = make_lambda(arg) if with_lambda else arg
        keys.append(arg_value)
        values.append(arg)
    del call.args[:]
    call.args.append(ast.Dict(keys=keys, values=values))

    return call 

def unpack_annotations(annotations):
    """Extract the information out of a bunch of annotations

    >>> code = '''
    ...
    ... __cell__('`foo.baz`', 'foo.baz', '1', '9')
    ... __cell__('Arguments', 'foo.baz', '1', '-1')
    ... __cell__('Body', 'foo.baz', '1', '-1')
    ... __cell__('pass', 'foo.baz', 'code', '10')
    ...
    ... '''
    ...
    >>> annotations = ast.parse(code)

    """
    info = []
    m = astor.to_source(annotations)
    print(m)
    for annotation in annotations.body:
        content, namespace, cell_type, lineno = [arg.s for arg in annotation.value.args]
        info.append([content, namespace, cell_type, int(lineno)])
    return info 

def promote_loop(*region):
    """

    >>> code = '''
    ...
    ... for i in range(5):
    ...     for j in range(5):
    ...         k = i = j
    ...         print(k)
    ...
    ... '''

    """
    code, namespace = region
    tree = ast.parse(code)
    m = FirstPassForSimple(buffer=namespace).visit(tree)
    return astor.to_source(m) 

def annotate_toplevel(*region):
    """

    >>> code = '''
    ...
    ... for i in range(5):
    ...     for j in range(5):
    ...         k = i = j
    ...         print(k)
    ...
    ... '''
    >>> namespace = 'foo'
    >>> region = [code, namespace]

    """
    code, namespace = region
    tree = ast.parse(code)
    annotated_tree = Annotator(buffer=namespace).visit(tree)
    return astor.to_source(annotated_tree) 

def convert_args_to_dict(call, with_lambda=False):
    """Convert all args to a dict such that every key and value in the dict is the same as the value of the arg.
    Return the AST Call node with only one arg that is the dictionary
    """
    keys, values = list(), list()
    for arg in call.args:
        if type(arg) in [ast.Str, ast.Num]:
            arg_value = arg
        else:
            # if arg is not a string or a number, we use its source code as the key
            arg_value = astor.to_source(arg).strip('\n"')
            arg_value = ast.Str(str(arg_value))
        arg = make_lambda(arg) if with_lambda else arg
        keys.append(arg_value)
        values.append(arg)
    del call.args[:]
    call.args.append(ast.Dict(keys=keys, values=values))

    return call 

def parse(code, para, module):
    """Annotate user code.
    Return annotated code (str) if annotation detected; return None if not.
    code: original user code (str)
    """
    global para_cfg
    global prefix_name
    para_cfg = para
    prefix_name = module
    try:
        ast_tree = ast.parse(code)
    except Exception:
        raise RuntimeError('Bad Python code')

    transformer = Transformer()
    try:
        transformer.visit(ast_tree)
    except AssertionError as exc:
        raise RuntimeError('%d: %s' % (ast_tree.last_line, exc.args[0]))

    if not transformer.annotated:
        return None

    return astor.to_source(ast_tree) 

def get_source(self):
        """
        Get a string containing the Python source code corresponding to the
        statements in the block.

        Returns:
            A string containing the source code of the statements.
        """
        src = ""
        for statement in self.statements:
            if type(statement) in [ast.If, ast.For, ast.While]:
                src += (astor.to_source(statement)).split('\n')[0] + "\n"
            elif type(statement) == ast.FunctionDef or\
                 type(statement) == ast.AsyncFunctionDef:
                src += (astor.to_source(statement)).split('\n')[0] + "...\n"
            else:
                src += astor.to_source(statement)
        return src 

def exposed_astor_roundtrip_parser_functions():
	'''
	Shove the feed-functions through the astor "round-trip"
	facility.

	Mostly, this homogenizes the indentation, and reformats the function.
	'''

	with db.session_context() as sess:
		res = sess.query(db.RssFeedEntry) \
			.all()

		for row in res:
			func = row.get_func()
			_ast = row._get_ast()
			src = astor.to_source(_ast, indent_with="	", pretty_source=better_pretty_source)

			if src.strip() != row.func.strip():
				try:
					rfdb.str_to_function(src, "testing_compile")
					print("Compiled OK")
					row.func = src
				except Exception:
					print("Compilation failed?")
		sess.commit() 

def exposed_import_feed_parse_funcs():
	'''
	Import the feed parsing functions into the database.
	'''


	# parse_map = WebMirror.OutputFilters.rss.FeedDataParser.RSS_PARSE_FUNCTION_MAP
	# for key, func in parse_map.items():
	# 	func_str = astor.to_source(astor.code_to_ast(func), indent_with="	")
	# 	update_func(sess, key, func_str)

	name_map = WebMirror.OutputFilters.util.feedNameLut.mapper

	with common.database.session_context() as sess:
		for key, val in name_map.items():
			add_name(sess, key, val) 

def run():
        asdl_text = open('asdl/lang/py/py_asdl.txt').read()
        grammar = ASDLGrammar.from_text(asdl_text)

        annot_file = 'data/django/all.anno'
        code_file = 'data/django/all.code'

        transition_system = PythonTransitionSystem(grammar)

        for idx, (src_query, tgt_code) in enumerate(zip(open(annot_file), open(code_file))):
            src_query = src_query.strip()
            tgt_code = tgt_code.strip()

            query_tokens, tgt_canonical_code, str_map = Django.canonicalize_example(src_query, tgt_code)
            python_ast = ast.parse(tgt_canonical_code).body[0]
            gold_source = astor.to_source(python_ast)
            tgt_ast = python_ast_to_asdl_ast(python_ast, grammar)
            tgt_actions = transition_system.get_actions(tgt_ast)

            # sanity check
            hyp = Hypothesis()
            hyp2 = Hypothesis()
            for action in tgt_actions:
                assert action.__class__ in transition_system.get_valid_continuation_types(hyp)
                if isinstance(action, ApplyRuleAction):
                    assert action.production in transition_system.get_valid_continuating_productions(hyp)
                hyp = hyp.clone_and_apply_action(action)
                hyp2.apply_action(action)

            src_from_hyp = astor.to_source(asdl_ast_to_python_ast(hyp.tree, grammar))
            assert src_from_hyp == gold_source
            assert hyp.tree == hyp2.tree and hyp.tree is not hyp2.tree

            print(idx) 

def node_to_string(node: ast.AST) -> str:
    """Returns the source code by doing ``ast`` to string convert."""
    return astor.to_source(node).strip() 

def main():
    parser = argparse.ArgumentParser(description="Kappa compiler")
    parser.add_argument("--auto-pause", action="store_true", help="Automatically insert pause points.")
    args = parser.parse_args()

    source = sys.stdin.read()
    mod = ast.parse(source)
    transformed = transform(mod, auto_pause=args.auto_pause)
    print(astor.to_source(transformed))  # Print out resulting AST as code. 

def canonicalize_raw_django_oneliner(code):
        # use the astor-style code
        code = Django.canonicalize_code(code)
        py_ast = ast.parse(code).body[0]
        code = astor.to_source(py_ast).strip()

        return code 

def dump_class(self):
		indent = "	"
		return astor.to_source(self.__to_module(), indent_with=indent) 

def is_hyp_correct(self, example, hyp):
        ref_code = example.tgt_code
        ref_py_ast = ast.parse(ref_code).body[0]
        ref_reformatted_code = astor.to_source(ref_py_ast).strip()

        ref_code_tokens = self.transition_system.tokenize_code(ref_reformatted_code)
        hyp_code_tokens = self.transition_system.tokenize_code(hyp.code)

        return ref_code_tokens == hyp_code_tokens 

def decanonicalize_code(code, slot_map):
    for slot_name, slot_val in slot_map.items():
        if is_enumerable_str(slot_name):
            code = code.replace(slot_name, slot_val['value'])

    slot2string = {x[0]: x[1]['value'] for x in list(slot_map.items())}
    py_ast = ast.parse(code)
    replace_identifiers_in_ast(py_ast, slot2string)
    raw_code = astor.to_source(py_ast).strip()
    # for slot_name, slot_info in slot_map.items():
    #     raw_code = raw_code.replace(slot_name, slot_info['value'])

    return raw_code 

def post_process_hypothesis(self, hyp, meta_info, utterance=None):
        """traverse the AST and replace slot ids with original strings"""
        hyp_ast = asdl_ast_to_python_ast(hyp.tree, self.transition_system.grammar)
        code_from_hyp = astor.to_source(hyp_ast).strip()
        hyp.code = decanonicalize_code(code_from_hyp, meta_info) 

def is_hyp_correct(self, example, hyp):
        ref_code = example.tgt_code
        ref_py_ast = ast.parse(ref_code)
        ref_reformatted_code = astor.to_source(ref_py_ast).strip()

        ref_code_tokens = self.transition_system.tokenize_code(ref_reformatted_code)
        hyp_code_tokens = self.transition_system.tokenize_code(hyp.code)

        return ref_code_tokens == hyp_code_tokens 

def ast_to_surface_code(self, asdl_ast):
        py_ast = asdl_ast_to_python_ast(asdl_ast, self.grammar)
        code = astor.to_source(py_ast).strip()

        if code.endswith(':'):
            code += ' pass'

        # first make sure the hypothesis code is parsable by `ast`
        # sometimes, the parser generates syntactically invalid surface code:
        # e.g., slot_0.1()
        # ast.parse(code)

        return code 

def ast_to_surface_code(self, asdl_ast):
        py_ast = asdl_ast_to_python_ast(asdl_ast, self.grammar)
        code = astor.to_source(py_ast).strip()

        return code