diff --git a/AGENTS.md b/AGENTS.md index f1bbb96..30aebb5 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 1677f35..b23865e 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.56.0] - 2026-07-16 + +### 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. + ### Fixed - `tilebox-workflows`: Allow task `execute()` methods to use a `-> None` return annotation when postponed annotation @@ -401,7 +412,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/tests/test_client.py b/tilebox-datasets/tests/test_client.py index b6abdbd..a411296 100644 --- a/tilebox-datasets/tests/test_client.py +++ b/tilebox-datasets/tests/test_client.py @@ -1,4 +1,6 @@ import os +import subprocess +import sys from datetime import datetime, timezone from pathlib import Path from unittest.mock import MagicMock, patch @@ -15,6 +17,24 @@ from tilebox.datasets.query.time_interval import us_to_datetime +def test_heavy_imports_are_lazy() -> None: + code = ( + "import sys\n" + "import tilebox.datasets as datasets\n" + "assert not {'pandas', 'xarray'} & sys.modules.keys()\n" + "assert set(datasets.__all__) <= set(dir(datasets))\n" + "from tilebox.datasets.query.time_interval import TimeInterval\n" + "assert not {'pandas', 'xarray'} & sys.modules.keys()\n" + "TimeInterval.parse('2026-01-01')\n" + "assert 'pandas' in sys.modules\n" + "assert 'xarray' not in sys.modules\n" + "client = datasets.Client\n" + "assert datasets.Client is client\n" + "assert 'Client' in vars(datasets)\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) # noqa: S603 + + @pytest.mark.parametrize("url", [_TILEBOX_API_URL, _TILEBOX_DEV_API_URL, f"{_TILEBOX_API_URL}/"]) @patch.dict(os.environ, {}, clear=True) @patch("tilebox.datasets.sync.client.open_channel") diff --git a/tilebox-datasets/tilebox/datasets/__init__.py b/tilebox-datasets/tilebox/datasets/__init__.py index 99984dc..ee5d8c1 100644 --- a/tilebox-datasets/tilebox/datasets/__init__.py +++ b/tilebox-datasets/tilebox/datasets/__init__.py @@ -1,16 +1,55 @@ 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 + + value = Client + case "CollectionClient": + from tilebox.datasets.sync.dataset import CollectionClient # noqa: PLC0415 + + value = CollectionClient + case "DatasetClient": + from tilebox.datasets.sync.dataset import DatasetClient # noqa: PLC0415 + + value = DatasetClient + case "TimeseriesCollection": + from tilebox.datasets.aio.timeseries import TimeseriesCollection # noqa: PLC0415 + + value = TimeseriesCollection + case "TimeseriesDataset": + from tilebox.datasets.aio.timeseries import TimeseriesDataset # noqa: PLC0415 + + value = TimeseriesDataset + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + # Cache the resolved export so subsequent access uses normal module lookup instead of calling __getattr__ again. + globals()[name] = value + return value + + +def __dir__() -> list[str]: + # Include public lazy exports in dir(module) before they have been loaded. + return sorted(set(globals()) | set(__all__)) + + 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..597a2e1 100644 --- a/tilebox-datasets/tilebox/datasets/aio/__init__.py +++ b/tilebox-datasets/tilebox/datasets/aio/__init__.py @@ -1,7 +1,45 @@ -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 + + value = Client + case "CollectionClient": + from tilebox.datasets.aio.dataset import CollectionClient # noqa: PLC0415 + + value = CollectionClient + case "DatasetClient": + from tilebox.datasets.aio.dataset import DatasetClient # noqa: PLC0415 + + value = DatasetClient + case "TimeseriesCollection": + from tilebox.datasets.aio.timeseries import TimeseriesCollection # noqa: PLC0415 + + value = TimeseriesCollection + case "TimeseriesDataset": + from tilebox.datasets.aio.timeseries import TimeseriesDataset # noqa: PLC0415 + + value = TimeseriesDataset + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + # Cache the resolved export so subsequent access uses normal module lookup instead of calling __getattr__ again. + globals()[name] = value + return value + + +def __dir__() -> list[str]: + # Include public lazy exports in dir(module) before they have been loaded. + return sorted(set(globals()) | set(__all__)) diff --git a/tilebox-datasets/tilebox/datasets/query/time_interval.py b/tilebox-datasets/tilebox/datasets/query/time_interval.py index 7c73be7..1d12bc4 100644 --- a/tilebox-datasets/tilebox/datasets/query/time_interval.py +++ b/tilebox-datasets/tilebox/datasets/query/time_interval.py @@ -1,21 +1,22 @@ 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 + _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 +134,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 +197,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/tests/test_storage_client.py b/tilebox-storage/tests/test_storage_client.py index 2aa77f5..9c770b7 100644 --- a/tilebox-storage/tests/test_storage_client.py +++ b/tilebox-storage/tests/test_storage_client.py @@ -1,4 +1,6 @@ import re +import subprocess +import sys from datetime import timedelta from pathlib import Path from tempfile import TemporaryDirectory @@ -9,6 +11,7 @@ from hypothesis import HealthCheck, given, settings from obstore.store import LocalStore +import tilebox.storage.aio as storage_aio from tests.storage_data import ers_granules, landsat_granules, s5p_granules, umbra_granules from tilebox.storage.aio import ( ASFStorageClient, @@ -25,11 +28,44 @@ USGSLandsatStorageGranule, ) +# Warm lazy storage attributes used in Hypothesis-timed tests. Otherwise the first generated example pays the one-time +# import cost and Hypothesis can report flaky deadline failures when replayed examples are much faster. +_ = storage_aio.Boto3CredentialProvider +_ = storage_aio.S3Store + pytestmark = pytest.mark.usefixtures("responses_mock") ASF_LOGIN_URL = "https://urs.earthdata.nasa.gov/oauth/authorize" +def test_heavy_imports_are_lazy() -> None: + code = ( + "import inspect\n" + "import pickle\n" + "import sys\n" + "import tilebox.storage as storage\n" + "heavy = {'tilebox.storage.aio', 'xarray', 'obstore', 'niquests', 'IPython'}\n" + "assert not heavy & sys.modules.keys()\n" + "assert set(storage.__all__) <= set(dir(storage))\n" + "client = storage.LocalFileSystemStorageClient\n" + "assert storage.LocalFileSystemStorageClient is client\n" + "assert inspect.isclass(client)\n" + "assert client.__module__ == 'tilebox.storage'\n" + "assert client.__init__.__qualname__ == 'LocalFileSystemStorageClient.__init__'\n" + "assert pickle.loads(pickle.dumps(client)) is client\n" + "assert 'tilebox.storage.aio' in sys.modules\n" + "assert not {'xarray', 'obstore', 'niquests', 'IPython'} & sys.modules.keys()\n" + "try:\n" + " from tilebox.storage.aio import Image\n" + "except ImportError:\n" + " pass\n" + "else:\n" + " assert inspect.isclass(Image)\n" + " assert 'IPython' in sys.modules\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) # noqa: S603 + + def _mock_asf_login(*, status: int = 200) -> None: responses.add(responses.GET, ASF_LOGIN_URL, status=status) @@ -101,7 +137,7 @@ async def test_quicklook( responses.add(responses.GET, granule.urls.quicklook, body=b"my-quicklook-image") client = _HttpClient(auth={"ASF": ("username", "password")}) - with patch("tilebox.storage.aio.Image"), patch("tilebox.storage.aio._display_quicklook") as display_mock: + with patch("tilebox.storage.aio._display_quicklook") as display_mock: await client.quicklook(granule) display_mock.assert_called_once() assert display_mock.call_args[0][0] == b"my-quicklook-image" @@ -220,7 +256,7 @@ async def test_umbra_storage_client_download(granule: UmbraStorageGranule) -> No await store.put_async(f"sar-data/tasks/{granule.location}/{granule.granule_name}_GEC.tif", b"content1") await store.put_async(f"sar-data/tasks/{granule.location}/{granule.granule_name}_CPHD.cphd", b"content2") - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store umbra = UmbraStorageClient(Path(tmp_path) / "cache") @@ -246,7 +282,7 @@ async def test_umbra_storage_client_list_objects(granule: UmbraStorageGranule) - f"sar-data/tasks/another_{granule.location}/another_{granule.granule_name}_CPHD.cphd", b"content2" ) - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store umbra = UmbraStorageClient(Path(tmp_path) / "cache") @@ -266,7 +302,7 @@ async def test_umbra_storage_client_download_objects(granule: UmbraStorageGranul await store.put_async(f"sar-data/tasks/{granule.location}/{granule.granule_name}_GEC.tif", b"content1") await store.put_async(f"sar-data/tasks/{granule.location}/{granule.granule_name}_CPHD.cphd", b"content2") - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store umbra = UmbraStorageClient(Path(tmp_path) / "cache") @@ -285,7 +321,7 @@ async def test_copernicus_storage_client_download(granule: CopernicusStorageGran store = LocalStore(store_path) await store.put_async(f"{granule.location.removeprefix('/eodata/')}/{granule.granule_name}", b"content1") - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store copernicus = CopernicusStorageClient( access_key="testing", @@ -311,7 +347,7 @@ async def test_copernicus_storage_client_list_objects(granule: CopernicusStorage await store.put_async( f"{granule.location.removeprefix('/eodata/')}_other_granule/{granule.granule_name}", b"content2" ) - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store copernicus = CopernicusStorageClient( access_key="testing", @@ -337,7 +373,7 @@ async def test_copernicus_storage_client_download_objects(granule: CopernicusSto await store.put_async( f"{granule.location.removeprefix('/eodata/')}/other_product_{granule.granule_name}", b"content2" ) - with patch("tilebox.storage.aio.S3Store") as store_mock: + with patch("obstore.store.S3Store") as store_mock: store_mock.return_value = store copernicus = CopernicusStorageClient( access_key="testing", @@ -362,7 +398,7 @@ async def test_landsat_storage_client_download(granule: USGSLandsatStorageGranul await store.put_async( f"{granule.location.removeprefix('s3://usgs-landsat/')}/{granule.granule_name}", b"content1" ) - with patch("tilebox.storage.aio.S3Store") as store_mock, patch("tilebox.storage.aio.Boto3CredentialProvider"): + with patch("obstore.store.S3Store") as store_mock, patch("obstore.auth.boto3.Boto3CredentialProvider"): store_mock.return_value = store landsat = USGSLandsatStorageClient(cache_directory=Path(tmp_path)) @@ -387,7 +423,7 @@ async def test_landsat_storage_client_list_objects(granule: USGSLandsatStorageGr f"{granule.location.removeprefix('s3://usgs-landsat/')}/{granule.granule_name}_thumb_small.jpeg", b"content2", ) - with patch("tilebox.storage.aio.S3Store") as store_mock, patch("tilebox.storage.aio.Boto3CredentialProvider"): + with patch("obstore.store.S3Store") as store_mock, patch("obstore.auth.boto3.Boto3CredentialProvider"): store_mock.return_value = store landsat = USGSLandsatStorageClient(cache_directory=Path(tmp_path)) @@ -411,7 +447,7 @@ async def test_landsat_storage_client_download_objects(granule: USGSLandsatStora await store.put_async( f"{granule.location.removeprefix('s3://usgs-landsat/')}/other_product_{granule.granule_name}", b"content2" ) - with patch("tilebox.storage.aio.S3Store") as store_mock, patch("tilebox.storage.aio.Boto3CredentialProvider"): + with patch("obstore.store.S3Store") as store_mock, patch("obstore.auth.boto3.Boto3CredentialProvider"): store_mock.return_value = store landsat = USGSLandsatStorageClient(cache_directory=Path(tmp_path)) diff --git a/tilebox-storage/tilebox/storage/__init__.py b/tilebox-storage/tilebox/storage/__init__.py index e7785be..e398809 100644 --- a/tilebox-storage/tilebox/storage/__init__.py +++ b/tilebox-storage/tilebox/storage/__init__.py @@ -1,80 +1,34 @@ -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 importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from tilebox.storage._sync import ( + ASFStorageClient, + CopernicusStorageClient, + LocalFileSystemStorageClient, + UmbraStorageClient, + USGSLandsatStorageClient, + ) + +__all__ = [ + "ASFStorageClient", + "CopernicusStorageClient", + "LocalFileSystemStorageClient", + "USGSLandsatStorageClient", + "UmbraStorageClient", +] + + +def __getattr__(name: str) -> Any: + if name not in __all__: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + + storage_client = getattr(import_module("tilebox.storage._sync"), name) + # Cache the resolved export so subsequent access uses normal module lookup instead of calling __getattr__ again. + globals()[name] = storage_client + return storage_client + + +def __dir__() -> list[str]: + # Include public lazy exports in dir(module) before they have been loaded. + return sorted(set(globals()) | set(__all__)) diff --git a/tilebox-storage/tilebox/storage/_sync.py b/tilebox-storage/tilebox/storage/_sync.py new file mode 100644 index 0000000..5172cfe --- /dev/null +++ b/tilebox-storage/tilebox/storage/_sync.py @@ -0,0 +1,93 @@ +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 + +# These classes live here to keep the package import light, but remain public tilebox.storage classes. Keeping their +# module identity stable preserves repr, introspection, and pickle lookup compatibility. + + +class ASFStorageClient(_ASFStorageClient): + __module__ = "tilebox.storage" + + 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): + __module__ = "tilebox.storage" + + 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): + __module__ = "tilebox.storage" + + 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): + __module__ = "tilebox.storage" + + 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): + __module__ = "tilebox.storage" + + 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() diff --git a/tilebox-storage/tilebox/storage/aio.py b/tilebox-storage/tilebox/storage/aio.py index 3bd3365..229638b 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,10 @@ from collections.abc import AsyncIterator from pathlib import Path from pathlib import PurePosixPath as ObjectPath -from typing import Any, TypeAlias +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 +26,52 @@ 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 + + +def __getattr__(name: str) -> Any: + match name: + case "S3Store": + from obstore.store import S3Store # noqa: PLC0415 + + value = S3Store + case "Boto3CredentialProvider": + from obstore.auth.boto3 import Boto3CredentialProvider # noqa: PLC0415 + + value = Boto3CredentialProvider + case "Image" | "HTML" | "display": + value = _lazy_load_ipython()[name] + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + globals()[name] = value + return value + + +def _lazy_load_ipython() -> dict[str, Any]: + try: + from IPython.display import HTML, Image, display # noqa: PLC0415 + except ImportError: + raise ImportError("IPython is not available, please use download_quicklook instead.") from None + return {"HTML": HTML, "Image": Image, "display": display} - class HTML: - def __init__(*_args: Any, **_kwargs: Any) -> None: ... - def display(*_args: Any, **_kwargs: Any) -> None: - raise RuntimeError("IPython is not available. Diagram can only be displayed in a notebook.") +def _s3_store_class() -> type[Any]: + from obstore.store import S3Store # noqa: PLC0415 + return S3Store -ObjectStore: TypeAlias = S3Store | LocalStore | GCSStore + +def _boto3_credential_provider_class() -> type[Any]: + from obstore.auth.boto3 import Boto3CredentialProvider # noqa: PLC0415 + + return Boto3CredentialProvider class _HttpClient(Syncifiable): @@ -210,6 +234,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,12 +275,18 @@ 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 def _display_quicklook(image_data: bytes | Path, width: int, height: int, image_caption: str | None = None) -> None: + display_attrs = _lazy_load_ipython() + Image = display_attrs["Image"] # noqa: N806 + HTML = display_attrs["HTML"] # noqa: N806 + display = display_attrs["display"] display(Image(image_data, width=width, height=height)) if image_caption is not None: display(HTML(image_caption)) @@ -283,6 +315,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 +358,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 +371,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 +476,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 +513,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 +642,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 +785,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 +804,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 +848,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 +962,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 +1050,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..c4cf9a5 100644 --- a/tilebox-storage/tilebox/storage/granule.py +++ b/tilebox-storage/tilebox/storage/granule.py @@ -1,11 +1,15 @@ +from __future__ import annotations + from dataclasses import dataclass from datetime import datetime from pathlib import PurePosixPath as ObjectPath - -import xarray as xr +from typing import TYPE_CHECKING from tilebox.storage.providers import StorageURLs +if TYPE_CHECKING: + import xarray as xr + @dataclass class ASFStorageGranule: @@ -16,7 +20,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 +72,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 +141,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 +180,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 +216,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..bc5446f 100644 --- a/tilebox-storage/tilebox/storage/providers.py +++ b/tilebox-storage/tilebox/storage/providers.py @@ -1,7 +1,11 @@ +from __future__ import annotations + from dataclasses import dataclass from platform import python_version +from typing import TYPE_CHECKING -from niquests import AsyncSession +if TYPE_CHECKING: + from niquests import AsyncSession @dataclass @@ -50,6 +54,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/tests/runner/test_runner.py b/tilebox-workflows/tests/runner/test_runner.py index 2aa53fb..780bfb4 100644 --- a/tilebox-workflows/tests/runner/test_runner.py +++ b/tilebox-workflows/tests/runner/test_runner.py @@ -1,5 +1,7 @@ import os import re +import subprocess +import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -14,6 +16,21 @@ from tilebox.workflows.runner.task_runner import TaskRunner +def test_task_authoring_imports_are_lazy() -> None: + code = ( + "import sys\n" + "import tilebox.workflows as workflows\n" + "heavy = {'pandas', 'xarray', 'boto3', 'google.cloud.storage', 'ipywidgets', 'opentelemetry.sdk'}\n" + "assert not heavy & sys.modules.keys()\n" + "assert set(workflows.__all__) <= set(dir(workflows))\n" + "task = workflows.Task\n" + "assert not heavy & sys.modules.keys()\n" + "assert workflows.Task is task\n" + "assert 'Task' in vars(workflows)\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) # noqa: S603 + + def int_to_bytes(n: int) -> bytes: return n.to_bytes(1, "big") # python3.10 still requires arguments for length and byteorder diff --git a/tilebox-workflows/tilebox/workflows/__init__.py b/tilebox-workflows/tilebox/workflows/__init__.py index 722288c..a18af17 100644 --- a/tilebox-workflows/tilebox/workflows/__init__.py +++ b/tilebox-workflows/tilebox/workflows/__init__.py @@ -1,16 +1,53 @@ 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 + + value = Client + case "ExecutionContext": + from tilebox.workflows.task import ExecutionContext # noqa: PLC0415 + + value = ExecutionContext + case "Job": + from tilebox.workflows.data import Job # noqa: PLC0415 + + value = Job + case "Runner": + from tilebox.workflows.runner.runner import Runner # noqa: PLC0415 + + value = Runner + case "Task": + from tilebox.workflows.task import Task # noqa: PLC0415 + + value = Task + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + # Cache the resolved export so subsequent access uses normal module lookup instead of calling __getattr__ again. + globals()[name] = value + return value + + +def __dir__() -> list[str]: + # Include public lazy exports in dir(module) before they have been loaded. + return sorted(set(globals()) | set(__all__)) + + 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..cef0419 100644 --- a/tilebox-workflows/tilebox/workflows/formatting/job.py +++ b/tilebox-workflows/tilebox/workflows/formatting/job.py @@ -1,19 +1,21 @@ """HTML formatting and ipywidgets for interactive display of Tilebox Workflow jobs.""" +from __future__ import annotations + # some CSS helpers for our Jupyter HTML snippets - inspired by xarray's interactive display import random from collections.abc import Callable 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 + class JobWidget: def __init__(self, refresh_callback: Callable[[UUID], Job] | None = None) -> None: @@ -32,6 +34,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 +63,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 +290,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 +302,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 +336,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..9422d89 100644 --- a/tilebox-workflows/tilebox/workflows/runner/__init__.py +++ b/tilebox-workflows/tilebox/workflows/runner/__init__.py @@ -1,4 +1,31 @@ -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 + + value = Runner + case "TaskRunner": + from tilebox.workflows.runner.task_runner import TaskRunner # noqa: PLC0415 + + value = TaskRunner + case _: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + # Cache the resolved export so subsequent access uses normal module lookup instead of calling __getattr__ again. + globals()[name] = value + return value + + +def __dir__() -> list[str]: + # Include public lazy exports in dir(module) before they have been loaded. + return sorted(set(globals()) | set(__all__)) diff --git a/tilebox-workflows/tilebox/workflows/runner/runner.py b/tilebox-workflows/tilebox/workflows/runner/runner.py index dc6e284..1d6b179 100644 --- a/tilebox-workflows/tilebox/workflows/runner/runner.py +++ b/tilebox-workflows/tilebox/workflows/runner/runner.py @@ -3,8 +3,8 @@ from typing import TYPE_CHECKING from tilebox.workflows.cache import JobCache -from tilebox.workflows.data import TaskIdentifier -from tilebox.workflows.task import RunnerContext, Task, TaskMeta +from tilebox.workflows.data import RunnerContext, TaskIdentifier +from tilebox.workflows.task import Task, TaskMeta if TYPE_CHECKING: from tilebox.workflows.client import Client diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 0ce86c5..3a7e0b7 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 diff --git a/uv.lock b/uv.lock index 148881c..222a13e 100644 --- a/uv.lock +++ b/uv.lock @@ -2141,26 +2141,26 @@ wheels = [ [[package]] name = "prek" -version = "0.4.9" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ea/df/94ed29398576e03494c5aacda8bfed9536edf348bed29cd09f382f2b9b23/prek-0.4.9.tar.gz", hash = "sha256:f8b86441484a5756f3fdb6f3b201d3d448f8845902a84653d78cbf6f875c424f", size = 492711, upload-time = "2026-07-11T11:04:04.332Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/02/632446b72103daf92526203bf83cb76043ebce70588632aa2aea3741da07/prek-0.4.9-py3-none-linux_armv6l.whl", hash = "sha256:7b240ad6f679104309a944c4dd427ccc46d9aaf4f53ee07379c02bf7578c2750", size = 5637604, upload-time = "2026-07-11T11:03:38.985Z" }, - { url = "https://files.pythonhosted.org/packages/57/bf/89daefc85a1db9b8c525731c77121de0a8f57731e81c99d823e919e1f6a5/prek-0.4.9-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:5cbb220d6d77cb047747dcabf421375fdfdd958e01a998a854758698d38fb239", size = 5960396, upload-time = "2026-07-11T11:03:40.953Z" }, - { url = "https://files.pythonhosted.org/packages/42/39/0e448da00671e77740be32cca96530ef64bcdbed178acce7f155b87f6057/prek-0.4.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fd85df4c186becdc47b2e6155de8cace99133c7517387e30f08ca15b93ead11f", size = 5524982, upload-time = "2026-07-11T11:03:42.605Z" }, - { url = "https://files.pythonhosted.org/packages/71/e5/1beceff9cfcf02817c06f38e726c9875cd003f62cbfa516f282fb3c3109b/prek-0.4.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:ba40511145e948d461d6641b556b6ff671b186b0c90743600b9dcb788dd5eb5c", size = 5793691, upload-time = "2026-07-11T11:03:44.106Z" }, - { url = "https://files.pythonhosted.org/packages/7e/17/99e4884c45c46be90e81e220d2bdd5020716963707ef6702361cc398dd91/prek-0.4.9-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0cc5c06ae448568d076bb5e7ce4d630879d6467b2a89fd79061c349a47826efa", size = 5544549, upload-time = "2026-07-11T11:03:45.564Z" }, - { url = "https://files.pythonhosted.org/packages/0c/4b/63f5fe22d867a6e07b12e8a7887a0f3b193c2ef0fa9ed80649f2d257f4ac/prek-0.4.9-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:31af616c216ec7e47913802b082ca816d707d24738fa58eae0463ae296cbe71e", size = 5948798, upload-time = "2026-07-11T11:03:47.424Z" }, - { url = "https://files.pythonhosted.org/packages/1f/89/90d5005436afb6ab99d3ba6599820fe0f0fcb776f5e9e362b5a19e10cbea/prek-0.4.9-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a1ef842d7f19879fac2135c42ba15508f2677346ba800bc8bb82e5f43fe6a6ec", size = 6696938, upload-time = "2026-07-11T11:03:49Z" }, - { url = "https://files.pythonhosted.org/packages/2a/62/068dd25e1106e262b0ce69ffcdc594cf52d85aa13f64628f851e458980d4/prek-0.4.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:442ecf6a454c692bc8a2d5a16936a0fe8e3419727ff208e15818932acca9923a", size = 6170870, upload-time = "2026-07-11T11:03:50.407Z" }, - { url = "https://files.pythonhosted.org/packages/8e/42/bb1f3f3d84af4d12bb675ff071305fa776d3525c10626465d9960ad93c21/prek-0.4.9-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:5b3c590252e3d5724ebed3774695712ff7cb554743b25184808aa5f42b06bd4c", size = 5800355, upload-time = "2026-07-11T11:03:51.882Z" }, - { url = "https://files.pythonhosted.org/packages/a5/ca/dfc8312b4a7ce8fa100ce37a843279d108eec7e7fe2ddbe069448acd1642/prek-0.4.9-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:1dd283b15dec4da29caa3910bd72c8c9d7df93209770503465ad109d6370cb8e", size = 5655126, upload-time = "2026-07-11T11:03:53.388Z" }, - { url = "https://files.pythonhosted.org/packages/25/8a/b19413cb64b81c18502ea7bbef32897f9126b8e53b58cd656996573dc53c/prek-0.4.9-py3-none-musllinux_1_1_armv7l.whl", hash = "sha256:cc25e30e1700a5c7dd9bad665c321e5589b51502bb1bbc6ada45e326d08b428b", size = 5517839, upload-time = "2026-07-11T11:03:54.884Z" }, - { url = "https://files.pythonhosted.org/packages/94/af/900df3f7535e87045df331646c6a01bef6e77dba7f2bdb53483f30cbb988/prek-0.4.9-py3-none-musllinux_1_1_i686.whl", hash = "sha256:4ff5947deeb9a92e6508dfba8b27962dfb927bf60fd36472c9ae862df96fb38c", size = 5802556, upload-time = "2026-07-11T11:03:56.787Z" }, - { url = "https://files.pythonhosted.org/packages/e9/7b/80e560cbe396d0f8687012769cb2d7d7f3428dd019549225ebf205dea806/prek-0.4.9-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:8320ca167d41855d9c4fed66df599f31f96307cbb0da1311a9fe465152e20bd5", size = 6285747, upload-time = "2026-07-11T11:03:58.099Z" }, - { url = "https://files.pythonhosted.org/packages/a3/d4/01ae3b99d09559a69befd128859d0036c604608dd9c6c99986592dde3c21/prek-0.4.9-py3-none-win32.whl", hash = "sha256:b1e8d3bc88ddce6414853468ed8126f45d4ae20f2f4677801ade20ad67a826fa", size = 5320862, upload-time = "2026-07-11T11:03:59.803Z" }, - { url = "https://files.pythonhosted.org/packages/28/68/d038b14f0220fed197be8bc93229e6ea7ad460803ff8a26e4b14a8c81f66/prek-0.4.9-py3-none-win_amd64.whl", hash = "sha256:ed1b4f87a13d1565e8731c60db7fa058966049cbb4d8872d160add510a286558", size = 5706850, upload-time = "2026-07-11T11:04:01.207Z" }, - { url = "https://files.pythonhosted.org/packages/02/4a/57d04de49f591088901794cc22f36563a102681e3512ae17ee6085cd2f30/prek-0.4.9-py3-none-win_arm64.whl", hash = "sha256:7eab3900d9ea614c8ea0d0d55a8b708f0c88e43c966dc8b13a4e36c1e398dd16", size = 5540477, upload-time = "2026-07-11T11:04:02.815Z" }, +version = "0.4.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7c/54/edc21e275f9fa3540d4d98cf349c2de11621d6729cc401bb7aedf563609e/prek-0.4.10.tar.gz", hash = "sha256:db3122f4e780eb4587635e6a83df881caf2dbb1eb7799d1cca51158216d6f33b", size = 502565, upload-time = "2026-07-16T10:13:00.788Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/e7/5a63528ba7b95b64f38db3e253aed49ee8e5e8ba16589889d2b7f809edb7/prek-0.4.10-py3-none-linux_armv6l.whl", hash = "sha256:023f302741d79301346c3088ba43a9592aff0ecdbe5ddc3019fa9b1183319c5e", size = 5694609, upload-time = "2026-07-16T10:12:26.352Z" }, + { url = "https://files.pythonhosted.org/packages/0b/ef/ee9e6bf9a5ce242e9e4e66ac4e2e9042a0f6fd9f367cee18ad404456e93d/prek-0.4.10-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:72adc707e16f97564bbae08d22b222ac3bb2491f8fbfb5a0754f80d472c28a71", size = 6044037, upload-time = "2026-07-16T10:12:28.539Z" }, + { url = "https://files.pythonhosted.org/packages/68/7e/da08cc39e5348ccb9234e63a21ee56861f72e8497d6a78f0db1ccae6515d/prek-0.4.10-py3-none-macosx_11_0_arm64.whl", hash = "sha256:04c9321957e1b32e1fc7cf60bb4f90bba3761f8659d5551ed04f96e25596de49", size = 5535983, upload-time = "2026-07-16T10:12:30.691Z" }, + { url = "https://files.pythonhosted.org/packages/30/c6/0486a35bb687a9beac7a5810bd1104c6da56d469b30b1eeaeefd03c99da2/prek-0.4.10-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:e66ccf6c5e4ebadd05cd98cb338d7f553e4d27aa243cf91279c5a569b3cdccc7", size = 5862085, upload-time = "2026-07-16T10:12:33.042Z" }, + { url = "https://files.pythonhosted.org/packages/52/39/277fe17ae1f121e532e3942456f5a6d01ddacfbc550e481dcb359be7a1b0/prek-0.4.10-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:63f9061d75a50ef0ca92c4b596ad352937a845df80758244950e513b27e9e18f", size = 5605697, upload-time = "2026-07-16T10:12:35.498Z" }, + { url = "https://files.pythonhosted.org/packages/ad/a1/08354af3e000f2656fad086690d834eab6c04631ff41313a219ea6232199/prek-0.4.10-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c2ff7110e4bfaafbbab13c2893a337081aca61ed797f14b6b224d2ea9741eef", size = 6034111, upload-time = "2026-07-16T10:12:37.545Z" }, + { url = "https://files.pythonhosted.org/packages/e4/74/4702396c8d486132e5ce009ab56a0b37f50cb6866830d371f2617b7bdfdc/prek-0.4.10-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2b696a05542e79aa27bcce68d1792e77f4fe6f9c6b012b34d74d62f964f3c72d", size = 6787203, upload-time = "2026-07-16T10:12:40.031Z" }, + { url = "https://files.pythonhosted.org/packages/90/29/b5d5d6fb87ebd64b37471e3e79761de9983f85e14d69c522efe7af6620ce/prek-0.4.10-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:431b44d6054e72815b4b05e1173596dfd02a7f7461211d40a2e3117e414642ad", size = 6261333, upload-time = "2026-07-16T10:12:42.216Z" }, + { url = "https://files.pythonhosted.org/packages/94/d6/54ba696d19f7efdc184093353cce713a850aef9c3556e23faeecafa22e94/prek-0.4.10-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ccbd2b4fd1df790087ba18b4506f680471922a5f13714f19801568434a040dee", size = 5867761, upload-time = "2026-07-16T10:12:44.329Z" }, + { url = "https://files.pythonhosted.org/packages/be/7d/3975098aa2baaabfc10f99f9fcf78045c4f10851beed8e9812b6a2688eab/prek-0.4.10-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:479e7480b447191aa5c6ed67e80f081d0f5ee4e878b140f4d2cee44165395f1c", size = 5714412, upload-time = "2026-07-16T10:12:46.297Z" }, + { url = "https://files.pythonhosted.org/packages/97/c0/3e0aac190fe95fdef98526343559b61d4d9fd54444c8c9137ba02412afe1/prek-0.4.10-py3-none-musllinux_1_1_armv7l.whl", hash = "sha256:0bb7451025cbd2b68e480a13cf665d7a5c87c8b87bf18549a78985c17df817ed", size = 5578145, upload-time = "2026-07-16T10:12:48.261Z" }, + { url = "https://files.pythonhosted.org/packages/d7/44/7b26035534204b8b8a9d5e625479201e616413d287262f557cb32e1f8d77/prek-0.4.10-py3-none-musllinux_1_1_i686.whl", hash = "sha256:4fb047e5776676805794574b2d7b178cb3ab536793aadf172419fcda56b34a57", size = 5889245, upload-time = "2026-07-16T10:12:50.818Z" }, + { url = "https://files.pythonhosted.org/packages/7e/6c/178a9d768876b4211a1bf63907fe308ae02d173639bcf41cea3c5eed35c1/prek-0.4.10-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:08318818d19caf79643babb89f872c92fda134a622b4731df1d6ed61e29d2d26", size = 6372849, upload-time = "2026-07-16T10:12:52.952Z" }, + { url = "https://files.pythonhosted.org/packages/4d/84/d5f5ac8193602883f9dd1d675d9d4084e34fbe3ed2ef50a0c336d8a53d8f/prek-0.4.10-py3-none-win32.whl", hash = "sha256:092872714dcde480a662bbdd98b980b248c2d3e10543d4d53a3a58cc9e5b35b0", size = 5413005, upload-time = "2026-07-16T10:12:55.113Z" }, + { url = "https://files.pythonhosted.org/packages/41/63/9e648fda10bc02c9b6ba305f93b6a6e4fd37d23d13a269a9d2d6bb44eaa1/prek-0.4.10-py3-none-win_amd64.whl", hash = "sha256:3d323a18d0f8c50e474a8fa29fb93bd2db680116d8afb19b76e72ad4667f58e6", size = 5799075, upload-time = "2026-07-16T10:12:56.963Z" }, + { url = "https://files.pythonhosted.org/packages/22/74/b34d8c80cec8dccc7b922c75b9dca62b18b603b5ed2eea93c9d7c2928d2d/prek-0.4.10-py3-none-win_arm64.whl", hash = "sha256:5e93865ef96756c4a26f37ece04ad514abbc19ae6a23ed1a507b6314e6a0d2fb", size = 5563955, upload-time = "2026-07-16T10:12:59.07Z" }, ] [[package]]