Python ast.FunctionDef() Examples

def get_line_number_from_function(fn):
    # type: (Union[ast.FunctionDef, ast.AsyncFunctionDef]) -> int
    """Get the line number for the end of the function signature.

    The function signature can be farther down when the parameter
    list is split across multiple lines.

    Args:
        fn: The function from which we are getting the line number.

    Returns:
        The line number for the start of the docstring for this
        function.

    """
    line_number = fn.lineno
    if hasattr(fn, 'args') and fn.args.args:
        last_arg = fn.args.args[-1]
        line_number = last_arg.lineno
    return line_number 

def find_globals(g, tree):
    """Uses AST to find globals in an ast tree"""
    for child in tree:
        if hasattr(child, 'body') and isinstance(child.body, list):
            find_globals(g, child.body)
        elif isinstance(child, (ast.FunctionDef, ast.ClassDef)):
            g.add(child.name)
            continue
        elif isinstance(child, ast.Assign):
            try:
                g.add(child.targets[0].id)
            except (IndexError, AttributeError):
                pass
        elif isinstance(child, ast.Import):
            g.add(child.names[0].name)
        elif isinstance(child, ast.ImportFrom):
            for name in child.names:
                g_name = name.asname or name.name
                if g_name == '*':
                    continue
                g.add(g_name) 

def indent_not_multiple_of_tab_size(project_folder, tab_size, *args, **kwargs):
    """
        Since there are cases for which col_offset is computed incorrectly,
        this validator must be nothing more than a simple warning.

        It compliments the pep8 validator which tends to fail in cases when
        the indent is incorrect.
    """
    node_types_to_validate = (ast.For, ast.If, ast.FunctionDef, ast.With)
    for parsed_file in project_folder.get_parsed_py_files():
        lines_offsets = file_helpers.get_line_offsets(parsed_file.content)
        for node in ast.walk(parsed_file.ast_tree):
            if not ast_helpers.is_node_offset_fine(
                node,
                lines_offsets,
                node_types_to_validate,
                tab_size,
            ):
                return parsed_file.get_name_with_line(node.lineno) 

def _fine_property_definition(self, property_name):
        """Find the lines in the source code that contain this property's name and definition.

        This function can find both attribute assignments as well as methods/functions.

        Args:
            property_name (str): the name of the property to look up in the template definition

        Returns:
            tuple: line numbers for the start and end of the attribute definition
        """
        for node in ast.walk(ast.parse(self._source)):
            if isinstance(node, ast.Assign) and node.targets[0].id == property_name:
                return node.targets[0].lineno - 1, self._get_node_line_end(node)
            elif isinstance(node, ast.FunctionDef) and node.name == property_name:
                return node.lineno - 1, self._get_node_line_end(node)
        raise ValueError('The requested node could not be found.') 

def CONTINUE(self, node):
        # Walk the tree up until we see a loop (OK), a function or class
        # definition (not OK), for 'continue', a finally block (not OK), or
        # the top module scope (not OK)
        n = node
        while hasattr(n, 'parent'):
            n, n_child = n.parent, n
            if isinstance(n, LOOP_TYPES):
                # Doesn't apply unless it's in the loop itself
                if n_child not in n.orelse:
                    return
            if isinstance(n, (ast.FunctionDef, ast.ClassDef)):
                break
            # Handle Try/TryFinally difference in Python < and >= 3.3
            if hasattr(n, 'finalbody') and isinstance(node, ast.Continue):
                if n_child in n.finalbody:
                    self.report(messages.ContinueInFinally, node)
                    return
        if isinstance(node, ast.Continue):
            self.report(messages.ContinueOutsideLoop, node)
        else:  # ast.Break
            self.report(messages.BreakOutsideLoop, node) 

def get_comments_str(file_name):
    with open(file_name) as fd:
        file_contents = fd.read()
    module = ast.parse(file_contents)

    function_definitions = [node for node in module.body if isinstance(node, ast.FunctionDef)]

    doc = get_func_comments(function_definitions)

    class_definitions = [node for node in module.body if isinstance(node, ast.ClassDef)]
    for class_def in class_definitions:
        temp_str = to_md(parse_func_string(ast.get_docstring(class_def)))

        # excludes private methods (start with '_')
        method_definitions = [
            node
            for node in class_def.body
            if isinstance(node, ast.FunctionDef) and (node.name[0] != '_' or node.name[:2] == '__')
        ]

        temp_str += get_func_comments(method_definitions)
        doc += '## class ' + class_def.name + '\n' + temp_str
    return doc 

