aereo 1.1.0__py3-none-any.whl

aereo/asset_downloader/__init__.py

@@ -0,0 +1,17 @@
+"""
+Safe asset downloading and cleanup routines
+"""
+
+from aereo.asset_downloader.core import (
+    cleanup_asset_safely,
+    download_asset_safely,
+    download_assets_safely,
+    extract_asset_safely,
+)
+
+__all__ = [
+    "cleanup_asset_safely",
+    "download_asset_safely",
+    "download_assets_safely",
+    "extract_asset_safely",
+]

aereo/asset_downloader/core.py

@@ -0,0 +1,194 @@
+import shutil
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from aereo.interfaces import Downloader
+
+
+def download_asset_safely(
+    href: str,
+    local_path: Path,
+    s3_client: Optional[Any] = None,
+    http_session: Optional[Any] = None,
+    downloader: Optional["Downloader"] = None,
+) -> None:
+    """Download asset with a filelock to avoid corruption in multi-processing.
+
+    Args:
+        href: URL or local path to the asset.
+        local_path: Destination path for the downloaded file.
+        s3_client: Optional authenticated S3FileSystem for S3 access.
+            Note: Only works when running in AWS us-west-2 region.
+        http_session: Optional authenticated requests.Session for HTTPS downloads.
+            Use earthaccess.get_requests_https_session() to create.
+            Works from anywhere (no AWS region requirement).
+        downloader: Optional callable that handles the download itself.
+            If provided, it is called unconditionally inside the file lock
+            and all built-in logic is skipped.
+    """
+    import filelock
+    import s3fs
+    import requests
+
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+    lock_path = local_path.with_suffix(".lock")
+
+    with filelock.FileLock(str(lock_path)):
+        if not local_path.exists():
+            if downloader is not None:
+                downloader(href, local_path)
+            elif href.startswith("s3://"):
+                fs = (
+                    s3_client if s3_client is not None else s3fs.S3FileSystem(anon=True)
+                )
+                fs.get(href.replace("s3://", ""), str(local_path))
+            elif href.startswith("http://") or href.startswith("https://"):
+                http = http_session if http_session is not None else requests
+                response = http.get(href, stream=True)
+                response.raise_for_status()
+                with open(local_path, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+            elif Path(href).exists():
+                if Path(href).absolute() != local_path.absolute():
+                    shutil.copy(href, local_path)
+            else:
+                raise FileNotFoundError(f"Source file not found at {href}")
+
+
+def download_assets_safely(
+    hrefs: list[str],
+    local_paths: list[Path],
+    s3_client: Optional[Any] = None,
+    http_session: Optional[Any] = None,
+    downloader: Optional["Downloader"] = None,
+    max_workers: Optional[int] = None,
+) -> None:
+    """Download multiple assets concurrently using a thread pool.
+
+    Args:
+        hrefs: List of URLs or local paths to the assets.
+        local_paths: List of destination paths for the downloaded files.
+        s3_client: Optional authenticated S3FileSystem for S3 access.
+        http_session: Optional authenticated requests.Session for HTTPS downloads.
+        downloader: Optional callable that handles the download itself.
+        max_workers: The maximum number of threads to use. Defaults to None.
+    """
+    if len(hrefs) != len(local_paths):
+        raise ValueError("hrefs and local_paths must have the same length")
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(
+                download_asset_safely,
+                href=href,
+                local_path=local_path,
+                s3_client=s3_client,
+                http_session=http_session,
+                downloader=downloader,
+            )
+            for href, local_path in zip(hrefs, local_paths, strict=True)
+        ]
+        for future in futures:
+            future.result()
+
+
+def extract_asset_safely(
+    archive_path: Path,
+    extract_dir: Optional[Path] = None,
+    lock_path: Optional[Path] = None,
+) -> None:
+    """Extract a zip archive safely with file locking.
+
+    Uses atomic extraction (temp directory + rename) so that other
+    processes never see a partially extracted directory.  An
+    ``.extracted`` marker file is written on success; if the marker
+    exists the function returns immediately.
+
+    Args:
+        archive_path: Path to the zip archive.
+        extract_dir: Destination directory.  Defaults to
+            ``archive_path.with_suffix("")``.
+        lock_path: Path to the lock file.  Defaults to
+            ``extract_dir.with_suffix(".lock")``.
+    """
+    import filelock
+    import shutil
+    import tempfile
+    import zipfile
+
+    archive_path = Path(archive_path)
+    if extract_dir is None:
+        extract_dir = archive_path.with_suffix("")
+    extract_dir = Path(extract_dir)
+
+    if lock_path is None:
+        lock_path = extract_dir.with_suffix(".lock")
+    lock_path = Path(lock_path)
+
+    marker_path = extract_dir.with_suffix(".extracted")
+
+    with filelock.FileLock(str(lock_path)):
+        # Already extracted and marked complete
+        if marker_path.exists() and extract_dir.exists():
+            return
+
+        # Remove any stale partial extraction
+        if extract_dir.exists():
+            shutil.rmtree(extract_dir, ignore_errors=True)
+
+        extract_dir.parent.mkdir(parents=True, exist_ok=True)
+
+        # Extract to a temporary directory so other processes never
+        # see an incomplete destination.
+        temp_dir = tempfile.mkdtemp(
+            prefix=extract_dir.name + "_tmp_",
+            dir=extract_dir.parent,
+        )
+        try:
+            with zipfile.ZipFile(archive_path, "r") as zf:
+                zf.extractall(temp_dir)
+
+            Path(temp_dir).rename(extract_dir)
+        except Exception:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            raise
+
+        marker_path.touch()
+
+
+def cleanup_asset_safely(
+    local_path: Path, chunk_id: Optional[int] = None, total_chunks: int = 1
+) -> None:
+    """Safely clean up the downloaded asset after all chunks are processed."""
+    import filelock
+
+    lock_path = local_path.with_suffix(".lock")
+    if total_chunks > 1 and chunk_id is not None:
+        done_file = local_path.with_suffix(f".chunk_{chunk_id}.done")
+        done_file.touch()
+        with filelock.FileLock(str(lock_path)):
+            done_files = list(local_path.parent.glob(f"{local_path.stem}.chunk_*.done"))
+            if len(done_files) >= total_chunks:
+                if local_path.exists():
+                    try:
+                        local_path.unlink()
+                    except Exception:
+                        pass
+                for df in done_files:
+                    try:
+                        df.unlink()
+                    except Exception:
+                        pass
+                try:
+                    lock_path.unlink()
+                except Exception:
+                    pass
+    else:
+        if local_path.exists():
+            try:
+                local_path.unlink()
+            except Exception:
+                pass

