terrakio-core 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

terrakio_core/__init__.py

@@ -5,7 +5,7 @@ Terrakio Core
 Core components for Terrakio API clients.
 """
 
-__version__ = "0.4.5"
+__version__ = "0.4.7"
 
 from .async_client import AsyncClient
 from .sync_client import SyncClient as Client

terrakio_core/async_client.py

@@ -196,6 +196,8 @@ class AsyncClient(BaseClient):
             out_crs: str = "epsg:4326",
             resolution: int = -1,
             geom_fix: bool = False,
+            mass_stats: bool = False,
+            id_column: Optional[str] = None,
     ):
         """
         Compute zonal statistics for all geometries in a GeoDataFrame.
@@ -208,6 +210,8 @@ class AsyncClient(BaseClient):
             out_crs (str): Output coordinate reference system
             resolution (int): Resolution parameter
             geom_fix (bool): Whether to fix the geometry (default False)
+            mass_stats (bool): Whether to use mass stats for processing (default False)
+            id_column (Optional[str]): Name of the ID column to use (default None)
 
         Returns:
             geopandas.GeoDataFrame: GeoDataFrame with added columns for results
@@ -225,10 +229,13 @@ class AsyncClient(BaseClient):
             out_crs=out_crs,
             resolution=resolution,
             geom_fix=geom_fix,
+            mass_stats=mass_stats,
+            id_column=id_column,
         )
 
     async def create_dataset_file(
         self,
+        name: str,
         aoi: str,
         expression: str,
         output: str,
@@ -241,6 +248,9 @@ class AsyncClient(BaseClient):
         non_interactive: bool = True,
         poll_interval: int = 30,
         download_path: str = "/home/user/Downloads",
+        mask = True,
+        max_file_size_mb: int = 5120,  # Default to 5GB
+        tile_size: int = 1024,
     ) -> dict:
         """
         Create a dataset file using mass stats operations.
@@ -280,6 +290,10 @@ class AsyncClient(BaseClient):
             non_interactive=non_interactive,
             poll_interval=poll_interval,
             download_path=download_path,
+            name=name,
+            mask=mask,
+            max_file_size_mb=max_file_size_mb,
+            tile_size=tile_size
         )
 
     async def geo_queries(

terrakio_core/convenience_functions/convenience_functions.py

@@ -24,6 +24,8 @@ import pyproj
 import pandas as pd
 import geopandas as gpd
 
+from typing import Optional
+
 def expand_on_time(gdf):
     """
     Expand datasets on time dimension - each time becomes a new row.
@@ -359,18 +361,119 @@ async def local_or_remote(
             "local_or_remote": "local",
             "reason": "The number of the requests is not too large, and the time taking for making these requests is not too long, and the size of the dataset is not too large",
         }
+    
+def gdf_to_json(
+    gdf: GeoDataFrame,
+    expr: str,
+    in_crs: str = "epsg:4326",
+    out_crs: str = "epsg:4326",
+    resolution: int = -1,
+    geom_fix: bool = False,
+    id_column: Optional[str] = None,
+):
+    """
+    Convert a GeoDataFrame to a list of JSON requests for mass_stats processing.
+    
+    Args:
+        gdf: GeoDataFrame containing geometries and optional metadata
+        expr: Expression to evaluate
+        in_crs: Input coordinate reference system
+        out_crs: Output coordinate reference system
+        resolution: Resolution parameter
+        geom_fix: Whether to fix geometry issues
+        id_column: Optional column name to use for group and file names
+        
+    Returns:
+        list: List of dictionaries formatted for mass_stats requests
+    """
+    mass_stats_requests = []
+    
+    # Loop through each row in the GeoDataFrame
+    for idx, row in gdf.iterrows():
+        # Create the request feature
+        request_feature = {
+            "expr": expr,
+            "feature": {
+                "type": "Feature",
+                "geometry": mapping(gdf.geometry.iloc[idx]),
+                "properties": {}
+            },
+            "in_crs": in_crs,
+            "out_crs": out_crs,
+            "resolution": resolution,
+            "geom_fix": geom_fix,
+        }
+        
+        # Determine group name and file name based on id_column
+        if id_column is not None and id_column in gdf.columns:
+            # Use the value from the specified column as group and file name
+            identifier = str(row[id_column])
+            group_name = f"group_{identifier}"
+            file_name = f"file_{identifier}"
+        else:
+            # Use the index as group and file name
+            group_name = f"group_{idx}"
+            file_name = f"file_{idx}"
+            
+        # Create the complete request entry
+        request_entry = {
+            "group": group_name,
+            "file": file_name,
+            "request": request_feature,
+        }
+        
+        # Add the request to our list
+        mass_stats_requests.append(request_entry)
+        
+    return mass_stats_requests
+    
+async def handle_mass_stats(
+    client,
+    gdf: GeoDataFrame,
+    expr: str,
+    in_crs: str = "epsg:4326",
+    out_crs: str = "epsg:4326",
+    resolution: int = -1,
+    geom_fix: bool = False,
+    id_column: Optional[str] = None,
+
+):
+    request_json = gdf_to_json(gdf = gdf, expr = expr, in_crs = in_crs, out_crs = out_crs, resolution = resolution, geom_fix = geom_fix, id_column = id_column)
+    job_id =await client.mass_stats.execute_job(
+        name = "zonal_stats_job",
+        output = "netcdf",
+        config = {},
+        request_json = request_json,
+        overwrite = True,
+    )
+    return job_id
 
 async def zonal_stats(
-        client,
-        gdf: GeoDataFrame,
-        expr: str,
-        conc: int = 20,
-        in_crs: str = "epsg:4326",
-        out_crs: str = "epsg:4326",
-        resolution: int = -1,
-        geom_fix: bool = False,
+    client,
+    gdf: GeoDataFrame,
+    expr: str,
+    conc: int = 20,
+    in_crs: str = "epsg:4326",
+    out_crs: str = "epsg:4326",
+    resolution: int = -1,
+    geom_fix: bool = False,
+    mass_stats: bool = False,
+    id_column: Optional[str] = None,
 ):
     """Compute zonal statistics for all geometries in a GeoDataFrame."""
