Repository: huggingface/datasets
File: tests/test_load.py (Python)

Code before:
    data_name, tmp_dir, datafiles = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, data_files=datafiles)
    return dataset


def test_load_dataset_distributed(tmp_path, csv_path):
    num_workers = 5
    args = "csv", str(tmp_path), csv_path
    with Pool(processes=num_workers) as pool:  # start num_workers processes
        datasets = pool.map(distributed_load_dataset, [args] * num_workers)
        assert len(datasets) == num_workers
        assert all(len(dataset) == len(datasets[0]) > 0 for dataset in datasets)
        assert len(datasets[0].cache_files) > 0
        assert all(dataset.cache_files == datasets[0].cache_files for dataset in datasets)


def distributed_load_dataset_with_script(args):
    data_name, tmp_dir, download_mode = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, download_mode=download_mode)
    return dataset


@require_not_windows  # windows doesn't support overwriting Arrow files from other processes
@pytest.mark.parametrize("download_mode", [None, "force_redownload"])
def test_load_dataset_distributed_with_script(tmp_path, download_mode):
    # we need to cehck in the "force_redownload" case
    # since in `_copy_script_and_other_resources_in_importable_dir()` we might delete the directory
    # containing the .py file while the other processes use it
    num_workers = 5
    args = (SAMPLE_DATASET_IDENTIFIER, str(tmp_path), download_mode)
    with Pool(processes=num_workers) as pool:  # start num_workers processes
        datasets = pool.map(distributed_load_dataset_with_script, [args] * num_workers)
        assert len(datasets) == num_workers
        assert all(len(dataset) == len(datasets[0]) > 0 for dataset in datasets)
        assert len(datasets[0].cache_files) > 0
        assert all(dataset.cache_files == datasets[0].cache_files for dataset in datasets)


def test_load_dataset_with_storage_options(mockfs):
    with mockfs.open("data.txt", "w") as f:
        f.write("Hello there\n")
        f.write("General Kenobi !")
    data_files = {"train": ["mock://data.txt"]}
    ds = load_dataset("text", data_files=data_files, storage_options=mockfs.storage_options)
    assert list(ds["train"]) == [{"text": "Hello there"}, {"text": "General Kenobi !"}]


@require_pil
def test_load_dataset_with_storage_options_with_decoding(mockfs, image_file):
    import PIL.Image


Review comment (suggestion):
Nit:

```suggestion
    # we need to check in the "force_redownload" case
```

Code after:
    data_name, tmp_dir, datafiles = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, data_files=datafiles)
    return dataset


def test_load_dataset_distributed(tmp_path, csv_path):
    num_workers = 5
    args = "csv", str(tmp_path), csv_path
    with Pool(processes=num_workers) as pool:  # start num_workers processes
        datasets = pool.map(distributed_load_dataset, [args] * num_workers)
        assert len(datasets) == num_workers
        assert all(len(dataset) == len(datasets[0]) > 0 for dataset in datasets)
        assert len(datasets[0].cache_files) > 0
        assert all(dataset.cache_files == datasets[0].cache_files for dataset in datasets)


def distributed_load_dataset_with_script(args):
    data_name, tmp_dir, download_mode = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, download_mode=download_mode)
    return dataset


@require_not_windows  # windows doesn't support overwriting Arrow files from other processes
@pytest.mark.parametrize("download_mode", [None, "force_redownload"])
def test_load_dataset_distributed_with_script(tmp_path, download_mode):
    # we need to check in the "force_redownload" case
    # since in `_copy_script_and_other_resources_in_importable_dir()` we might delete the directory
    # containing the .py file while the other processes use it
    num_workers = 5
    args = (SAMPLE_DATASET_IDENTIFIER, str(tmp_path), download_mode)
    with Pool(processes=num_workers) as pool:  # start num_workers processes
        datasets = pool.map(distributed_load_dataset_with_script, [args] * num_workers)
        assert len(datasets) == num_workers
        assert all(len(dataset) == len(datasets[0]) > 0 for dataset in datasets)
        assert len(datasets[0].cache_files) > 0
        assert all(dataset.cache_files == datasets[0].cache_files for dataset in datasets)


def test_load_dataset_with_storage_options(mockfs):
    with mockfs.open("data.txt", "w") as f:
        f.write("Hello there\n")
        f.write("General Kenobi !")
    data_files = {"train": ["mock://data.txt"]}
    ds = load_dataset("text", data_files=data_files, storage_options=mockfs.storage_options)
    assert list(ds["train"]) == [{"text": "Hello there"}, {"text": "General Kenobi !"}]


@require_pil
def test_load_dataset_with_storage_options_with_decoding(mockfs, image_file):
    import PIL.Image