def from_ast(cls, tree: ast.Module) -> List['Func']:
        funcs = []
        for expr in tree.body:
            if not isinstance(expr, ast.FunctionDef):
                continue
            contracts = []
            for category, args in get_contracts(expr.decorator_list):
                contract = Contract(
                    args=args,
                    category=Category(category),
                    func_args=expr.args,
                )
                contracts.append(contract)
            funcs.append(cls(
                name=expr.name,
                args=expr.args,
                body=expr.body,
                contracts=contracts,
                line=expr.lineno,
                col=expr.col_offset,
            ))
        return funcs 

def _get_decorator_names(fun):  # type: (Union[ast.FunctionDef, ast.AsyncFunctionDef]) -> List[str]  # noqa: E501
    """Get decorator names from the function.

    Args:
        fun: The function whose decorators we are getting.

    Returns:
        The names of the decorators. Does not include setters and
        getters.

    """
    ret = list()
    for decorator in fun.decorator_list:
        # Attributes (setters and getters) won't have an id.
        if hasattr(decorator, 'id'):
            ret.append(getattr(decorator, 'id'))
    return ret 

def _has_return(fun):
    # type: (Union[ast.FunctionDef, ast.AsyncFunctionDef]) -> bool
    """Return true if the function has a fruitful return.

    Args:
        fun: A function node to check.

    Returns:
        True if there is a fruitful return, otherwise False.

    """
    def skip(f):
        return f != fun and isinstance(f, FunctionDef)

    for node in _walk(fun, skip):
        if isinstance(node, ast.Return) and node.value is not None:
            return True
    return False 

def _computeInitMethodAstOrNone(pyClassObject):
    pyAst = pyAstFor(pyClassObject)

    assert len(pyAst.body) == 1

    classDef = pyAst.body[0]

    assert isinstance(classDef, ast.ClassDef)

    tr = None

    # recall ClassDef = (identifier name, expr* bases, stmt* body, expr* decorator_list)
    # we're taking the *last* __init__ here
    for stmt in classDef.body:
        if isinstance(stmt, ast.FunctionDef) and stmt.name == "__init__":
            tr = stmt

    return tr 

def visit_ClassDef(self, classdef):
        """Check the line number of each of the methods

        >>> self = DefunFinder(func_name='bar', lineno=4)
        >>> code = '''
        ...
        ... class Foo:
        ...     def bar():
        ...         \"\"\"function\"\"\"
        ...         pass
        ...     def biz():
        ...         \"\"\"function\"\"\"
        ...         pass
        ...
        ... '''
        >>>
        >>> tree = ast.parse(code)
        >>> classdef = tree.body[0]

        """
        methods = [stmt for stmt in classdef.body if isinstance(stmt, ast.FunctionDef)]
        for method in methods:
            if method.name == self.func_name and method.lineno == self.lineno:
                raise Exception(f'{classdef.name}.{method.name}')
        return classdef 

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 _find_sub_setup_call(
        self, elements
    ):  # type: (List[Any]) -> Tuple[Optional[ast.Call], Optional[List[Any]]]
        for element in elements:
            if not isinstance(element, (ast.FunctionDef, ast.If)):
                continue

            setup_call = self._find_setup_call(element.body)
            if setup_call != (None, None):
                setup_call, body = setup_call

                body = elements + body

                return setup_call, body

        return None, None 

def __init__(self,
                 symbol_table: SymbolTable,
                 custom_types_symbol_table: SymbolTable,
                 external_ir2_symbols_by_name_by_module: Dict[str, Dict[str, Union[ir2.FunctionDefn, ir2.CustomType]]],
                 filename: str,
                 source_lines: List[str],
                 identifier_generator: Iterator[str],
                 function_name: Optional[str] = None,
                 function_definition_line: Optional[int] = None,
                 first_enclosing_except_stmt_line: Optional[int] = None,
                 partially_typechecked_function_definitions_by_name: Dict[str, ast.FunctionDef] = None):
        assert (function_name is None) == (function_definition_line is None)
        self.symbol_table = symbol_table
        self.custom_types_symbol_table = custom_types_symbol_table
        self.external_ir2_symbols_by_name_by_module = external_ir2_symbols_by_name_by_module
        self.partially_typechecked_function_definitions_by_name = partially_typechecked_function_definitions_by_name or dict()
        self.filename = filename
        self.source_lines = source_lines
        self.current_function_name = function_name
        self.current_function_definition_line = function_definition_line
        self.first_enclosing_except_stmt_line = first_enclosing_except_stmt_line
        self.identifier_generator = identifier_generator 

