terrakio-core 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

terrakio_core/__init__.py

@@ -1,3 +1,4 @@
+# terrakio_core/__init__.py
 """
 Terrakio Core
 
@@ -5,3 +6,11 @@ Core components for Terrakio API clients.
 """
 
 __version__ = "0.3.4"
+
+from .async_client import AsyncClient
+from .sync_client import SyncClient
+
+__all__ = [
+    "AsyncClient", 
+    "SyncClient"
+]

terrakio_core/async_client.py

@@ -0,0 +1,304 @@
+import aiohttp
+import asyncio
+import json
+import pandas as pd
+import xarray as xr
+from io import BytesIO
+from typing import Optional, Dict, Any, Union
+from geopandas import GeoDataFrame
+from shapely.geometry.base import BaseGeometry as ShapelyGeometry
+from shapely.geometry import mapping
+from .client import BaseClient
+from .exceptions import APIError
+from .endpoints.dataset_management import DatasetManagement
+from .endpoints.user_management import UserManagement
+from .endpoints.mass_stats import MassStats
+from .endpoints.group_management import GroupManagement
+from .endpoints.space_management import SpaceManagement
+from .endpoints.model_management import ModelManagement
+from .endpoints.auth import AuthClient
+from .convenience_functions.convenience_functions import zonal_stats as _zonal_stats, create_dataset_file as _create_dataset_file
+
+class AsyncClient(BaseClient):
+    def __init__(self, url: Optional[str] = None, api_key: Optional[str] = None, verbose: bool = False, session: Optional[aiohttp.ClientSession] = None):
+        super().__init__(url, api_key, verbose)
+        self.datasets = DatasetManagement(self)
+        self.users = UserManagement(self)
+        self.mass_stats = MassStats(self)
+        self.groups = GroupManagement(self)
+        self.space = SpaceManagement(self)
+        self.model = ModelManagement(self)
+        self.auth = AuthClient(self)
+        self._session = session
+        self._owns_session = session is None
+
+    async def _terrakio_request(self, method: str, endpoint: str, **kwargs):
+        # if self.session is None:
+        if self.session is None:
+            headers = {
+                'Content-Type': 'application/json',
+                'x-api-key': self.key,
+                'Authorization': self.token
+            }
+            clean_headers = {k: v for k, v in headers.items() if v is not None}
+            async with aiohttp.ClientSession(headers=clean_headers, timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
+                return await self._make_request_with_retry(session, method, endpoint, **kwargs)
+        else:
+            return await self._make_request_with_retry(self._session, method, endpoint, **kwargs)
+
+
+    async def _make_request_with_retry(self, session: aiohttp.ClientSession, method: str, endpoint: str, **kwargs) -> Dict[Any, Any]:
+        url = f"{self.url}/{endpoint.lstrip('/')}"
+        for attempt in range(self.retry + 1):
+            try:                
+                async with session.request(method, url, **kwargs) as response:
+                    response_text = await response.text()
+                    if not response.ok:
+                        should_retry = False
+                        
+                        if response.status == 400:
+                            should_retry = False
+                        else:
+                            if response.status in [408, 502, 503, 504]:
+                                should_retry = True
+                            elif response.status == 500:
+                                try:
+                                    response_text = await response.text()
+                                    if "Internal server error" not in response_text:
+                                        should_retry = True
+                                except:
+                                    should_retry = True
+                        
+                        if should_retry and attempt < self.retry:
+                            self.logger.info(f"Request failed (attempt {attempt+1}/{self.retry+1}): {response.status} {response.reason}. Retrying...")
+                            continue
+                        else:
+                            error_msg = f"API request failed: {response.status} {response.reason}"
+                            try:
+                                error_data = await response.json()
+                                if "detail" in error_data:
+                                    error_msg += f" - {error_data['detail']}"
+                            except:
+                                pass
+                            raise APIError(error_msg, status_code=response.status)
+                    
+                    content_type = response.headers.get('content-type', '').lower()
+                    content = await response.read()
+                    if 'json' in content_type:
+                        return json.loads(content.decode('utf-8'))
+                    elif 'csv' in content_type:
+                        return pd.read_csv(BytesIO(content))
+                    elif 'image/' in content_type:
+                        return content
+                    elif 'text' in content_type:
+                        return content.decode('utf-8')
+                    else:
+                        try:
+                            return xr.open_dataset(BytesIO(content))
+                        except:
+                            raise APIError(f"Unknown response format. Content-Type: {response.headers.get('content-type', 'unknown')}", status_code=response.status)
+            except aiohttp.ClientError as e:
+                if attempt < self.retry:
+                    self.logger.info(f"Request failed (attempt {attempt+1}/{self.retry+1}): {e}. Retrying...")
+                    continue
+                else:
+                    raise APIError(f"Request failed after {self.retry+1} attempts: {e}", status_code=None)
+                
+
+    async def _regular_request(self, method: str, endpoint: str, **kwargs):
+        url = endpoint.lstrip('/')
+        if self._session is None:
+
+            async with aiohttp.ClientSession() as session:
+                try:
+                    async with session.request(method, url, **kwargs) as response:
+                        response.raise_for_status()
+                        return response
+                except aiohttp.ClientError as e:
+                    raise APIError(f"Request failed: {e}")
+        else:
+            # this means that we used the with statement, and we already have a session
+            try:
+                async with self._session.request(method, url, **kwargs) as response:
+                    response.raise_for_status()
+                    return response
+            except aiohttp.ClientError as e:
+                raise APIError(f"Request failed: {e}")
+        
+    async def geoquery(
+        self,
+        expr: str,
+        feature: Union[Dict[str, Any], ShapelyGeometry],
+        in_crs: str = "epsg:4326",
+        out_crs: str = "epsg:4326",
+        output: str = "csv",
+        resolution: int = -1,
+        geom_fix: bool = False,
+        **kwargs
+    ):
+        """
+        Compute WCS query for a single geometry.
+
+        Args:
+            expr (str): The WCS expression to evaluate
+            feature (Union[Dict[str, Any], ShapelyGeometry]): The geographic feature
+            in_crs (str): Input coordinate reference system
+            out_crs (str): Output coordinate reference system
+            output (str): Output format ('csv' or 'netcdf')
+            resolution (int): Resolution parameter
+            geom_fix (bool): Whether to fix the geometry (default False)
+            **kwargs: Additional parameters to pass to the WCS request
+            
+        Returns:
+            Union[pd.DataFrame, xr.Dataset, bytes]: The response data in the requested format
+
+        Raises:
+            APIError: If the API request fails
+        """
+        if hasattr(feature, 'is_valid'):
+            feature = {
+                "type": "Feature",
+                "geometry": mapping(feature),
+                "properties": {}
+            }
+        payload = {
+            "feature": feature,
+            "in_crs": in_crs,
+            "out_crs": out_crs,
+            "output": output,
+            "resolution": resolution,
+            "expr": expr,
+            "buffer": geom_fix,
+            **kwargs
+        }
+        return await self._terrakio_request("POST", "geoquery", json=payload)
+
+    async def zonal_stats(
+            self,
+            gdf: GeoDataFrame,
+            expr: str,
+            conc: int = 20,
+            inplace: bool = False,
+            in_crs: str = "epsg:4326",
+            out_crs: str = "epsg:4326",
+            resolution: int = -1,
+            geom_fix: bool = False,
+    ):
+        """
+        Compute zonal statistics for all geometries in a GeoDataFrame.
+
+        Args:
+            gdf (GeoDataFrame): GeoDataFrame containing geometries
+            expr (str): Terrakio expression to evaluate, can include spatial aggregations
+            conc (int): Number of concurrent requests to make
+            inplace (bool): Whether to modify the input GeoDataFrame in place
+            in_crs (str): Input coordinate reference system
+            out_crs (str): Output coordinate reference system
+            resolution (int): Resolution parameter
+            geom_fix (bool): Whether to fix the geometry (default False)
+        Returns:
+            geopandas.GeoDataFrame: GeoDataFrame with added columns for results, or None if inplace=True
+
+        Raises:
+            ValueError: If concurrency is too high
+            APIError: If the API request fails
+        """
+        return await _zonal_stats(
+            client=self,
+            gdf=gdf,
+            expr=expr,
+            conc=conc,
+            inplace=inplace,
+            in_crs=in_crs,
+            out_crs=out_crs,
+            resolution=resolution,
+            geom_fix=geom_fix
+        )
+
+    async def create_dataset_file(
+        self,
+        aoi: str,
+        expression: str,
+        output: str,
+        in_crs: str = "epsg:4326",
+        res: float = 0.0001,
+        region: str = "aus",
+        to_crs: str = "epsg:4326",
+        overwrite: bool = True,
+        skip_existing: bool = False,
+        non_interactive: bool = True,
+        poll_interval: int = 30,
+        download_path: str = "/home/user/Downloads",
+    ) -> dict:
+        """
+        Create a dataset file using mass stats operations.
+
+        Args:
+            aoi (str): Area of interest
+            expression (str): Terrakio expression to evaluate
+            output (str): Output format
+            in_crs (str): Input coordinate reference system (default "epsg:4326")
+            res (float): Resolution (default 0.0001)
+            region (str): Region (default "aus")
+            to_crs (str): Target coordinate reference system (default "epsg:4326")
+            overwrite (bool): Whether to overwrite existing files (default True)
+            skip_existing (bool): Whether to skip existing files (default False)
+            non_interactive (bool): Whether to run non-interactively (default True)
+            poll_interval (int): Polling interval in seconds (default 30)
+            download_path (str): Download path (default "/home/user/Downloads")
+
+        Returns:
+            dict: Dictionary containing generation_task_id and combine_task_id
+
+        Raises:
+            ConfigurationError: If mass stats client is not properly configured
+            RuntimeError: If job fails
+        """
+        return await _create_dataset_file(
+            client=self,
+            aoi=aoi,
+            expression=expression,
+            output=output,
+            in_crs=in_crs,
+            res=res,
+            region=region,
+            to_crs=to_crs,
+            overwrite=overwrite,
+            skip_existing=skip_existing,
+            non_interactive=non_interactive,
+            poll_interval=poll_interval,
+            download_path=download_path,
+        )
+
+    async def __aenter__(self):
+        # if there is no session, we create a session
+        if self._session is None:
+            headers = {
+                'Content-Type': 'application/json',
+                'x-api-key': self.key,
+                'Authorization': self.token
+            }
+            # Remove None values from headers
+            clean_headers = {k: v for k, v in headers.items() if v is not None}
+            # we are creating the header and clean any value that is none
+            # now we create the session
+            self._session = aiohttp.ClientSession(
+                headers=clean_headers,
+                timeout=aiohttp.ClientTimeout(total=self.timeout)
+            )
+        return self
+    # if there is no session, we create a session
+
+    # now lets create the aexit function, this function is used when a user uses with, and this function will be automatically called when the with statement is done
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        #  so if the session is not being passed in, and we created it by ourselves, we are responsible for closing the session
+        #  if the session is being passed in, we are not responsible for closing the session
+        if self._owns_session and self._session:
+            await self._session.close()
+            self._session = None
+            # we close the session and set the session value to none
+
+    async def close(self):
+        if self._owns_session and self._session:
+            await self._session.close()
+            self._session = None