Merge pull request #36 from d-v-b/chore/normalize-type-hints

chore/normalize type hints

Author: d-v-b

src/eopf_geozarr/conversion/fs_utils.py

@@ -2,11 +2,16 @@
 
 import json
 import os
-from typing import Any, Dict, Optional
+from collections.abc import Mapping
+from typing import Any
 from urllib.parse import urlparse
 
 import s3fs
 import zarr
+from fsspec.implementations.local import LocalFileSystem
+from s3fs import S3FileSystem
+
+from eopf_geozarr.types import S3Credentials, S3FsOptions
 
 
 def normalize_s3_path(s3_path: str) -> str:
@@ -88,7 +93,7 @@ def parse_s3_path(s3_path: str) -> tuple[str, str]:
     return bucket, key
 
 
-def get_s3_storage_options(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
+def get_s3_storage_options(s3_path: str, **s3_kwargs: Any) -> S3FsOptions:
     """
     Get storage options for S3 access with xarray.
 
@@ -105,7 +110,7 @@ def get_s3_storage_options(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
         Storage options dictionary for xarray
     """
     # Set up default S3 configuration
-    default_s3_kwargs = {
+    default_s3_kwargs: S3FsOptions = {
         "anon": False,  # Use credentials
         "use_ssl": True,
         "client_kwargs": {
@@ -121,12 +126,12 @@ def get_s3_storage_options(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
             client_kwargs["endpoint_url"] = os.environ["AWS_ENDPOINT_URL"]
 
     # Merge with user-provided kwargs
-    s3_config = {**default_s3_kwargs, **s3_kwargs}
+    s3_config: S3FsOptions = {**default_s3_kwargs, **s3_kwargs}  # type: ignore[typeddict-item]
 
     return s3_config
 
 
-def get_storage_options(path: str, **kwargs: Any) -> Optional[Dict[str, Any]]:
+def get_storage_options(path: str, **kwargs: Any) -> S3FsOptions | None:
     """
     Get storage options for any URL type, leveraging fsspec as the abstraction layer.
 
@@ -202,7 +207,7 @@ def create_s3_store(s3_path: str, **s3_kwargs: Any) -> str:
 
 
 def write_s3_json_metadata(
-    s3_path: str, metadata: Dict[str, Any], **s3_kwargs: Any
+    s3_path: str, metadata: Mapping[str, Any], **s3_kwargs: Any
 ) -> None:
     """
     Write JSON metadata directly to S3.
@@ -220,7 +225,7 @@ def write_s3_json_metadata(
         Additional keyword arguments for s3fs.S3FileSystem
     """
     # Set up default S3 configuration
-    default_s3_kwargs = {
+    default_s3_kwargs: S3FsOptions = {
         "anon": False,
         "use_ssl": True,
         "asynchronous": False,  # Force synchronous mode
@@ -245,7 +250,7 @@ def write_s3_json_metadata(
         f.write(json_content)
 
 
-def read_s3_json_metadata(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
+def read_s3_json_metadata(s3_path: str, **s3_kwargs: Any) -> dict[str, Any]:
     """
     Read JSON metadata from S3.
 
@@ -262,7 +267,7 @@ def read_s3_json_metadata(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
         Parsed JSON metadata
     """
     # Set up default S3 configuration
-    default_s3_kwargs = {
+    default_s3_kwargs: S3FsOptions = {
         "anon": False,
         "use_ssl": True,
         "asynchronous": False,  # Force synchronous mode
@@ -284,7 +289,7 @@ def read_s3_json_metadata(s3_path: str, **s3_kwargs: Any) -> Dict[str, Any]:
     with fs.open(s3_path, "r") as f:
         content = f.read()
 
-    result: Dict[str, Any] = json.loads(content)
+    result: dict[str, Any] = json.loads(content)
     return result
 
 
@@ -304,7 +309,7 @@ def s3_path_exists(s3_path: str, **s3_kwargs: Any) -> bool:
     bool
         True if the path exists
     """
-    default_s3_kwargs = {
+    default_s3_kwargs: S3FsOptions = {
         "anon": False,
         "use_ssl": True,
         "asynchronous": False,  # Force synchronous mode
@@ -351,7 +356,7 @@ def open_s3_zarr_group(s3_path: str, mode: str = "r", **s3_kwargs: Any) -> zarr.
     )
 
 
-def get_s3_credentials_info() -> Dict[str, Optional[str]]:
+def get_s3_credentials_info() -> S3Credentials:
     """
     Get information about available S3 credentials.
 
@@ -372,7 +377,7 @@ def get_s3_credentials_info() -> Dict[str, Optional[str]]:
     }
 
 
-def validate_s3_access(s3_path: str, **s3_kwargs: Any) -> tuple[bool, Optional[str]]:
+def validate_s3_access(s3_path: str, **s3_kwargs: Any) -> tuple[bool, str | None]:
     """
     Validate that we can access the S3 path.
 
@@ -389,9 +394,9 @@ def validate_s3_access(s3_path: str, **s3_kwargs: Any) -> tuple[bool, Optional[s
         Tuple of (success, error_message)
     """
     try:
-        bucket, key = parse_s3_path(s3_path)
+        bucket, _ = parse_s3_path(s3_path)
 
-        default_s3_kwargs = {
+        default_s3_kwargs: S3FsOptions = {
             "anon": False,
             "use_ssl": True,
             "asynchronous": False,  # Force synchronous mode
@@ -419,7 +424,7 @@ def validate_s3_access(s3_path: str, **s3_kwargs: Any) -> tuple[bool, Optional[s
         return False, str(e)
 
 
-def get_filesystem(path: str, **kwargs: Any) -> Any:
+def get_filesystem(path: str, **kwargs: Any) -> LocalFileSystem | S3FileSystem:
     """
     Get the appropriate fsspec filesystem for any path type.
 
@@ -435,18 +440,17 @@ def get_filesystem(path: str, **kwargs: Any) -> Any:
     fsspec.AbstractFileSystem
         Filesystem instance
     """
-    import fsspec
 
     if is_s3_path(path):
         # Get S3 storage options and use them for fsspec
         storage_options = get_s3_storage_options(path, **kwargs)
-        return fsspec.filesystem("s3", **storage_options)
+        return S3FileSystem(**storage_options)
     else:
         # For local paths, use the local filesystem
-        return fsspec.filesystem("file")
+        return LocalFileSystem(**kwargs)
 
 
-def write_json_metadata(path: str, metadata: Dict[str, Any], **kwargs: Any) -> None:
+def write_json_metadata(path: str, metadata: dict[str, Any], **kwargs: Any) -> None:
     """
     Write JSON metadata to any path type using fsspec.
 
@@ -473,7 +477,7 @@ def write_json_metadata(path: str, metadata: Dict[str, Any], **kwargs: Any) -> N
         f.write(json_content)
 
 
-def read_json_metadata(path: str, **kwargs: Any) -> Dict[str, Any]:
+def read_json_metadata(path: str, **kwargs: Any) -> dict[str, Any]:
     """
     Read JSON metadata from any path type using fsspec.
 
@@ -494,7 +498,7 @@ def read_json_metadata(path: str, **kwargs: Any) -> Dict[str, Any]:
     with fs.open(path, "r") as f:
         content = f.read()
 
-    result: Dict[str, Any] = json.loads(content)
+    result: dict[str, Any] = json.loads(content)
     return result
 
 

src/eopf_geozarr/conversion/geozarr.py

@@ -18,8 +18,8 @@
 import os
 import shutil
 import time
-from collections.abc import Hashable, Iterable, Mapping
-from typing import Any, Dict, List, Tuple
+from collections.abc import Hashable, Iterable, Mapping, Sequence
+from typing import Any, Dict, Tuple
 
 import numpy as np
 import xarray as xr
@@ -55,7 +55,7 @@ def create_geozarr_dataset(
     min_dimension: int = 256,
     tile_width: int = 256,
     max_retries: int = 3,
-    crs_groups: list[str] | None = None,
+    crs_groups: Iterable[str] | None = None,
     gcp_group: str | None = None,
 ) -> xr.DataTree:
     """
@@ -77,8 +77,8 @@ def create_geozarr_dataset(
         Tile width for TMS compatibility
     max_retries : int, default 3
         Maximum number of retries for network operations
-    crs_groups : list[str], optional
-        List of group names that need CRS information added on best-effort basis
+    crs_groups : Iterabl[str], optional
+        Iterable of group names that need CRS information added on best-effort basis
     gcp_group : str, optional
         Group name where GCPs (Ground Control Points) are located.
 
@@ -228,7 +228,7 @@ def iterative_copy(
     min_dimension: int = 256,
     tile_width: int = 256,
     max_retries: int = 3,
-    crs_groups: list[str] | None = None,
+    crs_groups: Iterable[str] | None = None,
     gcp_group: str | None = None,
 ) -> xr.DataTree:
     """
@@ -252,8 +252,8 @@ def iterative_copy(
         Tile width for TMS compatibility
     max_retries : int, default 3
         Maximum number of retries for network operations
-    crs_groups : list[str], optional
-        List of group names that need CRS information added on best-effort basis
+    crs_groups : Iterable[str], optional
+        Iterable of group names that need CRS information added on best-effort basis
     gcp_group : str, optional
         Group name where GCPs (Ground Control Points) are located
 
@@ -273,7 +273,7 @@ def iterative_copy(
         storage_options=storage_options,
     )
 
-    written_groups = set()
+    written_groups: set[str] = set()
     reference_crs = None
 
     # Process all groups in the tree using iterative approach
@@ -811,8 +811,8 @@ def create_native_crs_tile_matrix_set(
         Native CRS of the data
     native_bounds : tuple
         Native bounds (left, bottom, right, top)
-    overview_levels : list
-        List of overview level dictionaries
+    overview_levels : Iterable[OverViewLevelJSON]
+        Iterable of overview level dictionaries
     group_prefix : str, optional
         Group prefix for the tile matrix IDs
 
@@ -883,7 +883,7 @@ def create_overview_dataset_all_vars(
     height: int,
     native_crs: Any,
     native_bounds: Tuple[float, float, float, float],
-    data_vars: List[Hashable],
+    data_vars: Sequence[Hashable],
     ds_gcp: xr.Dataset | None = None,
 ) -> xr.Dataset:
     """
@@ -903,8 +903,8 @@ def create_overview_dataset_all_vars(
         Native CRS of the data
     native_bounds : tuple
         Native bounds (left, bottom, right, top)
-    data_vars : list
-        List of data variable names to include
+    data_vars : Sequence[Hashable]
+        Sequence of data variable names to include
     ds_gcp : xr.Dataset, optional
         Source dataset with Sentinel-1 ground control points
         at native resolution
@@ -1552,7 +1552,7 @@ def _get_lat_coord_attrs() -> StandardLatCoordAttrsJSON:
     }
 
 
-def _find_grid_mapping_var_name(ds: xr.Dataset, data_vars: list[Hashable]) -> str:
+def _find_grid_mapping_var_name(ds: xr.Dataset, data_vars: Sequence[Hashable]) -> str:
     """Find the grid_mapping variable name from the dataset."""
     grid_mapping_var_name = ds.attrs.get("grid_mapping", None)
     if not grid_mapping_var_name and data_vars:

src/eopf_geozarr/types.py

@@ -93,3 +93,38 @@ class TileMatrixSetJSON(TypedDict):
 ]
 """A string literal indicating a resampling method"""
 XARRAY_DIMS_KEY: Final = "_ARRAY_DIMENSIONS"
+
+
+# Why is endpoint URL specified twice?
+class S3ClientOptions(TypedDict):
+    """
+    S3 client options
+    """
+
+    region_name: NotRequired[str]
+    endpoint_url: NotRequired[str]
+
+
+class S3FsOptions(TypedDict):
+    """
+    S3FS options
+    """
+
+    anon: NotRequired[bool]
+    use_ssl: NotRequired[bool]
+    client_kwargs: NotRequired[S3ClientOptions]
+    endpoint_url: NotRequired[str]
+    asynchronous: NotRequired[bool]
+
+
+class S3Credentials(TypedDict):
+    """
+    S3 credentials
+    """
+
+    aws_access_key_id: str | None
+    aws_secret_access_key: str | None
+    aws_session_token: str | None
+    aws_default_region: str
+    aws_profile: str | None
+    AWS_ENDPOINT_URL: str | None