diff --git a/AGENTS.md b/AGENTS.md index 93581c4..b9c2c5b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,14 @@ Pre-commit hooks include YAML checks, EOF fixer, `sync-with-uv`, Ruff, and `ty`. - Logging uses `loguru` in several packages; workflows also supports explicit logger/tracer configuration. - Tests use `pytest`, with async coverage (`pytest-asyncio`) and property-based testing (`hypothesis`) in multiple packages. +### Import-Time Discipline + +- Keep `tilebox.workflows` task-authoring imports light. Release runners create fresh virtual environments, so avoid + importing heavy optional/runtime dependencies (`pandas`, `numpy`, `xarray`, cloud SDKs, `ipywidgets`, OpenTelemetry + SDK/exporters, cache backends) from package `__init__` modules or core `Runner`/`Task` import paths. +- Prefer lazy imports inside the methods that actually need those dependencies. Module-level `__getattr__` aliases are + acceptable for public package aliases and are supported by Python 3.7+. + ## Protobuf And Generated Code Generated files live under paths such as: diff --git a/CHANGELOG.md b/CHANGELOG.md index 935907e..c149e03 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.55.1] - 2026-07-02 + +### Changed + +- `tilebox-workflows`: Reduced import-time overhead for release runners by lazily loading optional heavy dependencies + such as datasets, pandas, cloud SDKs, notebook widgets, and runtime/observability modules until they are needed. +- `tilebox-datasets`: Reduced import-time overhead by lazily exporting the root and async client APIs and deferring + pandas/xarray imports in time interval parsing until parsing requires them. +- `tilebox-storage`: Reduced startup overhead by lazily exporting sync storage clients and deferring geospatial, + notebook, object-store, cloud SDK, HTTP, and progress-display dependencies until storage operations require them. + ## [0.55.0] - 2026-07-01 ### Added @@ -394,7 +405,8 @@ the first client that does not cache data (since it's already on the local file - Released under the [MIT](https://opensource.org/license/mit) license. - Released packages: `tilebox-datasets`, `tilebox-workflows`, `tilebox-storage`, `tilebox-grpc` -[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.55.0...HEAD +[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.55.1...HEAD +[0.55.1]: https://github.com/tilebox/tilebox-python/compare/v0.55.0...v0.55.1 [0.55.0]: https://github.com/tilebox/tilebox-python/compare/v0.54.0...v0.55.0 [0.54.0]: https://github.com/tilebox/tilebox-python/compare/v0.53.0...v0.54.0 [0.53.0]: https://github.com/tilebox/tilebox-python/compare/v0.52.0...v0.53.0 diff --git a/tilebox-datasets/tilebox/datasets/__init__.py b/tilebox-datasets/tilebox/datasets/__init__.py index 99984dc..e1923c9 100644 --- a/tilebox-datasets/tilebox/datasets/__init__.py +++ b/tilebox-datasets/tilebox/datasets/__init__.py @@ -1,16 +1,46 @@ import os import sys +from typing import TYPE_CHECKING, Any from loguru import logger -# only here for backwards compatibility, to preserve backwards compatibility with older imports -from tilebox.datasets.aio.timeseries import TimeseriesCollection, TimeseriesDataset -from tilebox.datasets.sync.client import Client -from tilebox.datasets.sync.dataset import CollectionClient, DatasetClient +if TYPE_CHECKING: + from tilebox.datasets.aio.timeseries import TimeseriesCollection, TimeseriesDataset + from tilebox.datasets.sync.client import Client + from tilebox.datasets.sync.dataset import CollectionClient, DatasetClient __all__ = ["Client", "CollectionClient", "DatasetClient", "TimeseriesCollection", "TimeseriesDataset"] +def __getattr__(name: str) -> Any: + # PEP 562 module __getattr__ is supported since Python 3.7. Keep these aliases lazy so importing a focused + # submodule like tilebox.datasets.query.id_interval does not also import the sync/aio clients and their data-model + # dependencies. + match name: + case "Client": + from tilebox.datasets.sync.client import Client # noqa: PLC0415 + + return Client + case "CollectionClient": + from tilebox.datasets.sync.dataset import CollectionClient # noqa: PLC0415 + + return CollectionClient + case "DatasetClient": + from tilebox.datasets.sync.dataset import DatasetClient # noqa: PLC0415 + + return DatasetClient + case "TimeseriesCollection": + from tilebox.datasets.aio.timeseries import TimeseriesCollection # noqa: PLC0415 + + return TimeseriesCollection + case "TimeseriesDataset": + from tilebox.datasets.aio.timeseries import TimeseriesDataset # noqa: PLC0415 + + return TimeseriesDataset + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + def _init_logging(level: str = "INFO") -> None: logger.remove() logger.add(sys.stdout, level=level, format="{message}", catch=True) diff --git a/tilebox-datasets/tilebox/datasets/aio/__init__.py b/tilebox-datasets/tilebox/datasets/aio/__init__.py index 6e5a314..82d1072 100644 --- a/tilebox-datasets/tilebox/datasets/aio/__init__.py +++ b/tilebox-datasets/tilebox/datasets/aio/__init__.py @@ -1,7 +1,36 @@ -from tilebox.datasets.aio.client import Client -from tilebox.datasets.aio.dataset import CollectionClient, DatasetClient +from typing import TYPE_CHECKING, Any -# only here for backwards compatibility, to preserve backwards compatibility with older imports -from tilebox.datasets.aio.timeseries import TimeseriesCollection, TimeseriesDataset +if TYPE_CHECKING: + from tilebox.datasets.aio.client import Client + from tilebox.datasets.aio.dataset import CollectionClient, DatasetClient + from tilebox.datasets.aio.timeseries import TimeseriesCollection, TimeseriesDataset __all__ = ["Client", "CollectionClient", "DatasetClient", "TimeseriesCollection", "TimeseriesDataset"] + + +def __getattr__(name: str) -> Any: + # PEP 562 module __getattr__ is supported since Python 3.7. Keep these aliases lazy so importing + # tilebox.datasets.aio does not also import xarray/pandas-backed dataset clients. + match name: + case "Client": + from tilebox.datasets.aio.client import Client # noqa: PLC0415 + + return Client + case "CollectionClient": + from tilebox.datasets.aio.dataset import CollectionClient # noqa: PLC0415 + + return CollectionClient + case "DatasetClient": + from tilebox.datasets.aio.dataset import DatasetClient # noqa: PLC0415 + + return DatasetClient + case "TimeseriesCollection": + from tilebox.datasets.aio.timeseries import TimeseriesCollection # noqa: PLC0415 + + return TimeseriesCollection + case "TimeseriesDataset": + from tilebox.datasets.aio.timeseries import TimeseriesDataset # noqa: PLC0415 + + return TimeseriesDataset + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py index 4eda495..517512e 100644 --- a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py +++ b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py @@ -231,7 +231,7 @@ def resize(self, buffer_size: int) -> None: elif buffer_size > len(self._data): # resize the data buffer to the new capacity, by just padding it with zeros at the end missing = buffer_size - len(self._data) - self._data = np.pad( # ty: ignore[no-matching-overload] + self._data = np.pad( self._data, ((0, missing), (0, 0)), constant_values=self._type.fill_value, @@ -309,7 +309,7 @@ def _resize(self) -> None: else: # resize the data buffer to the new capacity, by just padding it with zeros at the end missing_capacity = self._capacity - self._data.shape[0] missing_array_dim = self._array_dim - self._data.shape[1] - self._data = np.pad( # ty: ignore[no-matching-overload] + self._data = np.pad( self._data, ((0, missing_capacity), (0, missing_array_dim), (0, 0)), constant_values=self._type.fill_value, diff --git a/tilebox-datasets/tilebox/datasets/query/time_interval.py b/tilebox-datasets/tilebox/datasets/query/time_interval.py index 7c73be7..80eed0f 100644 --- a/tilebox-datasets/tilebox/datasets/query/time_interval.py +++ b/tilebox-datasets/tilebox/datasets/query/time_interval.py @@ -1,21 +1,26 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import TypeAlias +from typing import TYPE_CHECKING, Any, TypeAlias -import numpy as np -import xarray as xr from google.protobuf.duration_pb2 import Duration from google.protobuf.timestamp_pb2 import Timestamp -from pandas.core.tools.datetimes import DatetimeScalar, to_datetime from tilebox.datasets.tilebox.v1 import query_pb2 +if TYPE_CHECKING: + from pandas.core.tools.datetimes import DatetimeScalar + from xarray import DataArray, Dataset +else: + DataArray = Any + Dataset = Any + DatetimeScalar = Any + _SMALLEST_POSSIBLE_TIMEDELTA = timedelta(microseconds=1) _EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc) # A type alias for the different types that can be used to specify a time interval TimeIntervalLike: TypeAlias = ( - "DatetimeScalar | tuple[DatetimeScalar, DatetimeScalar] | xr.DataArray | xr.Dataset | TimeInterval" + "DatetimeScalar | tuple[DatetimeScalar, DatetimeScalar] | list[DatetimeScalar] | DataArray | Dataset | TimeInterval" ) # once we require python >= 3.12 we can replace this with a type statement, which doesn't require a string at all # type TimeIntervalLike = DatetimeScalar | tuple[DatetimeScalar ... | TimeInterval @@ -133,30 +138,34 @@ def parse(cls, arg: TimeIntervalLike) -> "TimeInterval": TimeInterval: The parsed time interval """ - match arg: - case TimeInterval(_, _, _, _): - return arg - case (start, end): - return TimeInterval(start=_convert_to_datetime(start), end=_convert_to_datetime(end)) - case point_in_time if isinstance(point_in_time, DatetimeScalar | int): - dt = _convert_to_datetime(point_in_time) - return TimeInterval(start=dt, end=dt, start_exclusive=False, end_inclusive=True) - case arr if ( - isinstance(arr, xr.DataArray) - and arr.ndim == 1 - and arr.size > 0 - and arr.dtype == np.dtype("datetime64[ns]") - ): - start = arr.data[0] - end = arr.data[-1] - return TimeInterval( - start=_convert_to_datetime(start), - end=_convert_to_datetime(end), - start_exclusive=False, - end_inclusive=True, - ) - case ds if isinstance(ds, xr.Dataset) and "time" in ds.coords: - return cls.parse(ds.time) + if isinstance(arg, TimeInterval): + return arg + + if isinstance(arg, list | tuple) and len(arg) == 2: + start, end = arg + return TimeInterval(start=_convert_to_datetime(start), end=_convert_to_datetime(end)) + + from pandas.core.tools.datetimes import DatetimeScalar # noqa: PLC0415 + + if isinstance(arg, DatetimeScalar | int): + dt = _convert_to_datetime(arg) + return TimeInterval(start=dt, end=dt, start_exclusive=False, end_inclusive=True) + + import numpy as np # noqa: PLC0415 + import xarray as xr # noqa: PLC0415 + + if isinstance(arg, xr.DataArray) and arg.ndim == 1 and arg.size > 0 and arg.dtype == np.dtype("datetime64[ns]"): + start = arg.data[0] + end = arg.data[-1] + return TimeInterval( + start=_convert_to_datetime(start), + end=_convert_to_datetime(end), + start_exclusive=False, + end_inclusive=True, + ) + + if isinstance(arg, xr.Dataset) and "time" in arg.coords: + return cls.parse(arg.time) raise ValueError(f"Failed to convert {arg} ({type(arg)}) to TimeInterval)") @@ -192,8 +201,10 @@ def to_message(self) -> query_pb2.TimeInterval: _EMPTY_TIME_INTERVAL = TimeInterval(_EPOCH, _EPOCH, start_exclusive=True, end_inclusive=False) -def _convert_to_datetime(arg: DatetimeScalar) -> datetime: +def _convert_to_datetime(arg: Any) -> datetime: """Convert the given datetime scalar to a datetime object in the UTC timezone""" + from pandas.core.tools.datetimes import to_datetime # noqa: PLC0415 + dt: datetime = to_datetime(arg, utc=True).to_pydatetime() if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) diff --git a/tilebox-storage/tilebox/storage/__init__.py b/tilebox-storage/tilebox/storage/__init__.py index e7785be..e057448 100644 --- a/tilebox-storage/tilebox/storage/__init__.py +++ b/tilebox-storage/tilebox/storage/__init__.py @@ -1,80 +1,128 @@ +from importlib import import_module from pathlib import Path - -from tilebox.storage.aio import ASFStorageClient as _ASFStorageClient -from tilebox.storage.aio import CopernicusStorageClient as _CopernicusStorageClient -from tilebox.storage.aio import LocalFileSystemStorageClient as _LocalFileSystemStorageClient -from tilebox.storage.aio import UmbraStorageClient as _UmbraStorageClient -from tilebox.storage.aio import USGSLandsatStorageClient as _USGSLandsatStorageClient - - -class ASFStorageClient(_ASFStorageClient): - def __init__(self, user: str, password: str, cache_directory: Path = Path.home() / ".cache" / "tilebox") -> None: - """A tilebox storage client that downloads data from the Alaska Satellite Facility. - - Args: - user: The username to use for authentication. - password: The password to use for authentication. - cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None - no cache is used and the `output_dir` parameter will need be set when downloading data. - """ - super().__init__(user, password, cache_directory) - self._syncify() - - -class UmbraStorageClient(_UmbraStorageClient): - def __init__(self, cache_directory: Path | None = Path.home() / ".cache" / "tilebox") -> None: - """A tilebox storage client that downloads data from the Umbra Open Data Catalog. - - Args: - cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None - no cache is used and the `output_dir` parameter will need be set when downloading data. - """ - super().__init__(cache_directory) - self._syncify() - - -class CopernicusStorageClient(_CopernicusStorageClient): - def __init__( - self, - access_key: str | None = None, - secret_access_key: str | None = None, - cache_directory: Path | None = Path.home() / ".cache" / "tilebox", - ) -> None: - """A tilebox storage client that downloads data from the Copernicus EO data. - - Args: - access_key: The S3 Copernicus access key. If not provided, the AWS_ACCESS_KEY_ID environment - variable will be used. - secret_access_key: The S3 Copernicus secret access key. If not provided, the AWS_SECRET_ACCESS_KEY - environment variable will be used. - cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None - no cache is used and the `output_dir` parameter will need be set when downloading data. - """ - super().__init__(access_key, secret_access_key, cache_directory) - self._syncify() - - -class USGSLandsatStorageClient(_USGSLandsatStorageClient): - def __init__(self, cache_directory: Path | None = Path.home() / ".cache" / "tilebox") -> None: - """A tilebox storage client that downloads data from the USGS Landsat S3 bucket. - - This client handles the requester-pays nature of the bucket and provides methods for listing and downloading - data. - - Args: - cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None - no cache is used and the `output_dir` parameter will need be set when downloading data. - """ - super().__init__(cache_directory) - self._syncify() - - -class LocalFileSystemStorageClient(_LocalFileSystemStorageClient): - def __init__(self, root: Path) -> None: - """A tilebox storage client for accessing data on a local file system, or a mounted network file system. - - Args: - root: The root directory of the file system to access. - """ - super().__init__(root) - self._syncify() +from typing import Any + +__all__ = [ + "ASFStorageClient", + "CopernicusStorageClient", + "LocalFileSystemStorageClient", + "USGSLandsatStorageClient", + "UmbraStorageClient", +] + + +def __getattr__(name: str) -> Any: + # PEP 562 module __getattr__ is supported since Python 3.7. Keep these sync aliases lazy so importing + # tilebox.storage does not also import the async storage stack and geospatial/notebook dependencies. + try: + init = _SYNC_CLIENT_INITIALIZERS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + + storage_client = type( + name, + (_aio_class(name),), + { + "__module__": __name__, + "__qualname__": name, + "__init__": init, + }, + ) + globals()[name] = storage_client + return storage_client + + +def _aio_class(name: str) -> type[Any]: + return getattr(import_module("tilebox.storage.aio"), name) + + +def _init_asf_storage_client( + self: Any, user: str, password: str, cache_directory: Path = Path.home() / ".cache" / "tilebox" +) -> None: + """A tilebox storage client that downloads data from the Alaska Satellite Facility. + + Args: + user: The username to use for authentication. + password: The password to use for authentication. + cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None + no cache is used and the `output_dir` parameter will need be set when downloading data. + """ + _aio_class("ASFStorageClient").__init__(self, user, password, cache_directory) + self._syncify() + + +def _init_umbra_storage_client(self: Any, cache_directory: Path | None = Path.home() / ".cache" / "tilebox") -> None: + """A tilebox storage client that downloads data from the Umbra Open Data Catalog. + + Args: + cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None + no cache is used and the `output_dir` parameter will need be set when downloading data. + """ + _aio_class("UmbraStorageClient").__init__(self, cache_directory) + self._syncify() + + +def _init_copernicus_storage_client( + self: Any, + access_key: str | None = None, + secret_access_key: str | None = None, + cache_directory: Path | None = Path.home() / ".cache" / "tilebox", +) -> None: + """A tilebox storage client that downloads data from the Copernicus EO data. + + Args: + access_key: The S3 Copernicus access key. If not provided, the AWS_ACCESS_KEY_ID environment + variable will be used. + secret_access_key: The S3 Copernicus secret access key. If not provided, the AWS_SECRET_ACCESS_KEY + environment variable will be used. + cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None + no cache is used and the `output_dir` parameter will need be set when downloading data. + """ + _aio_class("CopernicusStorageClient").__init__(self, access_key, secret_access_key, cache_directory) + self._syncify() + + +def _init_usgs_landsat_storage_client( + self: Any, cache_directory: Path | None = Path.home() / ".cache" / "tilebox" +) -> None: + """A tilebox storage client that downloads data from the USGS Landsat S3 bucket. + + This client handles the requester-pays nature of the bucket and provides methods for listing and downloading + data. + + Args: + cache_directory: The directory to store downloaded data in. Defaults to ~/.cache/tilebox. If set to None + no cache is used and the `output_dir` parameter will need be set when downloading data. + """ + _aio_class("USGSLandsatStorageClient").__init__(self, cache_directory) + self._syncify() + + +def _init_local_file_system_storage_client(self: Any, root: Path) -> None: + """A tilebox storage client for accessing data on a local file system, or a mounted network file system. + + Args: + root: The root directory of the file system to access. + """ + _aio_class("LocalFileSystemStorageClient").__init__(self, root) + self._syncify() + + +for _init in [ + _init_asf_storage_client, + _init_umbra_storage_client, + _init_copernicus_storage_client, + _init_usgs_landsat_storage_client, + _init_local_file_system_storage_client, +]: + _init.__name__ = "__init__" + _init.__qualname__ = "__init__" + + +_SYNC_CLIENT_INITIALIZERS = { + "ASFStorageClient": _init_asf_storage_client, + "CopernicusStorageClient": _init_copernicus_storage_client, + "LocalFileSystemStorageClient": _init_local_file_system_storage_client, + "UmbraStorageClient": _init_umbra_storage_client, + "USGSLandsatStorageClient": _init_usgs_landsat_storage_client, +} diff --git a/tilebox-storage/tilebox/storage/aio.py b/tilebox-storage/tilebox/storage/aio.py index 3bd3365..1a2ff58 100644 --- a/tilebox-storage/tilebox/storage/aio.py +++ b/tilebox-storage/tilebox/storage/aio.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import contextlib import hashlib @@ -9,16 +11,11 @@ from collections.abc import AsyncIterator from pathlib import Path from pathlib import PurePosixPath as ObjectPath -from typing import Any, TypeAlias +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, TypeAlias import anyio -import niquests -import obstore as obs -import xarray as xr from aiofile import async_open -from obstore.auth.boto3 import Boto3CredentialProvider -from obstore.store import GCSStore, LocalStore, S3Store -from tqdm.auto import tqdm from _tilebox.grpc.aio.producer_consumer import async_producer_consumer from _tilebox.grpc.aio.syncify import Syncifiable @@ -30,24 +27,70 @@ USGSLandsatStorageGranule, _is_copernicus_odata_url, ) -from tilebox.storage.providers import login -try: - from IPython.display import HTML, Image, display -except ImportError: - # IPython is not available, so we can't display the quicklook image - # but let's define stubs for the type checker - class Image: - def __init__(*_args: Any, **_kwargs: Any) -> None: ... +if TYPE_CHECKING: + import niquests + import xarray as xr + from obstore.store import GCSStore, LocalStore, S3Store + + ObjectStore: TypeAlias = S3Store | LocalStore | GCSStore +else: + ObjectStore = Any + niquests = SimpleNamespace(AsyncSession=Any) + xr = SimpleNamespace(Dataset=Any) + + +def Image(*args: Any, **kwargs: Any) -> Any: # noqa: N802 + return _ipython_display_attr("Image")(*args, **kwargs) + + +def HTML(*args: Any, **kwargs: Any) -> Any: # noqa: N802 + return _ipython_display_attr("HTML")(*args, **kwargs) + + +def display(*args: Any, **kwargs: Any) -> Any: + return _ipython_display_attr("display")(*args, **kwargs) + + +def _ipython_display_attr(name: str) -> Any: + try: + from IPython import display as ipython_display # noqa: PLC0415 + except ImportError: + raise ImportError("IPython is not available, please use download_quicklook instead.") from None + return getattr(ipython_display, name) + - class HTML: - def __init__(*_args: Any, **_kwargs: Any) -> None: ... +def __getattr__(name: str) -> Any: + # Keep these heavy dependencies lazy while preserving public module attributes that tests and users may patch. + match name: + case "S3Store": + from obstore.store import S3Store # noqa: PLC0415 - def display(*_args: Any, **_kwargs: Any) -> None: - raise RuntimeError("IPython is not available. Diagram can only be displayed in a notebook.") + return S3Store + case "Boto3CredentialProvider": + from obstore.auth.boto3 import Boto3CredentialProvider # noqa: PLC0415 + return Boto3CredentialProvider + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -ObjectStore: TypeAlias = S3Store | LocalStore | GCSStore + +def _s3_store_class() -> type[Any]: + try: + return globals()["S3Store"] + except KeyError: + from obstore.store import S3Store # noqa: PLC0415 + + return S3Store + + +def _boto3_credential_provider_class() -> type[Any]: + try: + return globals()["Boto3CredentialProvider"] + except KeyError: + from obstore.auth.boto3 import Boto3CredentialProvider # noqa: PLC0415 + + return Boto3CredentialProvider class _HttpClient(Syncifiable): @@ -210,6 +253,8 @@ async def downloader() -> AsyncIterator[bytes]: md5 = hashlib.md5() if verify else None # noqa: S324 progress = None if show_progress: + from tqdm.auto import tqdm # noqa: PLC0415 + progress = tqdm(total=granule.file_size, unit="B", unit_scale=True, unit_divisor=1024) async def writer(chunk: bytes) -> None: @@ -249,6 +294,8 @@ async def _client(self, storage_provider: str) -> niquests.AsyncSession: raise ValueError(f"Missing credentials for storage provider '{storage_provider}'") auth = self._auth[storage_provider] + from tilebox.storage.providers import login # noqa: PLC0415 + client = await login(storage_provider, auth) self._clients[storage_provider] = client return client @@ -283,6 +330,8 @@ async def destroy_cache(self) -> None: async def list_object_paths(store: ObjectStore, prefix: str) -> list[str]: + import obstore as obs # noqa: PLC0415 + objects = await obs.list(store, prefix).collect_async() prefix_path = ObjectPath(prefix) return sorted(str(ObjectPath(obj["path"]).relative_to(prefix_path)) for obj in objects) @@ -324,6 +373,8 @@ async def _download_worker( async def _download_object( store: ObjectStore, prefix: str, obj: str, output_dir: Path, show_progress: bool = True ) -> Path: + import obstore as obs # noqa: PLC0415 + key = str(ObjectPath(prefix) / obj) output_path = output_dir / obj if output_path.exists(): # already cached @@ -335,6 +386,8 @@ async def _download_object( file_size = response.meta["size"] with download_path.open("wb") as f: if show_progress: + from tqdm.auto import tqdm # noqa: PLC0415 + with tqdm(desc=obj, total=file_size, unit="B", unit_scale=True, unit_divisor=1024) as progress: async for bytes_chunk in response: f.write(bytes_chunk) @@ -438,8 +491,6 @@ async def quicklook(self, datapoint: xr.Dataset | ASFStorageGranule, width: int Image: The quicklook image. """ granule = ASFStorageGranule.from_data(datapoint) - if Image is None: - raise ImportError("IPython is not available, please use download_quicklook instead.") quicklook = await self._download_quicklook(datapoint) _display_quicklook(quicklook, width, height, f"Image {quicklook.name} © ASF {granule.time.year}") @@ -477,7 +528,8 @@ def __init__(self, cache_directory: Path | None = Path.home() / ".cache" / "tile """ super().__init__(cache_directory) - self._store: ObjectStore = S3Store(self._BUCKET, region=self._REGION, skip_signature=True) + s3_store = _s3_store_class() + self._store: ObjectStore = s3_store(self._BUCKET, region=self._REGION, skip_signature=True) async def list_objects(self, datapoint: xr.Dataset | UmbraStorageGranule) -> list[str]: """List all available objects for a given datapoint. @@ -605,7 +657,8 @@ def __init__( f"To get access to the Copernicus data, please visit: https://documentation.dataspace.copernicus.eu/APIs/S3.html" ) - self._store = S3Store( + s3_store = _s3_store_class() + self._store = s3_store( bucket=self._BUCKET, endpoint=self._ENDPOINT_URL, access_key_id=access_key, @@ -747,8 +800,6 @@ async def quicklook( ImportError: In case IPython is not available. ValueError: If no quicklook is available for the given datapoint. """ - if Image is None: - raise ImportError("IPython is not available, please use download_quicklook instead.") granule = CopernicusStorageGranule.from_data(datapoint) quicklook = await self._download_quicklook(granule) _display_quicklook(quicklook, width, height, f"{granule.granule_name} © ESA {granule.time.year}") @@ -768,6 +819,8 @@ async def _download_quicklook(self, datapoint: xr.Dataset | CopernicusStorageGra if _is_copernicus_odata_url(granule.thumbnail): # the thumbnail is not stored in the S3 bucket, but is accessible via a public URL. So download it # directly. + import niquests # noqa: PLC0415 + response = await niquests.aget( granule.thumbnail, allow_redirects=True ) # to check if the thumbnail is accessible, raises if not @@ -810,8 +863,10 @@ def __init__(self, cache_directory: Path | None = Path.home() / ".cache" / "tile """ super().__init__(cache_directory) - self._store = S3Store( - self._BUCKET, region=self._REGION, request_payer=True, credential_provider=Boto3CredentialProvider() + boto3_credential_provider = _boto3_credential_provider_class() + s3_store = _s3_store_class() + self._store = s3_store( + self._BUCKET, region=self._REGION, request_payer=True, credential_provider=boto3_credential_provider() ) async def list_objects(self, datapoint: xr.Dataset | USGSLandsatStorageGranule) -> list[str]: @@ -922,8 +977,6 @@ async def quicklook( ImportError: In case IPython is not available. ValueError: If no quicklook is available for the given datapoint. """ - if Image is None: - raise ImportError("IPython is not available, please use download_quicklook instead.") quicklook = await self._download_quicklook(datapoint) _display_quicklook(quicklook, width, height, f"Image {quicklook.name} © USGS") @@ -1012,6 +1065,4 @@ async def quicklook( height: Display height of the image in pixels. Defaults to 600. """ quicklook_path = await self._download_quicklook(datapoint) - if Image is None: - raise ImportError("IPython is not available, please use download_quicklook instead.") _display_quicklook(quicklook_path, width, height, None) diff --git a/tilebox-storage/tilebox/storage/granule.py b/tilebox-storage/tilebox/storage/granule.py index 03441d5..902de3c 100644 --- a/tilebox-storage/tilebox/storage/granule.py +++ b/tilebox-storage/tilebox/storage/granule.py @@ -1,11 +1,18 @@ +from __future__ import annotations + from dataclasses import dataclass from datetime import datetime from pathlib import PurePosixPath as ObjectPath - -import xarray as xr +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any from tilebox.storage.providers import StorageURLs +if TYPE_CHECKING: + import xarray as xr +else: + xr = SimpleNamespace(Dataset=Any) + @dataclass class ASFStorageGranule: @@ -16,7 +23,7 @@ class ASFStorageGranule: urls: StorageURLs @classmethod - def from_data(cls, dataset: "xr.Dataset | ASFStorageGranule") -> "ASFStorageGranule": + def from_data(cls, dataset: xr.Dataset | ASFStorageGranule) -> ASFStorageGranule: """Extract the granule information from a datapoint given as xarray dataset.""" if isinstance(dataset, ASFStorageGranule): return dataset @@ -68,7 +75,7 @@ class UmbraStorageGranule: location: str @classmethod - def from_data(cls, dataset: "xr.Dataset | UmbraStorageGranule") -> "UmbraStorageGranule": + def from_data(cls, dataset: xr.Dataset | UmbraStorageGranule) -> UmbraStorageGranule: """Extract the granule information from a datapoint given as xarray dataset.""" if isinstance(dataset, UmbraStorageGranule): return dataset @@ -137,7 +144,7 @@ class CopernicusStorageGranule: thumbnail: str | None = None @classmethod - def from_data(cls, dataset: "xr.Dataset | CopernicusStorageGranule") -> "CopernicusStorageGranule": + def from_data(cls, dataset: xr.Dataset | CopernicusStorageGranule) -> CopernicusStorageGranule: """Extract the granule information from a datapoint given as xarray dataset.""" if isinstance(dataset, CopernicusStorageGranule): return dataset @@ -176,7 +183,7 @@ class USGSLandsatStorageGranule: thumbnail: str | None = None @classmethod - def from_data(cls, dataset: "xr.Dataset | USGSLandsatStorageGranule") -> "USGSLandsatStorageGranule": + def from_data(cls, dataset: xr.Dataset | USGSLandsatStorageGranule) -> USGSLandsatStorageGranule: """Extract the granule information from a datapoint given as xarray dataset.""" if isinstance(dataset, USGSLandsatStorageGranule): return dataset @@ -212,7 +219,7 @@ class LocationStorageGranule: thumbnail: str | None = None @classmethod - def from_data(cls, dataset: "xr.Dataset | LocationStorageGranule") -> "LocationStorageGranule": + def from_data(cls, dataset: xr.Dataset | LocationStorageGranule) -> LocationStorageGranule: """Extract the granule information from a datapoint given as xarray dataset.""" if isinstance(dataset, LocationStorageGranule): return dataset diff --git a/tilebox-storage/tilebox/storage/providers.py b/tilebox-storage/tilebox/storage/providers.py index 7a3f1a7..ead8997 100644 --- a/tilebox-storage/tilebox/storage/providers.py +++ b/tilebox-storage/tilebox/storage/providers.py @@ -1,7 +1,13 @@ +from __future__ import annotations + from dataclasses import dataclass from platform import python_version +from typing import TYPE_CHECKING, Any -from niquests import AsyncSession +if TYPE_CHECKING: + from niquests import AsyncSession +else: + AsyncSession = Any @dataclass @@ -50,6 +56,8 @@ async def _asf_login(auth: tuple[str, str]) -> AsyncSession: "Client-Id": client_id, } + from niquests import AsyncSession # noqa: PLC0415 + client = AsyncSession(auth=auth, headers=headers) response = await client.get( login_url, diff --git a/tilebox-workflows/tilebox/workflows/__init__.py b/tilebox-workflows/tilebox/workflows/__init__.py index 722288c..949d5cb 100644 --- a/tilebox-workflows/tilebox/workflows/__init__.py +++ b/tilebox-workflows/tilebox/workflows/__init__.py @@ -1,16 +1,44 @@ import os import sys +from typing import TYPE_CHECKING, Any from loguru import logger -from tilebox.workflows.client import Client -from tilebox.workflows.data import Job -from tilebox.workflows.runner.runner import Runner -from tilebox.workflows.task import ExecutionContext, Task +if TYPE_CHECKING: + from tilebox.workflows.client import Client + from tilebox.workflows.data import Job + from tilebox.workflows.runner.runner import Runner + from tilebox.workflows.task import ExecutionContext, Task __all__ = ["Client", "ExecutionContext", "Job", "Runner", "Task"] +def __getattr__(name: str) -> Any: + match name: + case "Client": + from tilebox.workflows.client import Client # noqa: PLC0415 + + return Client + case "ExecutionContext": + from tilebox.workflows.task import ExecutionContext # noqa: PLC0415 + + return ExecutionContext + case "Job": + from tilebox.workflows.data import Job # noqa: PLC0415 + + return Job + case "Runner": + from tilebox.workflows.runner.runner import Runner # noqa: PLC0415 + + return Runner + case "Task": + from tilebox.workflows.task import Task # noqa: PLC0415 + + return Task + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + def _init_logging(level: str = "INFO") -> None: logger.remove() logger.add(sys.stdout, level=level, format="{process}: {level}: {message}", catch=True) diff --git a/tilebox-workflows/tilebox/workflows/cache.py b/tilebox-workflows/tilebox/workflows/cache.py index 0966b7a..50c64fc 100644 --- a/tilebox-workflows/tilebox/workflows/cache.py +++ b/tilebox-workflows/tilebox/workflows/cache.py @@ -5,13 +5,15 @@ from io import BytesIO from pathlib import Path from pathlib import PurePosixPath as ObjectPath +from typing import TYPE_CHECKING, Any -import boto3 -from botocore.exceptions import ClientError -from google.cloud.exceptions import NotFound -from google.cloud.storage import Blob, Bucket -from obstore.exceptions import GenericError -from obstore.store import ObjectStore +if TYPE_CHECKING: + from google.cloud.storage import Blob, Bucket + from obstore.store import ObjectStore +else: + Blob = Any + Bucket = Any + ObjectStore = Any class JobCache(ABC): @@ -98,6 +100,8 @@ def __delitem__(self, key: str) -> None: raise KeyError(f"{key} is not cached!") from None def __getitem__(self, key: str) -> bytes: + from obstore.exceptions import GenericError # noqa: PLC0415 + try: entry = self.store.get(str(self.prefix / key)) return bytes(entry.bytes()) @@ -262,6 +266,8 @@ def __setitem__(self, key: str, value: bytes) -> None: self._blob(key).upload_from_file(BytesIO(value)) def __getitem__(self, key: str) -> bytes: + from google.cloud.exceptions import NotFound # noqa: PLC0415 + try: # GCS library has some weird typing issues, so let's ignore them for now return self._blob(key).download_as_bytes() @@ -301,11 +307,15 @@ def __init__(self, bucket: str, prefix: str | ObjectPath = "jobs") -> None: self.bucket = bucket self.prefix = ObjectPath(prefix) with warnings.catch_warnings(): + import boto3 # noqa: PLC0415 + # https://github.com/boto/boto3/issues/3889 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*datetime.utcnow.*") self._s3 = boto3.client("s3") def __contains__(self, key: str) -> bool: + from botocore.exceptions import ClientError # noqa: PLC0415 + try: self._s3.head_object(Bucket=self.bucket, Key=str(self.prefix / key)) except ClientError as e: @@ -320,6 +330,8 @@ def __setitem__(self, key: str, value: bytes) -> None: self._s3.upload_fileobj(BytesIO(value), self.bucket, str(self.prefix / key)) def __getitem__(self, key: str) -> bytes: + from botocore.exceptions import ClientError # noqa: PLC0415 + item = BytesIO() try: self._s3.download_fileobj(self.bucket, str(self.prefix / key), item) diff --git a/tilebox-workflows/tilebox/workflows/data.py b/tilebox-workflows/tilebox/workflows/data.py index 0f89eda..450fb30 100644 --- a/tilebox-workflows/tilebox/workflows/data.py +++ b/tilebox-workflows/tilebox/workflows/data.py @@ -6,13 +6,9 @@ from enum import Enum from functools import lru_cache from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from uuid import UUID -import boto3 -import pandas as pd -from google.cloud.storage import Client as GoogleStorageClient -from google.cloud.storage.bucket import Bucket from google.protobuf.duration_pb2 import Duration from opentelemetry.proto.common.v1 import common_pb2 from opentelemetry.proto.logs.v1 import logs_pb2 @@ -30,13 +26,16 @@ ) from tilebox.datasets.uuid import must_uuid_to_uuid_message, uuid_message_to_uuid, uuid_to_uuid_message -try: - # let's not make this a hard dependency, but if it's installed we can use its types +if TYPE_CHECKING: + from google.cloud.storage.bucket import Bucket from mypy_boto3_s3.client import S3Client -except ModuleNotFoundError: - from typing import Any as S3Client -from tilebox.workflows.observability.tracing import NoopWorkflowTracer, WorkflowTracer + from tilebox.workflows.observability.tracing import WorkflowTracer +else: + Bucket = Any + S3Client = Any + WorkflowTracer = Any + from tilebox.workflows.workflows.v1 import automation_pb2 as automation_pb from tilebox.workflows.workflows.v1 import core_pb2, job_pb2, task_pb2, workflows_pb2 @@ -758,6 +757,8 @@ def to_message(self) -> logs_pb2.ResourceLogs: class LogRecords(list[LogRecord]): def to_pandas(self) -> Any: """Convert log records to a pandas DataFrame.""" + import pandas as pd # noqa: PLC0415 + return pd.DataFrame([asdict(record) for record in self]) @@ -836,6 +837,8 @@ def to_message(self) -> trace_pb2.ResourceSpans: class Spans(list[Span]): def to_pandas(self) -> Any: """Convert spans to a pandas DataFrame.""" + import pandas as pd # noqa: PLC0415 + rows = [] for span in self: row = asdict(span) @@ -1139,6 +1142,8 @@ def __init__( storage_locations: list[StorageLocation] | None = None, ) -> None: if tracer is None: + from tilebox.workflows.observability.tracing import NoopWorkflowTracer # noqa: PLC0415 + tracer = NoopWorkflowTracer() self.tracer = tracer self.storage_locations = { @@ -1150,6 +1155,8 @@ def gcs_client(self, location: str) -> Bucket: return _default_google_storage_client(location) def s3_client(self, location: str) -> S3Client: + import boto3 # noqa: PLC0415 + _ = location # we always use the default s3 client, regardless of the location with warnings.catch_warnings(): # https://github.com/boto/boto3/issues/3889 @@ -1162,6 +1169,8 @@ def local_path(self, location: str) -> Path: @lru_cache def _default_google_storage_client(location: str) -> Bucket: + from google.cloud.storage import Client as GoogleStorageClient # noqa: PLC0415 + project, bucket = location.split(":") return GoogleStorageClient(project=project).bucket(bucket) diff --git a/tilebox-workflows/tilebox/workflows/formatting/job.py b/tilebox-workflows/tilebox/workflows/formatting/job.py index b1e18d1..2b60604 100644 --- a/tilebox-workflows/tilebox/workflows/formatting/job.py +++ b/tilebox-workflows/tilebox/workflows/formatting/job.py @@ -6,14 +6,17 @@ from dataclasses import dataclass, field from datetime import datetime from threading import Event, Thread -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import UUID -from dateutil.tz import tzlocal -from ipywidgets import HTML, HBox, IntProgress, VBox - from tilebox.workflows.data import Job, JobState +if TYPE_CHECKING: + from ipywidgets import HBox, VBox +else: + HBox = Any + VBox = Any + class JobWidget: def __init__(self, refresh_callback: Callable[[UUID], Job] | None = None) -> None: @@ -32,6 +35,8 @@ def _repr_mimebundle_(self, *args: Any, **kwargs: Any) -> dict[str, str] | None: return None if self.layout is None: # initialize the widget the first time we want to interactively display it + from ipywidgets import HTML, VBox # noqa: PLC0415 + self.widgets.append(HTML(_render_job_details_html(self.job))) self.widgets.append(HTML(_render_job_progress(self.job, False))) self.widgets.extend( @@ -59,12 +64,16 @@ def _refresh_worker(self) -> None: progress = self.refresh_callback(self.job.id) updated = False if last_progress is None: # first time, don't add the refresh time + from ipywidgets import HTML # noqa: PLC0415 + self.widgets[1] = HTML(_render_job_progress(progress, False)) updated = True elif ( progress.state != last_progress.state or progress.execution_stats.first_task_started_at != last_progress.execution_stats.first_task_started_at ): + from ipywidgets import HTML # noqa: PLC0415 + self.widgets[1] = HTML(_render_job_progress(progress, True)) updated = True @@ -282,6 +291,8 @@ def _render_job_details_html(job: Job) -> str: def _render_datetime(dt: datetime) -> str: + from dateutil.tz import tzlocal # noqa: PLC0415 + local = dt.astimezone(tzlocal()) time_part = local.strftime("%Y-%m-%d %H:%M:%S") tz_part = local.strftime("%z") @@ -292,6 +303,8 @@ def _render_datetime(dt: datetime) -> str: def _render_job_progress(job: Job, include_refresh_time: bool) -> str: refresh = "" if include_refresh_time: + from dateutil.tz import tzlocal # noqa: PLC0415 + current_time = datetime.now(tzlocal()) refresh = f" (refreshed at {current_time.strftime('%H:%M:%S')}) {_info_icon}" @@ -324,6 +337,8 @@ def _render_job_progress(job: Job, include_refresh_time: bool) -> str: def _progress_indicator_bar(label: str, done: int, total: int, state: JobState) -> HBox: + from ipywidgets import HTML, HBox, IntProgress # noqa: PLC0415 + percentage = done / total if total > 0 else 0 if done <= total else 1 non_completed_color = ( _BAR_COLORS["failed"] if state in (JobState.FAILED, JobState.CANCELED) else _BAR_COLORS["running"] diff --git a/tilebox-workflows/tilebox/workflows/runner/__init__.py b/tilebox-workflows/tilebox/workflows/runner/__init__.py index bd7cb57..cf00572 100644 --- a/tilebox-workflows/tilebox/workflows/runner/__init__.py +++ b/tilebox-workflows/tilebox/workflows/runner/__init__.py @@ -1,4 +1,22 @@ -from tilebox.workflows.runner.runner import Runner -from tilebox.workflows.runner.task_runner import TaskRunner +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from tilebox.workflows.runner.runner import Runner + from tilebox.workflows.runner.task_runner import TaskRunner __all__ = ["Runner", "TaskRunner"] + + +def __getattr__(name: str) -> Any: + # PEP 562 module __getattr__ is supported since Python 3.7. + match name: + case "Runner": + from tilebox.workflows.runner.runner import Runner # noqa: PLC0415 + + return Runner + case "TaskRunner": + from tilebox.workflows.runner.task_runner import TaskRunner # noqa: PLC0415 + + return TaskRunner + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tilebox-workflows/tilebox/workflows/runner/runner.py b/tilebox-workflows/tilebox/workflows/runner/runner.py index dc6e284..e56eca8 100644 --- a/tilebox-workflows/tilebox/workflows/runner/runner.py +++ b/tilebox-workflows/tilebox/workflows/runner/runner.py @@ -1,15 +1,22 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from tilebox.workflows.cache import JobCache -from tilebox.workflows.data import TaskIdentifier -from tilebox.workflows.task import RunnerContext, Task, TaskMeta +from tilebox.workflows.task import Task, TaskMeta if TYPE_CHECKING: + from tilebox.workflows.cache import JobCache from tilebox.workflows.client import Client from tilebox.workflows.clusters.client import ClusterSlugLike + from tilebox.workflows.data import RunnerContext, TaskIdentifier from tilebox.workflows.runner.task_runner import TaskRunner +else: + Client = Any + ClusterSlugLike = Any + JobCache = Any + RunnerContext = Any + TaskIdentifier = Any + TaskRunner = Any class Runner: @@ -22,7 +29,7 @@ def __init__( ) -> None: self.cache = cache self.context = context - self._tasks_by_identifier: dict[TaskIdentifier, type[Task]] = {} + self._tasks_by_identifier: dict[Any, type[Task]] = {} for task in tasks or []: self.register(task) diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 9c73e14..5219b55 100644 --- a/tilebox-workflows/tilebox/workflows/task.py +++ b/tilebox-workflows/tilebox/workflows/task.py @@ -7,14 +7,19 @@ from collections.abc import Sequence from dataclasses import dataclass, fields, is_dataclass from types import NoneType, UnionType -from typing import Any, Generic, TypeVar, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin # from python 3.11 onwards this is available as typing.dataclass_transform: from typing_extensions import dataclass_transform from tilebox.workflows.data import RunnerContext, TaskIdentifier, TaskSubmissionGroup, TaskSubmissions -from tilebox.workflows.observability.logging import StructuredLogger -from tilebox.workflows.observability.tracing import WorkflowTracer + +if TYPE_CHECKING: + from tilebox.workflows.observability.logging import StructuredLogger + from tilebox.workflows.observability.tracing import WorkflowTracer +else: + StructuredLogger = Any + WorkflowTracer = Any META_ATTR = "__tilebox_task_meta__" # the name of the attribute we use to store task metadata on the class