diff --git a/src/earthkit/data/sources/__init__.py b/src/earthkit/data/sources/__init__.py index 06600e14c..5d66bfbc1 100644 --- a/src/earthkit/data/sources/__init__.py +++ b/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)