def find_method(module, namespace):
    """Filter away everything except the method

    Promote the method up to the global namespace so that it is
    indistinguishable from a regular function.

    Arguments:
        module (ast.Module): the entire parsed source code
        namespace (str): identifier for the method of interest

    Returns:
        module (ast.Module): the original module but with everything filtered
        away except the method name but with the name `namespace` and promoted
        to the global (i.e. top) level

    """
    module_name, class_name, method_name = namespace.split('.')
    classdefs = [stmt for stmt in module.body if isinstance(stmt, ast.ClassDef)]
    classdef, = [classdef for classdef in classdefs if classdef.name == class_name]
    methods = [stmt for stmt in classdef.body if isinstance(stmt, ast.FunctionDef)]
    for method in methods:
        if method.name == method_name:
            method.name = f'{module_name}.{class_name}.{method_name}'
            module.body = [method]
            return module 

def visit_FunctionDef(self, node):
        '''Visitor for AST FunctionDef nodes

        add relevant information about the node to
        the context for use in tests which inspect function definitions.
        Add the function name to the current namespace for all descendants.
        :param node: The node that is being inspected
        :return: -
        '''

        self.context['function'] = node
        qualname = self.namespace + '.' + b_utils.get_func_name(node)
        name = qualname.split('.')[-1]

        self.context['qualname'] = qualname
        self.context['name'] = name

        # For all child nodes and any tests run, add this function name to
        # current namespace
        self.namespace = b_utils.namespace_path_join(self.namespace, name)
        self.update_scores(self.tester.run_tests(self.context, 'FunctionDef')) 

def get_function_and_class_names(code):
    """Get all function and class names of the code block.

    Inspects the code walking through its AST and returns all
    the names of nodes matching ast.FunctionDef or ast.ClassDef.

    Args:
        code: Multiline string representing Python code

    Returns: List of string names
    """
    names = set()
    tree = ast.parse(code)
    for block in tree.body:
        for node in walk(block):
            if isinstance(node, (ast.FunctionDef, ast.ClassDef,)):
                names.add(node.name)
    return names 

def _find_sub_setup_call(
        self, elements
    ):  # type: (List[Any]) -> Tuple[Optional[ast.Call], Optional[List[Any]]]
        for element in elements:
            if not isinstance(element, (ast.FunctionDef, ast.If)):
                continue

            setup_call = self._find_setup_call(element.body)
            if setup_call != (None, None):
                setup_call, body = setup_call

                body = elements + body

                return setup_call, body

        return None, None 

def format_variables(self, event, last_statement):
        if last_statement:
            last_source_line = event.source.lines[last_statement.lineno - 1]
            dots = (get_leading_spaces(last_source_line)
                    .replace(' ', '.')
                    .replace('\t', '....'))
        else:
            dots = ''
            last_source_line = ''
        lines = []
        for var in event.variables:
            if (u'{} = {}'.format(*var) != last_source_line.strip()
                    and not (
                            isinstance(last_statement, ast.FunctionDef)
                            and not last_statement.decorator_list
                            and var[0] == last_statement.name
                    )
            ):
                lines += self.format_variable(var, dots, event.comprehension_type)
        return lines 

def find_func(module, namespace):
    """Filter away everything except the function

    Addionally rename the function for better readability.

    Args:
        module (ast.Module): the entire parsed code
        namespace (str): identifier for the function of interest

    `namspace` will be of the form <module_name>.<function_name>

    Returns:
        module (ast.Module): the original module but with everything filtered
        away except the function and the function with a more readable name

    """
    module_name, func_name = namespace.split('.')
    funcs = [stmt for stmt in module.body if isinstance(stmt, ast.FunctionDef)]
    func, = [func for func in funcs if func.name == func_name]
    func.name = f'{module_name}.{func_name}'
    module.body = [func]
    return module 

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 getRootInContext(pyAstNode, isClassContext):
    """Return the given node either as-is, or wrapped in a Module node."""
    if not NodeVisitorBases.isScopeNode(pyAstNode):
        raise Exceptions.InternalError(
            "Unsupported type of root node in Analysis (%s)."
            % type(pyAstNode))
    if isinstance(pyAstNode, ast.FunctionDef):
        if isClassContext is None:
            raise Exceptions.InternalError(
                "Value for `isClassContext is required when `type(pyAstNode) is `ast.FunctionDef.")
        elif isClassContext is False:
            pyAstNode = ast.Module([pyAstNode])
        # else nothing
    return pyAstNode 