aereo/cli/main.py

@@ -0,0 +1,491 @@
+"""AER CLI — command-line interface for the AER satellite data framework."""
+
+# ruff: noqa: E402
+from __future__ import annotations
+
+import json
+import pickle
+from datetime import datetime
+from pathlib import Path
+from typing import Annotated, Any, Optional
+
+import geopandas as gpd
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from aereo.client import AereoClient
+from aereo.execution import LocalProcessBackend
+from aereo.interfaces import AereoProfile, GridConfig
+from aereo.schemas import AssetSchema
+
+app = typer.Typer(
+    name="aereo",
+    help="AER — Modular satellite data discovery, extraction, and processing",
+    no_args_is_help=True,
+)
+console = Console()
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_geometry(geojson_path: Path) -> Optional[dict[str, Any]]:
+    """Load a GeoJSON file and return the geometry dict."""
+    data = json.loads(geojson_path.read_text())
+    if data.get("type") == "FeatureCollection":
+        if not data.get("features"):
+            raise ValueError("GeoJSON FeatureCollection has no features.")
+        return data["features"][0]["geometry"]
+    elif data.get("type") == "Feature":
+        return data["geometry"]
+    elif "type" in data and data["type"] in (
+        "Point",
+        "MultiPoint",
+        "LineString",
+        "MultiLineString",
+        "Polygon",
+        "MultiPolygon",
+        "GeometryCollection",
+    ):
+        return data
+    raise ValueError("Could not extract geometry from GeoJSON.")
+
+
+def _search_results_to_json(df: gpd.GeoDataFrame) -> list[dict[str, Any]]:
+    """Convert search results GeoDataFrame to JSON-serializable records."""
+    # Convert to GeoJSON feature collection then to simple records
+    df = df.copy()
+    df["geometry"] = df.geometry.apply(
+        lambda g: g.__geo_interface__ if g is not None else None
+    )
+    records = df.to_dict(orient="records")
+    # Convert datetime to ISO strings
+    for rec in records:
+        for key in ("start_time", "end_time"):
+            if key in rec and rec[key] is not None:
+                rec[key] = (
+                    rec[key].isoformat() if hasattr(rec[key], "isoformat") else rec[key]
+                )
+    return records
+
+
+def _search_results_from_json(records: list[dict[str, Any]]) -> Any:
+    """Reconstruct a GeoDataFrame from JSON records."""
+    df = gpd.GeoDataFrame.from_records(records)
+    if "geometry" in df.columns:
+        from shapely.geometry import shape
+
+        def _to_geom(g: Any) -> Any:
+            if isinstance(g, dict):
+                return shape(g)
+            return g
+
+        df["geometry"] = gpd.GeoSeries(df["geometry"].apply(_to_geom))
+        df = gpd.GeoDataFrame(df, geometry="geometry")
+        df.set_crs(epsg=4326, inplace=True)
+    for key in ("start_time", "end_time"):
+        if key in df.columns:
+            df[key] = gpd.pd.to_datetime(df[key])
+    return gpd.GeoDataFrame(AssetSchema.validate(df))
+
+
+def _print_search_table(df: gpd.GeoDataFrame) -> None:
+    """Pretty-print search results as a Rich table."""
+    table = Table(title=f"Search Results ({len(df)} scenes)")
+    table.add_column("ID", style="cyan", no_wrap=True)
+    table.add_column("Collection", style="magenta")
+    table.add_column("Start Time", style="green")
+    table.add_column("End Time", style="green")
+    table.add_column("Href", style="blue", overflow="fold")
+
+    for _, row in df.head(50).iterrows():
+        table.add_row(
+            str(row.get("id", "")),
+            str(row.get("collection", "")),
+            str(row.get("start_time", "")),
+            str(row.get("end_time", "")),
+            str(row.get("href", ""))[:60] + "...",
+        )
+    if len(df) > 50:
+        table.add_row("...", f"... and {len(df) - 50} more rows", "", "", "")
+    console.print(table)
+
+
+# ---------------------------------------------------------------------------
+# Commands
+# ---------------------------------------------------------------------------
+
+
+@app.command()
+def search(
+    profile: Annotated[
+        list[Path],
+        typer.Option("--profile", "-p", help="Path to profile YAML (repeatable)"),
+    ],
+    config: Annotated[
+        Optional[Path], typer.Option("--config", "-c", help="Path to grid config YAML")
+    ] = None,
+    geojson: Annotated[
+        Optional[Path], typer.Option("--geojson", "-g", help="Path to AOI GeoJSON file")
+    ] = None,
+    start: Annotated[
+        Optional[str], typer.Option("--start", "-s", help="Start datetime (ISO 8601)")
+    ] = None,
+    end: Annotated[
+        Optional[str], typer.Option("--end", "-e", help="End datetime (ISO 8601)")
+    ] = None,
+    output: Annotated[
+        Optional[Path],
+        typer.Option("--output", "-o", help="Output JSON file for search results"),
+    ] = None,
+    fmt: Annotated[
+        str, typer.Option("--format", help="Output format: table or json")
+    ] = "table",
+    verbose: Annotated[
+        bool, typer.Option("--verbose", "-v", help="Enable verbose logging")
+    ] = False,
+) -> None:
+    """Search for satellite data across configured profiles."""
+    if verbose:
+        import structlog
+
+        structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(10))
+
+    # Load profiles
+    profiles: list[AereoProfile] = []
+    for p in profile:
+        if not p.exists():
+            console.print(f"[red]Profile not found:[/red] {p}")
+            raise typer.Exit(code=1)
+        profiles.extend(AereoProfile.from_yaml(p))
+
+    # Load geometry
+    intersects = _load_geometry(geojson) if geojson and geojson.exists() else None
+
+    # Parse datetimes
+    start_dt = datetime.fromisoformat(start) if start else None
+    end_dt = datetime.fromisoformat(end) if end else None
+
+    client = AereoClient()
+    try:
+        results = client.search(
+            profiles=profiles,
+            intersects=intersects,
+            start_datetime=start_dt,
+            end_datetime=end_dt,
+        )
+    except Exception as exc:
+        console.print(f"[red]Search failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+
+    if results is None or len(results) == 0:
+        console.print("[yellow]No results found.[/yellow]")
+        raise typer.Exit(code=2)
+
+    if fmt == "json":
+        records = _search_results_to_json(results)
+        json_out = json.dumps(records, indent=2, default=str)
+        if output:
+            output.write_text(json_out)
+            console.print(f"[green]Wrote {len(records)} results to[/green] {output}")
+        else:
+            console.print(json_out)
+    else:
+        _print_search_table(results)
+        if output:
+            records = _search_results_to_json(results)
+            output.write_text(json.dumps(records, indent=2, default=str))
+            console.print(f"[green]Wrote results to[/green] {output}")
+
+
+@app.command()
+def prepare(
+    search_results: Annotated[Path, typer.Argument(help="Path to search results JSON")],
+    profile: Annotated[
+        list[Path],
+        typer.Option("--profile", "-p", help="Path to profile YAML (repeatable)"),
+    ],
+    config: Annotated[
+        Optional[Path], typer.Option("--config", "-c", help="Path to grid config YAML")
+    ] = None,
+    output_dir: Annotated[
+        Path, typer.Option("--output-dir", "-d", help="Output directory for extraction")
+    ] = Path("./out"),
+    output: Annotated[
+        Optional[Path],
+        typer.Option("--output", "-o", help="Output pickle file for tasks"),
+    ] = None,
+    cells_per_chunk: Annotated[
+        int, typer.Option("--cells-per-chunk", help="Max grid cells per task")
+    ] = 50,
+    verbose: Annotated[
+        bool, typer.Option("--verbose", "-v", help="Enable verbose logging")
+    ] = False,
+) -> None:
+    """Prepare search results for extraction."""
+    if verbose:
+        import structlog
+
+        structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(10))
+
+    if not search_results.exists():
+        console.print(f"[red]Search results not found:[/red] {search_results}")
+        raise typer.Exit(code=1)
+
+    records = json.loads(search_results.read_text())
+    df = _search_results_from_json(records)
+
+    profiles: list[AereoProfile] = []
+    for p in profile:
+        if not p.exists():
+            console.print(f"[red]Profile not found:[/red] {p}")
+            raise typer.Exit(code=1)
+        profiles.extend(AereoProfile.from_yaml(p))
+
+    grid_config = (
+        GridConfig.from_yaml(config) if config and config.exists() else GridConfig()
+    )
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    client = AereoClient()
+    try:
+        tasks = client.prepare_for_extraction(
+            search_results=df,  # type: ignore[arg-type]
+            grid_config=grid_config,
+            profiles=profiles,
+            uri=str(output_dir),
+            cells_per_chunk=cells_per_chunk,
+        )
+    except Exception as exc:
+        console.print(f"[red]Prepare failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+
+    task_file = output or (output_dir / "tasks.pkl")
+    task_file.write_bytes(pickle.dumps(tasks))
+    console.print(
+        f"[green]✓ Prepared {len(tasks)} tasks (chunk size: {cells_per_chunk}).[/green]"
+    )
+    console.print(f"[green]Wrote tasks to[/green] {task_file}")
+
+
+@app.command()
+def extract(
+    tasks: Annotated[Path, typer.Argument(help="Path to prepared tasks pickle file")],
+    output_dir: Annotated[
+        Path, typer.Option("--output-dir", "-d", help="Output directory")
+    ] = Path("./out"),
+    workers: Annotated[
+        int, typer.Option("--workers", "-w", help="Max batch workers")
+    ] = 1,
+    verbose: Annotated[
+        bool, typer.Option("--verbose", "-v", help="Enable verbose logging")
+    ] = False,
+) -> None:
+    """Run extraction on prepared tasks."""
+    if verbose:
+        import structlog
+
+        structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(10))
+
+    if not tasks.exists():
+        console.print(f"[red]Tasks file not found:[/red] {tasks}")
+        raise typer.Exit(code=1)
+
+    task_list = pickle.loads(tasks.read_bytes())
+
+    backend = LocalProcessBackend(max_workers=workers)
+    client = AereoClient()
+    try:
+        artifacts = client.execute_tasks(task_list, backend=backend)
+    except Exception as exc:
+        console.print(f"[red]Extraction failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+
+    # Write full GeoDataFrame to parquet
+    parquet_path = output_dir / "artifacts.parquet"
+    artifacts.to_parquet(parquet_path)
+
+    console.print(f"[green]✓ Extracted {len(artifacts)} artifacts.[/green]")
+    console.print(f"[green]Parquet saved to:[/green] {parquet_path}")
+    console.print(f"[green]Output directory:[/green] {output_dir}")
+
+
+@app.command()
+def run(
+    profile: Annotated[
+        list[Path],
+        typer.Option("--profile", "-p", help="Path to profile YAML (repeatable)"),
+    ],
+    config: Annotated[
+        Optional[Path], typer.Option("--config", "-c", help="Path to grid config YAML")
+    ] = None,
+    geojson: Annotated[
+        Optional[Path], typer.Option("--geojson", "-g", help="Path to AOI GeoJSON file")
+    ] = None,
+    start: Annotated[
+        Optional[str], typer.Option("--start", "-s", help="Start datetime (ISO 8601)")
+    ] = None,
+    end: Annotated[
+        Optional[str], typer.Option("--end", "-e", help="End datetime (ISO 8601)")
+    ] = None,
+    output_dir: Annotated[
+        Path, typer.Option("--output-dir", "-d", help="Output directory")
+    ] = Path("./out"),
+    workers: Annotated[
+        int, typer.Option("--workers", "-w", help="Max batch workers")
+    ] = 1,
+    cells_per_chunk: Annotated[
+        int, typer.Option("--cells-per-chunk", help="Max grid cells per task")
+    ] = 50,
+    verbose: Annotated[
+        bool, typer.Option("--verbose", "-v", help="Enable verbose logging")
+    ] = False,
+) -> None:
+    """One-shot: search → prepare → extract."""
+    if verbose:
+        import structlog
+
+        structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(10))
+
+    # Load profiles
+    profiles: list[AereoProfile] = []
+    for p in profile:
+        if not p.exists():
+            console.print(f"[red]Profile not found:[/red] {p}")
+            raise typer.Exit(code=1)
+        profiles.extend(AereoProfile.from_yaml(p))
+
+    grid_config = (
+        GridConfig.from_yaml(config) if config and config.exists() else GridConfig()
+    )
+    intersects = _load_geometry(geojson) if geojson and geojson.exists() else None
+    start_dt = datetime.fromisoformat(start) if start else None
+    end_dt = datetime.fromisoformat(end) if end else None
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    client = AereoClient()
+
+    # Search
+    console.print("[bold blue]🔍 Searching...[/bold blue]")
+    try:
+        results = client.search(
+            profiles=profiles,
+            intersects=intersects,
+            start_datetime=start_dt,
+            end_datetime=end_dt,
+        )
+    except Exception as exc:
+        console.print(f"[red]Search failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+
+    if results is None or len(results) == 0:
+        console.print("[yellow]No results found.[/yellow]")
+        raise typer.Exit(code=2)
+    console.print(f"[green]✓ Found {len(results)} scenes.[/green]")
+
+    # Prepare
+    console.print("[bold blue]📦 Preparing...[/bold blue]")
+    try:
+        tasks = client.prepare_for_extraction(
+            search_results=results,
+            grid_config=grid_config,
+            profiles=profiles,
+            uri=str(output_dir),
+            cells_per_chunk=cells_per_chunk,
+        )
+    except Exception as exc:
+        console.print(f"[red]Prepare failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+    console.print(
+        f"[green]✓ Prepared {len(tasks)} tasks (chunk size: {cells_per_chunk}).[/green]"
+    )
+
+    # Extract
+    console.print("[bold blue]⛏️ Extracting...[/bold blue]")
+    backend = LocalProcessBackend(max_workers=workers)
+    try:
+        artifacts = client.execute_tasks(tasks, backend=backend)
+    except Exception as exc:
+        console.print(f"[red]Extraction failed:[/red] {exc}")
+        raise typer.Exit(code=1)
+
+    parquet_path = output_dir / "artifacts.parquet"
+    artifacts.to_parquet(parquet_path)
+
+    console.print(f"[green]✓ Extracted {len(artifacts)} artifacts.[/green]")
+    console.print(f"[green]Parquet saved to:[/green] {parquet_path}")
+    console.print(f"[green]Output:[/green] {output_dir}")
+
+
+@app.command()
+def validate(
+    config: Annotated[
+        Optional[Path],
+        typer.Option("--config", "-c", help="Path to config YAML to validate"),
+    ] = None,
+    profile: Annotated[
+        Optional[Path],
+        typer.Option("--profile", "-p", help="Path to profile YAML to validate"),
+    ] = None,
+) -> None:
+    """Validate a config or profile YAML against AER schemas."""
+    if config:
+        if not config.exists():
+            console.print(f"[red]Config not found:[/red] {config}")
+            raise typer.Exit(code=1)
+        try:
+            GridConfig.from_yaml(config)
+            console.print(f"[green]✓ Config valid:[/green] {config}")
+        except Exception as exc:
+            console.print(f"[red]✗ Config invalid:[/red] {config}\n{exc}")
+            raise typer.Exit(code=1)
+
+    if profile:
+        if not profile.exists():
+            console.print(f"[red]Profile not found:[/red] {profile}")
+            raise typer.Exit(code=1)
+        try:
+            AereoProfile.from_yaml(profile)
+            console.print(f"[green]✓ Profile valid:[/green] {profile}")
+        except Exception as exc:
+            console.print(f"[red]✗ Profile invalid:[/red] {profile}\n{exc}")
+            raise typer.Exit(code=1)
+
+    if not config and not profile:
+        console.print("[yellow]Provide --config or --profile to validate.[/yellow]")
+        raise typer.Exit(code=1)
+
+
+@app.command()
+def plugins() -> None:
+    """List installed AER plugins."""
+    from aereo.registry import AereoRegistry
+
+    registry = AereoRegistry()
+
+    table = Table(title="Installed AER Plugins")
+    table.add_column("Type", style="cyan")
+    table.add_column("Name", style="magenta")
+    table.add_column("Collections", style="green")
+
+    for name, cls in registry._searchers.items():
+        cols = ", ".join(cls.supported_collections[:3])
+        if len(cls.supported_collections) > 3:
+            cols += " ..."
+        table.add_row("Searcher", name, cols)
+
+    for name, cls in registry._extractors.items():
+        cols = ", ".join(cls.supported_collections[:3])
+        if len(cls.supported_collections) > 3:
+            cols += " ..."
+        table.add_row("Extractor", name, cols)
+
+    console.print(table)
+
+
+# Entry point for `python -m aer.cli`
+if __name__ == "__main__":
+    app()