+
+    if mass_stats:
+        mass_stats_id = await handle_mass_stats(
+            client = client,
+            gdf = gdf,
+            expr = expr,
+            in_crs = in_crs,
+            out_crs = out_crs,
+            resolution = resolution,
+            geom_fix = geom_fix,
+            id_column = id_column
+        )
+        return mass_stats_id
     quries = []
     for i in range(len(gdf)):
         quries.append({
@@ -385,6 +488,7 @@ async def zonal_stats(
             "resolution": resolution,
             "geom_fix": geom_fix,
         })
+
     local_or_remote_result = await local_or_remote(client= client, quries = quries)
     if local_or_remote_result["local_or_remote"] == "remote":
         raise ValueError(local_or_remote_result["reason"])
@@ -398,30 +502,35 @@ async def create_dataset_file(
     aoi: str,
     expression: str,
     output: str,
+    download_path: str,
     in_crs: str = "epsg:4326",
-    res: float = 0.0001,
-    region: str = "aus",
     to_crs: str = "epsg:4326",
-    overwrite: bool = True,
+    res: float = 0.0001,
+    region: str = None,
+    overwrite: bool = False,
     skip_existing: bool = False,
     non_interactive: bool = True,
+    name: str | None = None,
     poll_interval: int = 30,
-    download_path: str = "/home/user/Downloads",
+    max_file_size_mb: int = 5120,
+    tile_size: int = 1024,
+    mask: bool = True
 ) -> dict:
     
-    name = f"tiles-{uuid.uuid4().hex[:8]}"
+    if not name:
+        name = f"file-gen-{uuid.uuid4().hex[:8]}"
     
     body, reqs, groups = tiles(
         name = name, 
         aoi = aoi, 
         expression = expression,
         output = output,
-        tile_size = 128,
+        tile_size = tile_size,
         crs = in_crs,
         res = res,
         region = region,
         to_crs = to_crs,
-        fully_cover = True,
+        mask = mask,
         overwrite = overwrite,
         skip_existing = skip_existing,
         non_interactive = non_interactive
@@ -429,9 +538,6 @@ async def create_dataset_file(
     with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempreq:
         tempreq.write(reqs)
         tempreqname = tempreq.name
-    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempmanifest:
-        tempmanifest.write(groups)
-        tempmanifestname = tempmanifest.name
 
     task_id = await client.mass_stats.execute_job(
         name=body["name"],
@@ -441,28 +547,24 @@ async def create_dataset_file(
         overwrite=body["overwrite"],
         skip_existing=body["skip_existing"],
         request_json=tempreqname,
-        manifest_json=tempmanifestname,
     )
 
     start_time = time.time()
     status = None
-    
+    client.logger.info(f"Tracking data generation job {task_id['task_id']}...")
     while True:
         try:
             taskid = task_id['task_id']
             trackinfo = await client.mass_stats.track_job([taskid])
-            client.logger.info("the trackinfo is: ", trackinfo)
             status = trackinfo[taskid]['status']
-            
             if status == 'Completed':
-                client.logger.info('Tiles generated successfully!')
+                client.logger.info('Data generated successfully!')
                 break
             elif status in ['Failed', 'Cancelled', 'Error']:
                 raise RuntimeError(f"Job {taskid} failed with status: {status}")
             else:
                 elapsed_time = time.time() - start_time
-                client.logger.info(f"Job status: {status} - Elapsed time: {elapsed_time:.1f}s", end='\r')
-                
+                client.logger.info(f"Job status: {status} - Elapsed time: {elapsed_time:.1f}s")
                 await asyncio.sleep(poll_interval)
                 
                 
@@ -474,16 +576,15 @@ async def create_dataset_file(
             raise
 
     os.unlink(tempreqname)
-    os.unlink(tempmanifestname)
 
-    combine_result = await client.mass_stats.combine_tiles(body["name"], body["overwrite"], body["output"])
+    combine_result = await client.mass_stats.combine_tiles(body["name"], body["overwrite"], body["output"], max_file_size_mb=max_file_size_mb)
     combine_task_id = combine_result.get("task_id")
 
     combine_start_time = time.time()
+    client.logger.info(f"Tracking file generation job {combine_task_id}...")
     while True:
         try:
             trackinfo = await client.mass_stats.track_job([combine_task_id])
-            client.logger.info('client create dataset file track info:', trackinfo)
             if body["output"] == "netcdf":
                 download_file_name = trackinfo[combine_task_id]['folder'] + '.nc'
             elif body["output"] == "geotiff":
@@ -491,19 +592,19 @@ async def create_dataset_file(
             bucket = trackinfo[combine_task_id]['bucket']
             combine_status = trackinfo[combine_task_id]['status']
             if combine_status == 'Completed':
-                client.logger.info('Tiles combined successfully!')
+                client.logger.info('File/s generated successfully!')
                 break
             elif combine_status in ['Failed', 'Cancelled', 'Error']:
-                raise RuntimeError(f"Combine job {combine_task_id} failed with status: {combine_status}")
+                raise RuntimeError(f"File generation job {combine_task_id} failed with status: {combine_status}")
             else:
                 elapsed_time = time.time() - combine_start_time
-                client.logger.info(f"Combine job status: {combine_status} - Elapsed time: {elapsed_time:.1f}s", end='\r')
+                client.logger.info(f"File generation job status: {combine_status} - Elapsed time: {elapsed_time:.1f}s")
                 time.sleep(poll_interval)
         except KeyboardInterrupt:
-            client.logger.info(f"\nInterrupted! Combine job {combine_task_id} is still running in the background.")
+            client.logger.info(f"\nInterrupted! File generation job {combine_task_id} is still running in the background.")
             raise
         except Exception as e:
-            client.logger.info(f"\nError tracking combine job: {e}")
+            client.logger.info(f"\nError tracking file generation job: {e}")
             raise
 
     if download_path:
@@ -511,11 +612,12 @@ async def create_dataset_file(
             job_name=body["name"],
             bucket=bucket,
             file_type='processed',
-            page_size=10,
+            folder='file-gen',
+            page_size=100,
             output_path=download_path,
         )
     else:
         path = f"{body['name']}/outputs/merged/{download_file_name}"
-        client.logger.info(f"Combined file is available at {path}")
+        client.logger.info(f"Dataset file/s is available at {path}")
 
     return {"generation_task_id": task_id, "combine_task_id": combine_task_id}

terrakio_core/endpoints/mass_stats.py

@@ -7,6 +7,13 @@ from urllib.parse import urlparse
 from ..helper.decorators import require_token, require_api_key, require_auth
 import aiohttp
 from typing import Dict, Any, Optional, List, Union
+import asyncio
+import xarray as xr
+from io import BytesIO
+import geopandas as gpd
+from shapely.geometry import shape
+from ..convenience_functions.convenience_functions import expand_on_variables_and_time
+
 class MassStats:
     def __init__(self, client):
         self._client = client
@@ -19,6 +26,7 @@ class MassStats:
         sample: str,
         output: str,
         config: Dict[str, Any],
+        region: str = None,
         overwrite: bool = False,
         skip_existing: bool = False,
         location: Optional[str] = None,
@@ -55,7 +63,8 @@ class MassStats:
             "config": config,
             "overwrite": overwrite,
             "skip_existing": skip_existing,
-            "server": server
+            "server": server,
+            "region": region
         }
         payload_mapping = {
             "location": location,
@@ -66,7 +75,6 @@ class MassStats:
                 payload[key] = str(value).lower()
         return await self._client._terrakio_request("POST", "mass_stats/upload", json=payload)
 
-
     @require_api_key
     async def start_job(self, id: str) -> Dict[str, Any]:
         """
@@ -276,6 +284,7 @@ class MassStats:
                         bucket: str,
                         file_type: str,
                         output_path: str,
+                        folder: str = None,
                         page_size: int = None,
                       ) -> list:
         """
@@ -303,7 +312,8 @@ class MassStats:
         request_body = {
             "job_name": job_name,
             "bucket": bucket,
-            "file_type": file_type
+            "file_type": file_type,
+            "folder": folder
         }
 
         output_dir = Path(output_path)
@@ -311,8 +321,7 @@ class MassStats:
         output_files = []
 
         async def download_urls_batch(download_urls, session):
-            for url in download_urls:
-                self._client.logger.info(f"Processing download URL: {url}")
+            for i, url in enumerate(download_urls):
                 parsed = urlparse(url)
                 path_parts = Path(parsed.path).parts
                 try:
@@ -322,7 +331,7 @@ class MassStats:
                     subpath = Path(path_parts[-1])
                 file_save_path = output_dir / subpath
                 file_save_path.parent.mkdir(parents=True, exist_ok=True)
-                self._client.logger.info(f"Downloading file to {file_save_path}")
+                self._client.logger.info(f"Downloading file to {file_save_path} ({i+1}/{len(download_urls)})")
 
                 async with session.get(url) as resp:
                     resp.raise_for_status()
@@ -352,7 +361,6 @@ class MassStats:
                     response = await self._client._terrakio_request("POST", "mass_stats/download_files", json=request_body, params=params)
                     data = response
 
-                    self._client.logger.info(f'processed, endpoint response is {data}')
                     download_urls = data.get('download_urls')
                     if not download_urls:
                         break
@@ -389,18 +397,16 @@ class MassStats:
                 raise ValueError("Request must be a dictionary")
             if not isinstance(request["file"], (str, int, list)):
                 raise ValueError("'file' must be a string or a list of strings")
-            # Only check the first 3 requests
             if i == 3:
                 break
 
-    @require_api_key
     async def execute_job(
         self, 
         name: str, 
-        # region: str, 
         output: str, 
         config: Dict[str, Any],
         request_json: str,  # Path to request JSON file
+        region: str = None, 
         overwrite: bool = False, 
         skip_existing: bool = False,  
         location: str = None, 
@@ -474,7 +480,7 @@ class MassStats:
         upload_result = await self._upload_request(
             name=name, 
             size=size, 
-            # region=region, 
+            region=region, 
             sample = first_expression,
             output=output, 
             config=config, 
@@ -557,7 +563,6 @@ class MassStats:
         tile_size: int,
         res: float,
         output: str,
-        # region: str,
         year_range: list[int] = None,
         overwrite: bool = False,
         server: str = None,
@@ -600,7 +605,6 @@ class MassStats:
         payload_mapping = {
             "year_range": year_range,
             "server": server,
-            # "region": region,
             "bucket": bucket,
         }
         for key, value in payload_mapping.items():
@@ -630,7 +634,7 @@ class MassStats:
         return self._client._terrakio_request("POST", "pyramids/create", json=payload)
     
     @require_api_key
-    async def combine_tiles(self, data_name: str, overwrite: bool = True, output: str = "netcdf") -> Dict[str, Any]:
+    async def combine_tiles(self, data_name: str, overwrite: bool = True, output: str = "netcdf", max_file_size_mb = 5120) -> Dict[str, Any]:
         """
         Combine tiles for a dataset.
 
@@ -647,7 +651,130 @@ class MassStats:
         """
         payload = {
             'data_name': data_name,
+            'folder': "file-gen",
             'output': output,
-            'overwrite': str(overwrite).lower()
+            'overwrite': str(overwrite).lower(),
+            'max_file_size_mb': max_file_size_mb
         }
-        return await self._client._terrakio_request("POST", "mass_stats/combine_tiles", json=payload)
+        return await self._client._terrakio_request("POST", "mass_stats/combine_tiles", json=payload)
+
+    @require_api_key
+    async def load_zonal_stats(self, job_id: str, max_files: int = 5, poll_interval: int = 30):
+        """
+        Load zonal stats results from a completed mass stats job.
+        
+        Args:
+            job_id: The job ID returned from the mass stats execution
+            max_files: Maximum number of files to download (default: 5)
+            poll_interval: Seconds to wait between status checks (default: 30)
+            
+        Returns:
+            GeoDataFrame with geometry and dataset columns, or None if failed
+        """
+        try:
+            while True:
+                try:
+                    track_info = await self.track_job([job_id])
+                    job_info = track_info[job_id]
+                    status = job_info['status']
+                    
+                    self._client.logger.info(f"Job {job_id} status: {status}")
+                    
+                    if status == 'Completed':
+                        self._client.logger.info('Job completed successfully!')
+                        break
+                    elif status in ['Failed', 'Cancelled', 'Error']:
+                        raise RuntimeError(f"Job {job_id} failed with status: {status}")
+                    
+                    await asyncio.sleep(poll_interval)
+                    
+                except KeyboardInterrupt:
+                    self._client.logger.info(f"\nInterrupted! Job {job_id} is still running.")
+                    raise
+            
+            async with aiohttp.ClientSession() as session:
+                payload = {
+                    "job_name": job_info['name'],
+                    "file_type": "raw",
+                    "bucket": job_info['bucket']
+                }
+                
+                result = await self._client._terrakio_request("POST", "mass_stats/download_files", json=payload)
+                download_urls = result['download_urls'][:max_files]
+                
+                self._client.logger.info(f"Downloading {len(download_urls)} dataset files...")
+                
+                datasets = []
+                for i, url in enumerate(download_urls):
+                    try:
+                        self._client.logger.info(f"Downloading dataset {i+1}/{len(download_urls)}...")
+                        async with session.get(url) as response:
+                            if response.status == 200:
+                                content = await response.read()
+                                dataset = xr.open_dataset(BytesIO(content))
+                                datasets.append(dataset)
+                                self._client.logger.info(f"Successfully processed dataset {i+1}")
+                            else:
+                                self._client.logger.warning(f"Failed to download dataset {i+1}: HTTP {response.status}")
+                    except Exception as e:
+                        self._client.logger.error(f"Error downloading dataset {i+1}: {e}")
+                        continue
+                
+                if not datasets:
+                    self._client.logger.warning("No datasets were successfully downloaded")
+                    return gpd.GeoDataFrame({'geometry': [], 'dataset': []})
+                
+                try:
+                    json_response = await self._client._terrakio_request(
+                        "POST", "mass_stats/download_json", 
+                        params={"job_name": job_info['name']}
+                    )
+                    json_url = json_response["download_url"]
+                    
+                    async with session.get(json_url) as response:
+                        if response.status == 200:
+                            json_data = await response.json()
+                            self._client.logger.info("Successfully downloaded geometry data")
+                            
+                            geometries = []
+                            max_geometries = min(max_files, len(json_data), len(datasets))
+                            
+                            for i in range(max_geometries):
+                                try:
+                                    geom_dict = json_data[i]["request"]["feature"]["geometry"]
+                                    shapely_geom = shape(geom_dict)
+                                    geometries.append(shapely_geom)
+                                except (KeyError, ValueError) as e:
+                                    self._client.logger.warning(f"Error parsing geometry {i}: {e}")
+                                    continue
+                            
+                            min_length = min(len(datasets), len(geometries))
+                            if min_length == 0:
+                                self._client.logger.warning("No matching datasets and geometries found")
+                                return gpd.GeoDataFrame({'geometry': [], 'dataset': []})
+                            
+                            gdf = gpd.GeoDataFrame({
+                                'geometry': geometries[:min_length],
+                                'dataset': datasets[:min_length]
+                            })
+                            
+                            self._client.logger.info(f"Created GeoDataFrame with {len(gdf)} rows")
+                            
+                            try:
+                                expanded_gdf = expand_on_variables_and_time(gdf)
+                                return expanded_gdf
+                            except NameError:
+                                self._client.logger.warning("expand_on_variables_and_time function not found, returning raw GeoDataFrame")
+                                return gdf
+                            
+                        else:
+                            self._client.logger.warning(f"Failed to download geometry data: HTTP {response.status}")
+                            return gpd.GeoDataFrame({'geometry': [], 'dataset': []})
+                            
+                except Exception as e:
+                    self._client.logger.error(f"Error downloading geometry data: {e}")
+                    return gpd.GeoDataFrame({'geometry': [], 'dataset': []})
+        
+        except Exception as e:
+            self._client.logger.error(f"Failed to load zonal stats for job {job_id}: {e}")
+            return None

terrakio_core/helper/tiles.py

@@ -16,20 +16,21 @@ def get_bounds(aoi, crs, to_crs = None):
     bounds = aoi.geometry[0].bounds
     return *bounds, aoi
 
-def tile_generator(x_min, y_min, x_max, y_max, aoi, crs, res, tile_size, expression, output, fully_cover=True):
-    i_max = int((x_max-x_min)/(tile_size*res))
-    j_max = int((y_max-y_min)/(tile_size*res))
-    if fully_cover:
-        i_max += 1
-        j_max += 1
+def tile_generator(x_min, y_min, x_max, y_max, aoi, crs, res, tile_size, expression, output, mask):
+    i_max = int((x_max-x_min)/(tile_size*res)) + 1
+    j_max = int((y_max-y_min)/(tile_size*res)) + 1
     for j in range(0, int(j_max)):
         for i in range(0, int(i_max)):
             x = x_min + i*(tile_size*res)
             y = y_max - j*(tile_size*res)
-            bbox = shapely.geometry.box(x, y-(tile_size*res), x + (tile_size*res), y)
-            if not aoi.geometry[0].intersects(bbox):
+            geom = shapely.geometry.box(x, y-(tile_size*res), x + (tile_size*res), y)
+            if not aoi.geometry[0].intersects(geom):
                 continue
-            feat  = {"type": "Feature", "geometry": bbox.__geo_interface__}
+            if mask:
+                geom = geom.intersection(aoi.geometry[0])
+                if geom.is_empty:
+                    continue
+            feat  = {"type": "Feature", "geometry": geom.__geo_interface__}
             data = {
                 "feature": feat,
                 "in_crs": crs,
@@ -46,15 +47,15 @@ def tiles(
     aoi : str, 
     expression: str = "red=S2v2#(year,median).red@(year =2024) \n red",
     output: str = "netcdf",
-    tile_size : float = 512,
+    tile_size : float = 1024,
     crs : str = "epsg:3577",
     res: float = 10,
     region : str = "eu",
     to_crs: str = None,
-    fully_cover: bool = True,
     overwrite: bool = False,
     skip_existing: bool = False,
     non_interactive: bool = False,
+    mask: bool = True,
 ):
     
     reqs = []
@@ -62,7 +63,7 @@ def tiles(
 
     if to_crs is None:
         to_crs = crs
-    for tile_req, i, j in tile_generator(x_min, y_min, x_max, y_max, aoi, to_crs, res, tile_size, expression, output, fully_cover):
+    for tile_req, i, j in tile_generator(x_min, y_min, x_max, y_max, aoi, to_crs, res, tile_size, expression, output, mask):
         req_name = f"{name}_{i:02d}_{j:02d}"
         reqs.append({"group": "tiles", "file": req_name, "request": tile_req})
 

terrakio_core/sync_client.py

@@ -339,38 +339,41 @@
 
 
 import asyncio
-import functools
 import concurrent.futures
-from typing import Optional, Dict, Any, Union
+import threading
+import functools
+import inspect
+from typing import Optional, Dict, Any, Union, TYPE_CHECKING
 from geopandas import GeoDataFrame
 from shapely.geometry.base import BaseGeometry as ShapelyGeometry
 from .async_client import AsyncClient
 
+# Add type checking imports for better IDE support
+if TYPE_CHECKING:
+    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
+
 
 class SyncWrapper:
-    """
-    Generic synchronous wrapper with __dir__ support for runtime autocomplete.
-    """
+    """Generic synchronous wrapper with __dir__ support for runtime autocomplete."""
     
     def __init__(self, async_obj, sync_client):
         self._async_obj = async_obj
         self._sync_client = sync_client
     
     def __dir__(self):
-        """
-        Return list of attributes for autocomplete in interactive environments.
-        This enables autocomplete in Jupyter/iPython after instantiation.
-        """
+        """Return list of attributes for autocomplete in interactive environments."""
         async_attrs = [attr for attr in dir(self._async_obj) if not attr.startswith('_')]
-        
         wrapper_attrs = [attr for attr in object.__dir__(self) if not attr.startswith('_')]
-        
         return list(set(async_attrs + wrapper_attrs))
     
     def __getattr__(self, name):
-        """
-        Dynamically wrap any method call to convert async to sync.
-        """
+        """Dynamically wrap any method call to convert async to sync."""
         attr = getattr(self._async_obj, name)
         
         if callable(attr):
@@ -387,15 +390,32 @@ class SyncWrapper:
 
 class SyncClient:
     """
-    Synchronous wrapper with __dir__ support for runtime autocomplete.
-    Works best in interactive environments like Jupyter/iPython.
+    Thread-safe synchronous wrapper for AsyncClient.
+    Uses a persistent event loop in a dedicated thread to avoid event loop conflicts.
     """
+
+    # Add explicit type annotations for endpoint managers
+    datasets: 'DatasetManagement'
+    users: 'UserManagement' 
+    mass_stats: 'MassStats'
+    groups: 'GroupManagement'
+    space: 'SpaceManagement'
+    model: 'ModelManagement'
+    auth: 'AuthClient'
     
     def __init__(self, url: Optional[str] = None, api_key: Optional[str] = None, verbose: bool = False):
         self._async_client = AsyncClient(url=url, api_key=api_key, verbose=verbose)
         self._context_entered = False
         self._closed = False
         
+        # Thread and event loop management
+        self._loop = None
+        self._thread = None
+        self._loop_ready = None
+        self._loop_exception = None
+        
+        # Initialize endpoint managers with proper typing
+
         self.datasets = SyncWrapper(self._async_client.datasets, self)
         self.users = SyncWrapper(self._async_client.users, self)
         self.mass_stats = SyncWrapper(self._async_client.mass_stats, self)
@@ -407,21 +427,150 @@ class SyncClient:
         import atexit
         atexit.register(self._cleanup)
     
-    def __dir__(self):
+    def _ensure_event_loop(self) -> None:
+        """Ensure we have a persistent event loop in a dedicated thread."""
+        if self._loop is None or self._loop.is_closed():
+            self._loop_ready = threading.Event()
+            self._loop_exception = None
+            
+            def run_loop():
+                """Run the event loop in a dedicated thread."""
+                try:
+                    # Create a new event loop for this thread
+                    self._loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(self._loop)
+                    
+                    # Signal that the loop is ready
+                    self._loop_ready.set()
+                    
+                    # Run the loop forever (until stopped)
+                    self._loop.run_forever()
+                except Exception as e:
+                    self._loop_exception = e
+                    self._loop_ready.set()
+                finally:
+                    # Clean up when the loop stops
+                    if self._loop and not self._loop.is_closed():
+                        self._loop.close()
+            
+            # Start the thread
+            self._thread = threading.Thread(target=run_loop, daemon=True)
+            self._thread.start()
+            
+            # Wait for the loop to be ready
+            self._loop_ready.wait(timeout=10)
+            
+            if self._loop_exception:
+                raise self._loop_exception
+            
+            if not self._loop_ready.is_set():
+                raise RuntimeError("Event loop failed to start within timeout")
+    
+    def _run_async(self, coro):
         """
-        Return list of attributes for autocomplete in interactive environments.
-        This includes all methods from the async client plus the endpoint managers.
+        Run async coroutine using persistent event loop.
+        This is the core method that makes everything work.
         """
-        default_attrs = [attr for attr in object.__dir__(self) if not attr.startswith('_')]
+        # Ensure we have an event loop
+        self._ensure_event_loop()
         
-        async_client_attrs = [attr for attr in dir(self._async_client) if not attr.startswith('_')]
+        if self._loop.is_closed():
+            raise RuntimeError("Event loop is closed")
         
-        endpoint_attrs = ['datasets', 'users', 'mass_stats', 'groups', 'space', 'model', 'auth']
+        # Create a future to get the result back from the event loop thread
+        future = concurrent.futures.Future()
         
-        all_attrs = default_attrs + async_client_attrs + endpoint_attrs
+        async def run_with_context():
+            """Run the coroutine with proper context management."""
+            try:
+                # Ensure the async client is properly initialized
+                await self._ensure_context()
+                
+                # Run the actual coroutine
+                result = await coro
+                
+                # Set the result on the future
+                future.set_result(result)
+            except Exception as e:
+                # Set the exception on the future
+                future.set_exception(e)
+        
+        # Schedule the coroutine on the persistent event loop
+        self._loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(run_with_context())
+        )
+
         
+        # Wait for the result (with timeout to avoid hanging)
+        try:
+            return future.result(timeout=300)  # 5 minute timeout
+        except concurrent.futures.TimeoutError:
+            raise RuntimeError("Async operation timed out after 5 minutes")
+    
+    async def _ensure_context(self) -> None:
+        """Ensure the async client context is entered."""
+        if not self._context_entered and not self._closed:
+            await self._async_client.__aenter__()
+            self._context_entered = True
+    
+    async def _exit_context(self) -> None:
+        """Exit the async client context."""
+        if self._context_entered and not self._closed:
+            await self._async_client.__aexit__(None, None, None)
+            self._context_entered = False
+    
+    def close(self) -> None:
+        """Close the underlying async client session and stop the event loop."""
+        if not self._closed:
+            if self._loop and not self._loop.is_closed():
+                # Schedule cleanup on the event loop
+                future = concurrent.futures.Future()
+                
+                async def cleanup():
+                    """Clean up the async client."""
+                    try:
+                        await self._exit_context()
+                        future.set_result(None)
+                    except Exception as e:
+                        future.set_exception(e)
+                
+                # Run cleanup
+                self._loop.call_soon_threadsafe(
+                    lambda: asyncio.create_task(cleanup())
+                )
+                
+                # Wait for cleanup to complete
+                try:
+                    future.result(timeout=10)
+                except:
+                    pass  # Ignore cleanup errors
+                
+                # Stop the event loop
+                self._loop.call_soon_threadsafe(self._loop.stop)
+                
+                # Wait for thread to finish
+                if self._thread and self._thread.is_alive():
+                    self._thread.join(timeout=5)
+            
+            self._closed = True
+    
+    def _cleanup(self) -> None:
+        """Internal cleanup method called by atexit."""
+        if not self._closed:
+            try:
+                self.close()
+            except Exception:
+                pass  # Ignore cleanup errors
+    
+    def __dir__(self):
+        """Return list of attributes for autocomplete in interactive environments."""
+        default_attrs = [attr for attr in object.__dir__(self) if not attr.startswith('_')]
+        async_client_attrs = [attr for attr in dir(self._async_client) if not attr.startswith('_')]
+        endpoint_attrs = ['datasets', 'users', 'mass_stats', 'groups', 'space', 'model', 'auth']
+        all_attrs = default_attrs + async_client_attrs + endpoint_attrs
         return list(set(all_attrs))
     
+    # Your existing methods with proper type annotations
     def geoquery(
         self,
         expr: str,
@@ -446,15 +595,18 @@ class SyncClient:
         return self._run_async(coro)
 
     def zonal_stats(
-            self,
-            gdf: GeoDataFrame,
-            expr: str,
-            conc: int = 20,
-            in_crs: str = "epsg:4326",
-            out_crs: str = "epsg:4326",
-            resolution: int = -1,
-            geom_fix: bool = False,
-    ):
+        self,
+        gdf: GeoDataFrame,
+        expr: str,
+        conc: int = 20,
+        in_crs: str = "epsg:4326",
+        out_crs: str = "epsg:4326",
+        resolution: int = -1,
+        geom_fix: bool = False,
+        mass_stats: bool = False,
+        id_column: Optional[str] = None,
+    ) -> GeoDataFrame:
+
         """
         Compute zonal statistics for all geometries in a GeoDataFrame (synchronous version).
         
@@ -466,6 +618,8 @@ class SyncClient:
             out_crs (str): Output coordinate reference system
             resolution (int): Resolution parameter
             geom_fix (bool): Whether to fix the geometry (default False)
+            mass_stats (bool): Whether to use mass stats for processing (default False)
+            id_column (Optional[str]): Name of the ID column to use (default None)
 
         Returns:
             geopandas.GeoDataFrame: GeoDataFrame with added columns for results
@@ -482,11 +636,14 @@ class SyncClient:
             out_crs=out_crs,
             resolution=resolution,
             geom_fix=geom_fix,
+            mass_stats=mass_stats,
+            id_column=id_column,
         )
         return self._run_async(coro)
     
     def create_dataset_file(
         self,
+        name: str,
         aoi: str,
         expression: str,
         output: str,
@@ -499,6 +656,9 @@ class SyncClient:
         non_interactive: bool = True,
         poll_interval: int = 30,
         download_path: str = "/home/user/Downloads",
+        mask = True,
+        max_file_size_mb: int = 5120,  # Default to 5GB
+        tile_size: int = 1024,
     ) -> dict:
         """Create a dataset file using mass stats operations (synchronous version)."""
         coro = self._async_client.create_dataset_file(
@@ -514,14 +674,19 @@ class SyncClient:
             non_interactive=non_interactive,
             poll_interval=poll_interval,
             download_path=download_path,
+            name=name,
+            mask=mask,
+            max_file_size_mb=max_file_size_mb,
+            tile_size=tile_size
         )
         return self._run_async(coro)
 
+
     def geo_queries(
         self,
         queries: list[dict],
         conc: int = 20,
-    ):
+    ) -> Union[float, GeoDataFrame]:
         """
         Execute multiple geo queries concurrently (synchronous version).
 
@@ -559,120 +724,17 @@ class SyncClient:
             conc=conc,
         )
         return self._run_async(coro)
-
-    async def _ensure_context(self):
-        """Ensure the async client context is entered."""
-        if not self._context_entered and not self._closed:
-            await self._async_client.__aenter__()
-            self._context_entered = True
-    
-    async def _exit_context(self):
-        """Exit the async client context."""
-        if self._context_entered and not self._closed:
-            await self._async_client.__aexit__(None, None, None)
-            self._context_entered = False
-    
-    def _run_async(self, coro):
-        """
-        Run an async coroutine and return the result synchronously.
-        This version handles both Jupyter notebook environments and regular Python environments.
-        """
-        async def run_with_context():
-            await self._ensure_context()
-            return await coro
-        
-        try:
-            # Check if we're in a running event loop (like Jupyter)
-            loop = asyncio.get_running_loop()
-            
-            # Method 1: Try using nest_asyncio if available
-            try:
-                import nest_asyncio
-                nest_asyncio.apply()
-                return asyncio.run(run_with_context())
-            except ImportError:
-                pass
-            
-            # Method 2: Use ThreadPoolExecutor to run in a separate thread
-            def run_in_thread():
-                return asyncio.run(run_with_context())
-            
-            with concurrent.futures.ThreadPoolExecutor() as executor:
-                future = executor.submit(run_in_thread)
-                return future.result()
-                
-        except RuntimeError:
-            # No running loop, safe to use asyncio.run()
-            return asyncio.run(run_with_context())
     
-    def close(self):
-        """Close the underlying async client session."""
-        if not self._closed:
-            async def close_async():
-                await self._exit_context()
-            
-            try:
-                loop = asyncio.get_running_loop()
-                
-                # Try nest_asyncio first
-                try:
-                    import nest_asyncio
-                    nest_asyncio.apply()
-                    asyncio.run(close_async())
-                except ImportError:
-                    # Fall back to ThreadPoolExecutor
-                    def run_in_thread():
-                        return asyncio.run(close_async())
-                    
-                    with concurrent.futures.ThreadPoolExecutor() as executor:
-                        future = executor.submit(run_in_thread)
-                        future.result()
-                        
-            except RuntimeError:
-                asyncio.run(close_async())
-            
-            self._closed = True
-    
-    def _cleanup(self):
-        """Internal cleanup method called by atexit."""
-        if not self._closed:
-            try:
-                self.close()
-            except Exception:
-                pass
-    
-    def __enter__(self):
+    # Context manager support
+    def __enter__(self) -> 'SyncClient':
         """Context manager entry."""
-        async def enter_async():
-            await self._ensure_context()
-        
-        try:
-            loop = asyncio.get_running_loop()
-            
-            # Try nest_asyncio first
-            try:
-                import nest_asyncio
-                nest_asyncio.apply()
-                asyncio.run(enter_async())
-            except ImportError:
-                # Fall back to ThreadPoolExecutor
-                def run_in_thread():
-                    return asyncio.run(enter_async())
-                
-                with concurrent.futures.ThreadPoolExecutor() as executor:
-                    future = executor.submit(run_in_thread)
-                    future.result()
-                    
-        except RuntimeError:
-            asyncio.run(enter_async())
-        
         return self
     
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         """Context manager exit."""
         self.close()
     
-    def __del__(self):
+    def __del__(self) -> None:
         """Destructor to ensure session is closed."""
         if not self._closed:
             try:

{terrakio_core-0.4.5.dist-info → terrakio_core-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.4.5
+Version: 0.4.7
 Summary: Core components for Terrakio API clients
 Author-email: Yupeng Chao <yupeng@haizea.com.au>
 Project-URL: Homepage, https://github.com/HaizeaAnalytics/terrakio-python-api
@@ -25,6 +25,10 @@ Requires-Dist: google-cloud-storage>=2.0.0
 Requires-Dist: scipy>=1.7.0
 Requires-Dist: nest_asyncio
 Requires-Dist: onnxruntime>=1.10.0
+Requires-Dist: psutil>=5.0.0
+Requires-Dist: h5netcdf>=1.0.0
+Requires-Dist: netcdf4>=1.5.0
+Requires-Dist: aiofiles>=24.1.0
 Provides-Extra: ml
 Requires-Dist: torch>=2.7.1; extra == "ml"
 Requires-Dist: scikit-learn>=1.7.0; extra == "ml"

{terrakio_core-0.4.5.dist-info → terrakio_core-0.4.7.dist-info}/RECORD

@@ -1,22 +1,22 @@
-terrakio_core/__init__.py,sha256=snOTpM6-GxHvvDJ6hjlsWywgIago1JbE8QDfsWMRlPI,273
+terrakio_core/__init__.py,sha256=oBlKpY5_oKnNAR3VXehRw4a0mmg7W0SvMGh2QEwqO0A,273
 terrakio_core/accessors.py,sha256=qWLljU83YO7EUOefo_f6_P6ba6uiYMXwou0ihAHKBHQ,23706
-terrakio_core/async_client.py,sha256=M7C6by-GCW9AvDRhIwpBG12fO5--NG5qrHdSC6Fjkrc,14082
+terrakio_core/async_client.py,sha256=FNl1K3g6UoB1aKPAG6gFc-25ukNLCmyI-m-UpLawzFI,14656
 terrakio_core/client.py,sha256=VXP7BtJWIfpPPZR7_yNdSTcGwNgTwhb7KorusqkQrzk,5603
 terrakio_core/config.py,sha256=r8NARVYOca4AuM88VP_j-8wQxOk1s7VcRdyEdseBlLE,4193
 terrakio_core/exceptions.py,sha256=4qnpOM1gOxsNIXDXY4qwY1d3I4Myhp7HBh7b2D0SVrU,529
-terrakio_core/sync_client.py,sha256=wsVKjaIDZW4AC2-81RXseE9mFbXu90T5Bbwrgw37x6k,24864
-terrakio_core/convenience_functions/convenience_functions.py,sha256=B7qbObjP4OuAUtrVf4Gi58c0q2EkTwOCOE568P0Q-EE,18607
+terrakio_core/sync_client.py,sha256=ATM-2aeLj4-V2x92Fc_Fxsr4UiI9iT8G0PqrDXh_8XU,27759
+terrakio_core/convenience_functions/convenience_functions.py,sha256=tbzG7j6o0HojAEJqAxg2KPqCpSASJlLnayFQ_Cxfips,21700
 terrakio_core/endpoints/auth.py,sha256=FdLsPScPIBo-Gxl6ZnE-46cp2molggAJtL72LssN3fg,6049
 terrakio_core/endpoints/dataset_management.py,sha256=BUm8IIlW_Q45vDiQp16CiJGeSLheI8uWRVRQtMdhaNk,13161
 terrakio_core/endpoints/group_management.py,sha256=VFl3jakjQa9OPi351D3DZvLU9M7fHdfjCzGhmyJsx3U,6309
-terrakio_core/endpoints/mass_stats.py,sha256=6PSWWCpKLKMcFOoXfOAuROX8iSuG_dLyTYcw7gzhhZ4,23464
+terrakio_core/endpoints/mass_stats.py,sha256=Gz70eyRrmQykrr7zFC5Pd5RTNvVpHgyTTF3rr2lNuFE,30038
 terrakio_core/endpoints/model_management.py,sha256=LH_gHPrqYA-_45KWpDBRcFbwHgm-Kg0zk1ealy7P_C0,52379
 terrakio_core/endpoints/space_management.py,sha256=YWb55nkJnFJGlALJ520DvurxDqVqwYtsvqQPWzxzhDs,2266
 terrakio_core/endpoints/user_management.py,sha256=WlFr3EfK8iI6DfkpMuYLHZUPk2n7_DHHO6z1hndmZB4,3816
 terrakio_core/helper/bounded_taskgroup.py,sha256=wiTH10jhKZgrsgrFUNG6gig8bFkUEPHkGRT2XY7Rgmo,677
 terrakio_core/helper/decorators.py,sha256=L6om7wmWNgCei3Wy5U0aZ-70OzsCwclkjIf7SfQuhCg,2289
-terrakio_core/helper/tiles.py,sha256=xNtp3oDD912PN_FQV5fb6uQYhwfHANuXyIcxoVCCfZU,2632
-terrakio_core-0.4.5.dist-info/METADATA,sha256=5_6bPROpM38AJ7NJDqFzORJFG2HJR-j_X0Rg4KwriUY,1791
-terrakio_core-0.4.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-terrakio_core-0.4.5.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
-terrakio_core-0.4.5.dist-info/RECORD,,
+terrakio_core/helper/tiles.py,sha256=lcLCO6KiP05lCI9vngo3zCZJ6Z9C0pUxHSQS4H58EHc,2699
+terrakio_core-0.4.7.dist-info/METADATA,sha256=t59FmeOR6RkUZ9XufM4qHDwaGREQjl2__VOAXFDAV_Y,1913
+terrakio_core-0.4.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+terrakio_core-0.4.7.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
+terrakio_core-0.4.7.dist-info/RECORD,,