Removes the dataset directory and archive file from the data home directory.

    Parameters
    ----------
    dataset : str
        The name of the dataset; should either be a folder in data home or
        specified in the yellowbrick.datasets.DATASETS variable.

    data_home : str, optional
        The path on disk where data is stored. If not passed in, it is looked
        up from YELLOWBRICK_DATA or the default returned by ``get_data_home``.

    ext : str, default: ".zip"
        The extension of the archive file.

    Returns
    -------
    removed : int
        The number of objects removed from data_home.

def cleanup_dataset(dataset, data_home=None, ext=".zip"):
    """
    Removes the dataset directory and archive file from the data home directory.

    Parameters
    ----------
    dataset : str
        The name of the dataset; should either be a folder in data home or
        specified in the yellowbrick.datasets.DATASETS variable.

    data_home : str, optional
        The path on disk where data is stored. If not passed in, it is looked
        up from YELLOWBRICK_DATA or the default returned by ``get_data_home``.

    ext : str, default: ".zip"
        The extension of the archive file.

    Returns
    -------
    removed : int
        The number of objects removed from data_home.
    """
    removed = 0
    data_home = get_data_home(data_home)

    # Paths to remove
    datadir = os.path.join(data_home, dataset)
    archive = os.path.join(data_home, dataset+ext)

    # Remove directory and contents
    if os.path.exists(datadir):
        shutil.rmtree(datadir)
        removed += 1

    # Remove the archive file
    if os.path.exists(archive):
        os.remove(archive)
        removed += 1

    return removed