code: def find_function(root, function_name):
      
      finder = _FindFunctionVisitor(function_name)
      finder.visit(root)
  
      if finder.function is None:
          raise KeyError('function not found', function_name)
      return finder.function
docs: 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)