Repository: huggingface/datasets
File: src/datasets/download/streaming_download_manager.py (Python)

Code before:
) -> Tuple[str, dict]:
    prepared_urlpath = []
    prepared_storage_options = {}
    for hop in urlpath.split("::"):
        hop, storage_options = _prepare_single_hop_path_and_storage_options(hop, download_config=download_config)
        prepared_urlpath.append(hop)
        prepared_storage_options.update(storage_options)
    return "::".join(prepared_urlpath), storage_options


def _prepare_single_hop_path_and_storage_options(
    urlpath: str, download_config: Optional[DownloadConfig] = None
) -> Tuple[str, Dict[str, Dict[str, Any]]]:
    """
    Prepare the URL and the kwargs that must be passed to the HttpFileSystem or to requests.get/head

    In particular it resolves google drive URLs
    It also adds the authentication headers for the Hugging Face Hub, for both https:// and hf:// paths.
    """
    token = None if download_config is None else download_config.token
    protocol = urlpath.split("://")[0] if "://" in urlpath else "file"
    if download_config is not None and protocol in download_config.storage_options:
        storage_options = download_config.storage_options[protocol]
    elif download_config is not None and protocol not in download_config.storage_options:
        storage_options = {
            option_name: option_value
            for option_name, option_value in download_config.storage_options.items()
            if option_name not in fsspec.available_protocols()
        }
    else:
        storage_options = {}
    if storage_options:
        storage_options = {protocol: storage_options}
    if protocol in ["http", "https"]:
        storage_options[protocol] = {
            "headers": {
                **get_authentication_headers_for_url(urlpath, token=token),
                "user-agent": get_datasets_user_agent(),
            },
            "client_kwargs": {"trust_env": True},  # Enable reading proxy env variables.
            **(storage_options.get(protocol, {})),
        }
        if "drive.google.com" in urlpath:
            response = http_head(urlpath)
            cookies = None
            for k, v in response.cookies.items():
                if k.startswith("download_warning"):
                    urlpath += "&confirm=" + v
                    cookies = response.cookies
                    storage_options[protocol] = {"cookies": cookies, **storage_options.get(protocol, {})}
        # Fix Google Drive URL to avoid Virus scan warning

Review comment (suggestion):
I think we also need to update `DownloadConfig.storage_options`' type hint.



Code after:
) -> Tuple[str, Dict[str, Dict[str, Any]]]:
    prepared_urlpath = []
    prepared_storage_options = {}
    for hop in urlpath.split("::"):
        hop, storage_options = _prepare_single_hop_path_and_storage_options(hop, download_config=download_config)
        prepared_urlpath.append(hop)
        prepared_storage_options.update(storage_options)
    return "::".join(prepared_urlpath), storage_options


def _prepare_single_hop_path_and_storage_options(
    urlpath: str, download_config: Optional[DownloadConfig] = None
) -> Tuple[str, Dict[str, Dict[str, Any]]]:
    """
    Prepare the URL and the kwargs that must be passed to the HttpFileSystem or to requests.get/head

    In particular it resolves google drive URLs
    It also adds the authentication headers for the Hugging Face Hub, for both https:// and hf:// paths.

    Storage options are formatted in the form {protocol: storage_options_for_protocol}
    """
    token = None if download_config is None else download_config.token
    protocol = urlpath.split("://")[0] if "://" in urlpath else "file"
    if download_config is not None and protocol in download_config.storage_options:
        storage_options = download_config.storage_options[protocol]
    elif download_config is not None and protocol not in download_config.storage_options:
        storage_options = {
            option_name: option_value
            for option_name, option_value in download_config.storage_options.items()
            if option_name not in fsspec.available_protocols()
        }
    else:
        storage_options = {}
    if storage_options:
        storage_options = {protocol: storage_options}
    if protocol in ["http", "https"]:
        storage_options[protocol] = {
            "headers": {
                **get_authentication_headers_for_url(urlpath, token=token),
                "user-agent": get_datasets_user_agent(),
            },
            "client_kwargs": {"trust_env": True},  # Enable reading proxy env variables.
            **(storage_options.get(protocol, {})),
        }
        if "drive.google.com" in urlpath:
            response = http_head(urlpath)
            cookies = None
            for k, v in response.cookies.items():
                if k.startswith("download_warning"):
                    urlpath += "&confirm=" + v
                    cookies = response.cookies