feat: add overloads for file, file-pattern and url for from_source

src/earthkit/data/sources/__init__.py

@@ -19,7 +19,13 @@
 from earthkit.data.core.plugins import register as register_plugin
 
 if TYPE_CHECKING:
+    import numpy.array
+
+    from earthkit.data.core.fieldlist import FieldList
     from earthkit.data.data import Data  # type: ignore[import]
+    from earthkit.data.sources import Source
+
+from typing import Any, Callable, Literal, Union, overload
 
 
 class Source(Loader):
@@ -108,6 +114,198 @@ def __getattr__(self, name):
 get_source = SourceMaker()
 
 
+@overload
+def from_source(
+    name: Literal["file"],
+    path: str,
+    expand_user: Union[bool, list, tuple] = True,
+    expand_vars: bool = False,
+    unix_glob: bool = True,
+    recursive_glob: bool = True,
+    filter: Union[str, Callable] = None,
+    parts: list = None,
+    stream: bool = False,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["file-pattern"], pattern: str, *args, hive_partitioning: bool = False, **kwargs
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["url"],
+    url: str,
+    unpack: bool = True,
+    parts: Union[list, tuple] = None,
+    stream: bool = False,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["url-pattern"], url: str, unpack: bool = True, **kwargs) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["sample"], name_or_path: str) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["stream"], stream: Union[list, tuple]) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["memory"], buffer) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["forcings"], source_or_dataset=Union["Source", "FieldList"], *, request: dict = {}, **kwargs
+) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["list-of-dicts"], list_of_dicts: list[dict]) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["multi"], *sources, merger: Union[str, Callable, tuple[str, dict], Any], **kwargs
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["ads"], dataset: str, *args, request: Union[dict, list[dict], tuple[dict]] = None, **kwargs
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["cds"],
+    dataset: str,
+    *args,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    prompt: bool = True,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["ecfs"], path: str) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["ecmwf-open-data"],
+    *args,
+    source: Literal["azure", "ecmwf"] = "ecmwf",
+    model: Literal["ifs", "aifs"] = "ifs",
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["fdb"],
+    *args,
+    config: Union[dict, str] = None,
+    userconfig: Union[dict, str] = None,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    stream: bool = True,
+    lazy: bool = False,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["gribjump"],
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    *,
+    ranges: list[tuple[int, int]] = None,
+    mask: "numpy.array" = None,
+    indices: "numpy.array" = None,
+    fetch_coords_from_fdb: bool = False,
+    fdb_kwargs: dict = None,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["mars"],
+    *args,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    prompt: bool = True,
+    log: Union[str, Callable, dict, None] = "default",
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(name: Literal["opendap"], url: str) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["polytope"],
+    collection: str,
+    *args,
+    address: str = None,
+    user_email: str = None,
+    user_key: str = None,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    stream: bool = True,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["s3"],
+    *args,
+    anon: bool = True,
+    aws_access_key: str = None,
+    aws_secret_access_key: str = None,
+    aws_token: str = None,
+    stream: bool = True,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["wekeo"],
+    dataset: str,
+    *args,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    prompt: bool = True,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["wekeocds"],
+    dataset: str,
+    *args,
+    request: Union[dict, list[dict], tuple[dict]] = None,
+    prompt: bool = True,
+    **kwargs,
+) -> "Data": ...
+
+
+@overload
+def from_source(
+    name: Literal["zarr"],
+    path: str,
+) -> "Data": ...
+
+
 def from_source(name: str, *args, lazily=False, **kwargs) -> "Data":
     if lazily:
         return from_source_lazily(name, *args, **kwargs)