def test_async_removed_by_transformer(self):
        self.maxDiff = 99999
        async_tree = ast.parse("\n".join([
            "async def a():",
            "   async for b in c():",
            "       await b()",
            "   async with d() as e:",
            "       pass",
            "   return await y()"
        ]))
        self.assertIsInstance(async_tree.body[0], ast.AsyncFunctionDef)
        self.assertIsInstance(async_tree.body[0].body[-1], ast.Return)
        self.assertIsInstance(async_tree.body[0].body[-1].value, ast.Await)

        sync_tree = ast.parse("\n".join([
            "def a():",
            "   for b in c():",
            "       b()",
            "   with d() as e:",
            "       pass",
            "   return y()"
        ]))
        self.assertIsInstance(sync_tree.body[0], ast.FunctionDef)

        transformed = PytTransformer().visit(async_tree)
        self.assertIsInstance(transformed.body[0], ast.FunctionDef)

        self.assertEqual(ast.dump(transformed), ast.dump(sync_tree)) 

def filter_away_except(tree, namespace):
    """Filter away everything except the namespace

    In the case that `namespace` is a method promote the method to the global
    (i.e. top) level.

    Arguments:
        tree (ast.Module): the entire parsed code from the code buffer
        namespace (str): the namespace (i.e. region of code) of interest

    Returns:
        small_tree (ast.Module): `tree` but everything filtered except the
        namespace region

    """
    ns_tokens = namespace.split('.')
    if len(ns_tokens) == 1: # module-level
        tree.body = [stmt for stmt in tree.body if not isinstance(stmt, ast.FunctionDef) and not isinstance(stmt, ast.ClassDef)]
        e = codebook.node_transformers.make_annotation(buffer=namespace, content=f'`{namespace}`', cell_type='1')
        tree.body.insert(0, e)
        small_tree = tree
    elif len(ns_tokens) == 2: # function
        small_tree = find_func(tree, namespace)
    else: # method
        assert len(ns_tokens) == 3
        small_tree = find_method(tree, namespace)
    return small_tree 

def _get_func_nodes():
    """Get all function nodes."""
    return [definition for definition in project_definitions.values()
            if isinstance(definition.node, ast.FunctionDef)] 

def visit_AsyncFunctionDef(self, node):
        return self.visit(ast.FunctionDef(**node.__dict__)) 

def visit_FunctionDef(self, node):
        transformed = ast.FunctionDef(
            name=node.name,
            args=node.args,
            body=self.visit_body(node.body),
            decorator_list=node.decorator_list,
            returns=node.returns
        )
        ast.copy_location(transformed, node)
        return self.generic_visit(transformed) 

def _resolve_name(contract):
        if not isinstance(contract, astroid.Name):
            return contract
        definitions = contract.lookup(contract.name)[1]
        if not definitions:
            return contract
        definition = definitions[0]
        if isinstance(definition, astroid.FunctionDef):
            return definition
        if isinstance(definition, astroid.AssignName):
            return definition.parent.value
        # resolved into something tricky, live with it
        return contract  # pragma: no cover 

def bytecode(self):
        module = ast.parse(TEMPLATE)

        # inject function signature
        func = ast.parse('lambda:0').body[0].value
        func.args = self.func_args
        module.body[3].value = func

        # inject contract
        contract = self.body
        if isinstance(contract, ast.FunctionDef):
            # if contract is function, add it's definition and assign it's name
            # to `contract` variable.
            module.body = [contract] + module.body
            module.body[3].value = ast.Name(
                id=contract.name,
                lineno=1,
                col_offset=1,
                ctx=ast.Load(),
            )
        else:
            if isinstance(contract, ast.Expr):
                contract = contract.value
            module.body[2].value = contract

        return compile(module, filename='<ast>', mode='exec') 

def main():
    file_list = [
        f
        for f in os.listdir(".")
        if re.search(r"^test_.*\.py$", f) and f != ("test_mypy.py" or "test_engine.py")
    ]
    temp_dir = tempfile.gettempdir()

    for f in file_list:
        f_buf = open(f)
        f_text = f_buf.read()
        f_buf.close()
        file_name = "mypy_" + f

        nodes = ast.walk(ast.parse(f_text))
        tests = [
            node.name
            for node in nodes
            if type(node) == ast.FunctionDef and node.name.startswith("test")
        ]

        test_file = temp_dir + "/" + file_name
        mypy_test = open(test_file, "w")
        mypy_test.write("import {}\n\n".format(f[:-3]))
        middle = "()\n" + f[:-3] + "."
        funcs = middle.join(tests)
        funcs += "()"
        funcs = f[:-3] + "." + funcs
        mypy_test.write(funcs)
        mypy_test.close()
        os.system("mypy {}".format(test_file))
        os.remove(test_file)