<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def find_function(root, function_name):
 """Search an AST for a function with the given name. A KeyError is raised if the function is not found. root: an AST node function_name: function to search for (string) """

    finder = _FindFunctionVisitor(function_name)
    finder.visit(root)

    if finder.function is None:
        raise KeyError('function not found', function_name)
    return finder.function