def set_extension(path, ext):
    """Set the extension of file.
    Given the path to a file and the desired extension, set the extension of the
    file to the given one.

    Args:
        path: path to a file
        ext: desired extension

    Return:
        path with the correct extension
    """
    ext = ext.replace(".", "")
    current_ext = path.split(".")
    if len(current_ext) == 2:
        path = path.replace(current_ext[1], ext)
    else:
        path = path + "." + ext
    return path