Add type overloads to load_dataset for better static type inference (#7888)

* Add type overloads to load_dataset for better static type inference

Fixes #7883

This PR adds @overload decorators to load_dataset() to help type checkers
like Pylance and mypy correctly infer the return type based on the split
and streaming parameters.

Changes:
- Added typing imports (Literal, overload) to load.py
- Added 4 @overload signatures that map argument combinations to specific return types:
  * split=None, streaming=False -> DatasetDict
  * split specified, streaming=False -> Dataset
  * split=None, streaming=True -> IterableDatasetDict
  * split specified, streaming=True -> IterableDataset

This resolves the Pylance error where to_csv() was not recognized on Dataset
objects returned by load_dataset(..., split='train'), since the type checker
previously saw the return type as a Union that included types without to_csv().

No runtime behavior changes - this is purely a static typing improvement.

* make style

---------

Co-authored-by: Quentin Lhoest <[email protected]>

Author: Aditya2755

src/datasets/load.py

@@ -25,7 +25,7 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Union, overload
 
 import fsspec
 import httpx
@@ -1187,6 +1187,101 @@ def load_dataset_builder(
     return builder_instance
 
 
+@overload
+def load_dataset(
+    path: str,
+    name: Optional[str] = None,
+    data_dir: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    split: None = None,
+    cache_dir: Optional[str] = None,
+    features: Optional[Features] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    verification_mode: Optional[Union[VerificationMode, str]] = None,
+    keep_in_memory: Optional[bool] = None,
+    save_infos: bool = False,
+    revision: Optional[Union[str, Version]] = None,
+    token: Optional[Union[bool, str]] = None,
+    streaming: Literal[False] = False,
+    num_proc: Optional[int] = None,
+    storage_options: Optional[dict] = None,
+    **config_kwargs: Any,
+) -> DatasetDict: ...
+
+
+@overload
+def load_dataset(
+    path: str,
+    name: Optional[str] = None,
+    data_dir: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    *,
+    split: Union[str, Split, list[str], list[Split]],
+    cache_dir: Optional[str] = None,
+    features: Optional[Features] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    verification_mode: Optional[Union[VerificationMode, str]] = None,
+    keep_in_memory: Optional[bool] = None,
+    save_infos: bool = False,
+    revision: Optional[Union[Version, str]] = None,
+    token: Optional[Union[bool, str]] = None,
+    streaming: Literal[False] = False,
+    num_proc: Optional[int] = None,
+    storage_options: Optional[dict] = None,
+    **config_kwargs: Any,
+) -> Dataset: ...
+
+
+@overload
+def load_dataset(
+    path: str,
+    name: Optional[str] = None,
+    data_dir: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    split: None = None,
+    cache_dir: Optional[str] = None,
+    features: Optional[Features] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    verification_mode: Optional[Union[VerificationMode, str]] = None,
+    keep_in_memory: Optional[bool] = None,
+    save_infos: bool = False,
+    revision: Optional[Union[Version, str]] = None,
+    token: Optional[Union[bool, str]] = None,
+    *,
+    streaming: Literal[True],
+    num_proc: Optional[int] = None,
+    storage_options: Optional[dict] = None,
+    **config_kwargs: Any,
+) -> IterableDatasetDict: ...
+
+
+@overload
+def load_dataset(
+    path: str,
+    name: Optional[str] = None,
+    data_dir: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    *,
+    split: Union[str, Split, list[str], list[Split]],
+    cache_dir: Optional[str] = None,
+    features: Optional[Features] = None,
+    download_config: Optional[DownloadConfig] = None,
+    download_mode: Optional[Union[DownloadMode, str]] = None,
+    verification_mode: Optional[Union[VerificationMode, str]] = None,
+    keep_in_memory: Optional[bool] = None,
+    save_infos: bool = False,
+    revision: Optional[Union[Version, str]] = None,
+    token: Optional[Union[bool, str]] = None,
+    streaming: Literal[True],
+    num_proc: Optional[int] = None,
+    storage_options: Optional[dict] = None,
+    **config_kwargs: Any,
+) -> IterableDataset: ...
+
+
 def load_dataset(
     path: str,
     name: Optional[str] = None,