datamint 2.3.3__py3-none-any.whl → 2.9.0__py3-none-any.whl

datamint/api/endpoints/resources_api.py

@@ -1,10 +1,11 @@
-from typing import Any, Optional, Sequence, TypeAlias, Literal, IO
+from typing import TypeAlias, Literal, IO
+from collections.abc import Sequence
 from ..base_api import ApiConfig, BaseApi
 from ..entity_base_api import CreatableEntityApi, DeletableEntityApi
 from datamint.entities.resource import Resource
-from datamint.entities.project import Project
-from datamint.entities.annotation import Annotation
+from datamint.entities.annotations.annotation import Annotation
 from datamint.exceptions import DatamintException, ResourceNotFoundError
+from datamint.api.dto import AnnotationType
 import httpx
 from datetime import date
 import json
@@ -21,10 +22,10 @@ from tqdm.auto import tqdm
 import asyncio
 import aiohttp
 from pathlib import Path
-import nest_asyncio  # For running asyncio in jupyter notebooks
 from PIL import Image
 import io
 from datamint.types import ImagingData
+from collections import defaultdict
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -52,9 +53,11 @@ def _open_io(file_path: str | Path | IO, mode: str = 'rb') -> IO:
 class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
     """API handler for resource-related endpoints."""
 
+    _ENDPOINT_BASE = 'resources'
+
     def __init__(self,
                  config: ApiConfig,
-                 client: Optional[httpx.Client] = None,
+                 client: httpx.Client | None = None,
                  annotations_api=None,
                  projects_api=None
                  ) -> None:
@@ -66,25 +69,24 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
         """
         from .annotations_api import AnnotationsApi
         from .projects_api import ProjectsApi
-        super().__init__(config, Resource, 'resources', client)
-        nest_asyncio.apply()
+        super().__init__(config, Resource, ResourcesApi._ENDPOINT_BASE, client)
         self.annotations_api = AnnotationsApi(
             config, client, resources_api=self) if annotations_api is None else annotations_api
-        self.projects_api = ProjectsApi(config, client) if projects_api is None else projects_api
+        self.projects_api = projects_api or ProjectsApi(config, client, resources_api=self)
 
     def get_list(self,
-                 status: Optional[ResourceStatus] = None,
+                 status: ResourceStatus | None = None,
                  from_date: date | str | None = None,
                  to_date: date | str | None = None,
-                 tags: Optional[Sequence[str]] = None,
-                 modality: Optional[str] = None,
-                 mimetype: Optional[str] = None,
+                 tags: Sequence[str] | None = None,
+                 modality: str | None = None,
+                 mimetype: str | None = None,
                  #  return_ids_only: bool = False,
-                 order_field: Optional[ResourceFields] = None,
-                 order_ascending: Optional[bool] = None,
-                 channel: Optional[str] = None,
+                 order_field: ResourceFields | None = None,
+                 order_ascending: bool | None = None,
+                 channel: str | None = None,
                  project_name: str | list[str] | None = None,
-                 filename: Optional[str] = None,
+                 filename: str | None = None,
                  limit: int | None = None
                  ) -> Sequence[Resource]:
         """Get resources with optional filtering.
@@ -146,7 +148,8 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
 
         return super().get_list(limit=limit, params=payload)
 
-    def get_annotations(self, resource: str | Resource) -> Sequence[Annotation]:
+    def get_annotations(self, resource: str | Resource,
+                        annotation_type: AnnotationType | str | None = None) -> Sequence[Annotation]:
         """Get annotations for a specific resource.
 
         Args:
@@ -155,7 +158,9 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
         Returns:
             A sequence of Annotation objects associated with the specified resource.
         """
-        return self.annotations_api.get_list(resource=resource)
+        return self.annotations_api.get_list(resource=resource,
+                                             load_ai_segmentations=True,
+                                             annotation_type=annotation_type)
 
     @staticmethod
     def __process_files_parameter(file_path: str | Sequence[str | IO | pydicom.Dataset]
@@ -229,16 +234,16 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
 
     async def _upload_single_resource_async(self,
                                             file_path: str | IO,
-                                            mimetype: Optional[str] = None,
+                                            mimetype: str | None = None,
                                             anonymize: bool = False,
                                             anonymize_retain_codes: Sequence[tuple] = [],
                                             tags: list[str] = [],
                                             mung_filename: Sequence[int] | Literal['all'] | None = None,
-                                            channel: Optional[str] = None,
+                                            channel: str | None = None,
                                             session=None,
-                                            modality: Optional[str] = None,
+                                            modality: str | None = None,
                                             publish: bool = False,
-                                            metadata_file: Optional[str | dict] = None,
+                                            metadata_file: str | dict | None = None,
                                             ) -> str:
         if is_io_object(file_path):
             source_filepath = os.path.abspath(os.path.expanduser(file_path.name))
@@ -263,8 +268,8 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
             filename = new_filename
 
         is_a_dicom_file = None
-        if mimetype is None:
-            mimetype_list, ext = guess_typez(file_path, use_magic=True)
+        if not mimetype or mimetype == DEFAULT_MIME_TYPE:
+            mimetype_list, ext = guess_typez(file_path, use_magic=False)
             for mime in mimetype_list:
                 if mime in NIFTI_MIMES:
                     mimetype = DEFAULT_NIFTI_MIME
@@ -273,7 +278,11 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                 if ext == '.nii.gz' or filename.lower().endswith('nii.gz'):
                     mimetype = DEFAULT_NIFTI_MIME
                 else:
-                    mimetype = mimetype_list[-1] if mimetype_list else DEFAULT_MIME_TYPE
+                    mimetype = mimetype_list[-1] if mimetype_list and mimetype_list[-1] else DEFAULT_MIME_TYPE
+        if (not mimetype) or mimetype == DEFAULT_MIME_TYPE:
+            msg = f"Could not determine mimetype for file {source_filepath}."
+            _LOGGER.warning(msg)
+            _USER_LOGGER.warning(msg)
 
         mimetype = standardize_mimetype(mimetype)
 
@@ -349,9 +358,12 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                 except Exception as e:
                     _LOGGER.warning(f"Failed to add metadata to form: {e}")
 
+            timeout = aiohttp.ClientTimeout(total=300, connect=60, sock_read=300)
             resp_data = await self._make_request_async_json('POST',
                                                             endpoint=self.endpoint_base,
-                                                            data=form)
+                                                            data=form,
+                                                            session=session,
+                                                            timeout=timeout)
             if 'error' in resp_data:
                 raise DatamintException(resp_data['error'])
             _LOGGER.debug(f"Response on uploading {filename}: {resp_data}")
@@ -367,20 +379,33 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
 
     async def _upload_resources_async(self,
                                       files_path: Sequence[str | IO],
-                                      mimetype: Optional[str] = None,
+                                      mimetype: str | None = None,
                                       anonymize: bool = False,
                                       anonymize_retain_codes: Sequence[tuple] = [],
                                       on_error: Literal['raise', 'skip'] = 'raise',
                                       tags=None,
                                       mung_filename: Sequence[int] | Literal['all'] | None = None,
-                                      channel: Optional[str] = None,
-                                      modality: Optional[str] = None,
+                                      channel: str | None = None,
+                                      modality: str | None = None,
                                       publish: bool = False,
                                       segmentation_files: Sequence[dict] | None = None,
                                       transpose_segmentation: bool = False,
                                       metadata_files: Sequence[str | dict | None] | None = None,
                                       progress_bar: tqdm | None = None,
+                                      session: aiohttp.ClientSession | None = None,
                                       ) -> list[str]:
+        """Upload multiple resources asynchronously.
+
+        Args:
+            session: Optional aiohttp session. If None, uses a shared session managed by
+                this API instance. Callers should NOT close a session they pass in until
+                all async operations complete, as nested calls may reuse it.
+        
+        Note:
+            Session ownership: When session=None (default), a long-lived shared session
+            is automatically used and managed. When explicitly passed, the caller retains
+            ownership and must ensure it remains open for the duration of this call.
+        """
         if on_error not in ['raise', 'skip']:
             raise ValueError("on_error must be either 'raise' or 'skip'")
 
@@ -390,52 +415,61 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
         if metadata_files is None:
             metadata_files = _infinite_gen(None)
 
-        async with aiohttp.ClientSession() as session:
-            async def __upload_single_resource(file_path, segfiles: dict[str, list | dict],
-                                               metadata_file: str | dict | None):
-                name = file_path.name if is_io_object(file_path) else file_path
-                name = os.path.basename(name)
-                rid = await self._upload_single_resource_async(
-                    file_path=file_path,
-                    mimetype=mimetype,
-                    anonymize=anonymize,
-                    anonymize_retain_codes=anonymize_retain_codes,
-                    tags=tags,
-                    session=session,
-                    mung_filename=mung_filename,
-                    channel=channel,
-                    modality=modality,
-                    publish=publish,
-                    metadata_file=metadata_file,
-                )
-                if progress_bar:
-                    progress_bar.update(1)
-                    progress_bar.set_postfix(file=name)
-                else:
-                    _USER_LOGGER.info(f'"{name}" uploaded')
-
-                if segfiles is not None:
-                    fpaths = segfiles['files']
-                    names = segfiles.get('names', _infinite_gen(None))
-                    if isinstance(names, dict):
-                        names = _infinite_gen(names)
-                    frame_indices = segfiles.get('frame_index', _infinite_gen(None))
-                    for f, name, frame_index in tqdm(zip(fpaths, names, frame_indices),
-                                                     desc=f"Uploading segmentations for {file_path}",
-                                                     total=len(fpaths)):
-                        if f is not None:
-                            await self.annotations_api._upload_segmentations_async(
-                                rid,
-                                file_path=f,
-                                name=name,
-                                frame_index=frame_index,
-                                transpose_segmentation=transpose_segmentation
-                            )
-                return rid
+        # Use shared session by default for stability across multiple sequential calls.
+        # This prevents connection churn and intermittent SSL shutdown timeouts.
+        session = session or self._get_aiohttp_session()
+
+        async def __upload_single_resource(file_path, segfiles: dict[str, list | dict],
+                                           metadata_file: str | dict | None):
+            name = file_path.name if is_io_object(file_path) else file_path
+            name = os.path.basename(name)
+            rid = await self._upload_single_resource_async(
+                file_path=file_path,
+                mimetype=mimetype,
+                anonymize=anonymize,
+                anonymize_retain_codes=anonymize_retain_codes,
+                tags=tags,
+                session=session,
+                mung_filename=mung_filename,
+                channel=channel,
+                modality=modality,
+                publish=publish,
+                metadata_file=metadata_file,
+            )
+            if progress_bar:
+                progress_bar.update(1)
+                progress_bar.set_postfix(file=name)
+            else:
+                _USER_LOGGER.info(f'"{name}" uploaded')
+
+            if segfiles is not None:
+                fpaths = segfiles['files']
+                names = segfiles.get('names', _infinite_gen(None))
+                if isinstance(names, dict):
+                    names = _infinite_gen(names)
+                frame_indices = segfiles.get('frame_index', _infinite_gen(None))
+                for f, name, frame_index in tqdm(zip(fpaths, names, frame_indices),
+                                                 desc=f"Uploading segmentations for {file_path}",
+                                                 total=len(fpaths)):
+                    if f is not None:
+                        await self.annotations_api._upload_segmentations_async(
+                            rid,
+                            file_path=f,
+                            name=name,
+                            frame_index=frame_index,
+                            transpose_segmentation=transpose_segmentation
+                        )
+            return rid
 
+        try:
             tasks = [__upload_single_resource(f, segfiles, metadata_file)
                      for f, segfiles, metadata_file in zip(files_path, segmentation_files, metadata_files)]
-            return await asyncio.gather(*tasks, return_exceptions=on_error == 'skip')
+        except ValueError:
+            msg = f"Error preparing upload tasks. Try `assemble_dicom=False`."
+            _LOGGER.error(msg)
+            _USER_LOGGER.error(msg)
+            raise
+        return await asyncio.gather(*tasks, return_exceptions=on_error == 'skip')
 
     def upload_resources(self,
                          files_path: Sequence[str | IO | pydicom.Dataset],
@@ -503,6 +537,18 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
             raise ValueError(
                 "upload_resources() only accepts multiple resources. For single resource upload, use upload_resource() instead.")
 
+        if publish_to:
+            publish = True
+            # Check if project exists
+            proj = self.projects_api.get_by_name(publish_to)
+            if proj is None:
+                try:
+                    proj = self.projects_api.get_by_id(publish_to)
+                except Exception:
+                    pass
+                if proj is None:
+                    raise ResourceNotFoundError('Project', {'name_or_id': publish_to})
+
         files_path = ResourcesApi.__process_files_parameter(files_path)
 
         # Discard DICOM reports
@@ -682,7 +728,7 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                 )
         """
         # Convert segmentation_files to the format expected by upload_resources
-        segmentation_files_list: Optional[list[list[str] | dict]] = None
+        segmentation_files_list: list[list[str] | dict] | None = None
         if segmentation_files is not None:
             segmentation_files_list = [segmentation_files]
 
@@ -737,22 +783,32 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
         """
         save_path = str(save_path)  # Ensure save_path is a string for file operations
         resource_id = self._entid(resource)
+
+        # Disable total timeout for large file downloads, keep connect timeout
+        timeout = aiohttp.ClientTimeout(total=None, sock_connect=self.config.timeout)
+
         try:
             async with self._make_request_async('GET',
                                                 f'{self.endpoint_base}/{resource_id}/file',
                                                 session=session,
-                                                headers={'accept': 'application/octet-stream'}) as resp:
-                data_bytes = await resp.read()
+                                                headers={'accept': 'application/octet-stream'},
+                                                timeout=timeout) as resp:
+
+                final_save_path = save_path
+                target_path = final_save_path
+                if add_extension:
+                    target_path = f"{save_path}.tmp"
+
+                with open(target_path, 'wb') as f:
+                    async for chunk in resp.content.iter_chunked(1024 * 1024):  # 1MB chunks
+                        f.write(chunk)
 
-            final_save_path = save_path
             if add_extension:
-                # Save to temporary file first to determine mimetype from content
-                temp_path = f"{save_path}.tmp"
-                with open(temp_path, 'wb') as f:
-                    f.write(data_bytes)
+                # Determine mimetype from file content (read first 2KB)
+                with open(target_path, 'rb') as f:
+                    head_content = f.read(2048)
 
-                # Determine mimetype from file content
-                mimetype, ext = BaseApi._determine_mimetype(content=data_bytes,
+                mimetype, ext = BaseApi._determine_mimetype(content=head_content,
                                                             declared_mimetype=resource.mimetype if isinstance(resource, Resource) else None)
 
                 # Generate final path with extension if needed
@@ -763,11 +819,7 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                         final_save_path = save_path + ext
 
                 # Move file to final location
-                os.rename(temp_path, final_save_path)
-            else:
-                # Standard save without extension detection
-                with open(final_save_path, 'wb') as f:
-                    f.write(data_bytes)
+                os.rename(target_path, final_save_path)
 
             if progress_bar:
                 progress_bar.update(1)
@@ -801,7 +853,9 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
             raise ValueError("resources must be a list of resources")
 
         async def _download_all_async():
-            async with aiohttp.ClientSession() as session:
+            connector = self._create_aiohttp_connector(force_close=True)
+            timeout = aiohttp.ClientTimeout(total=None, connect=60, sock_read=600)
+            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
                 tasks = [
                     self._async_download_file(
                         resource=r,
@@ -838,7 +892,7 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
 
     def download_resource_file(self,
                                resource: str | Resource,
-                               save_path: Optional[str] = None,
+                               save_path: str | None = None,
                                auto_convert: bool = True,
                                add_extension: bool = False
                                ) -> ImagingData | tuple[ImagingData, str] | bytes:
@@ -916,6 +970,53 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                 return resource_file, save_path
         return resource_file
 
+    def cache_resources(
+        self,
+        resources: Sequence[Resource],
+        progress_bar: bool = True,
+    ) -> None:
+        """Cache multiple resources in parallel, skipping already cached ones.
+
+        This method downloads and caches resource file data concurrently,
+        significantly improving efficiency when working with large datasets.
+        Only resources that are not already cached will be downloaded.
+
+        Args:
+            resources: Sequence of Resource instances to cache.
+            progress_bar: Whether to show a progress bar. Default is True.
+
+        Example:
+            >>> resources = api.resources.get_list(limit=100)
+            >>> api.resources.cache_resources(resources)
+            Caching resources: 100%|██████████| 85/85 [00:12<00:00,  6.8files/s]
+        """
+        # Filter out already cached resources
+        resources_to_cache = [res for res in resources if not res.is_cached()]
+
+        if not resources_to_cache:
+            _LOGGER.info("All resources are already cached.")
+            return
+
+        _LOGGER.info(f"Caching {len(resources_to_cache)} of {len(resources)} resources...")
+
+        if progress_bar:
+            pbar = tqdm(total=len(resources_to_cache), desc="Caching resources", unit="file")
+        else:
+            pbar = None
+
+        try:
+            for res in resources_to_cache:
+                res._api = self  # Ensure the resource has a reference to the API
+                res.fetch_file_data(auto_convert=False, use_cache=True)
+                if pbar:
+                    pbar.set_postfix(filename=res.filename)
+                    pbar.update(1)
+
+        finally:
+            if pbar:
+                pbar.close()
+        _LOGGER.info(f"Successfully cached {len(resources_to_cache)} resources.")
+
     def download_resource_frame(self,
                                 resource: str | Resource,
                                 frame_index: int) -> Image.Image:
@@ -940,7 +1041,7 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
             resource = self.get_by_id(resource)
         if resource.mimetype.startswith('image/') or resource.storage == 'ImageResource':
             if frame_index != 0:
-                raise DatamintException(f"Resource {resource.id} is a single frame image, "
+                raise DatamintException(f"Resource {resource.id} is not a multi-frame resource, "
                                         f"but frame_index is {frame_index}.")
             return self.download_resource_file(resource, auto_convert=True)
 
@@ -985,22 +1086,30 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
                     raise
 
     def set_tags(self,
-                 resource: str | Resource,
+                 resource: str | Resource | Sequence[str | Resource],
                  tags: Sequence[str],
                  ):
         """
         Set tags for a resource, IMPORTANT: This replaces all existing tags.
         Args:
-            resource: The resource unique id or Resource object.
+            resource: The resource object or a list of resources.
             tags: The tags to set.
         """
-        data = {'tags': tags}
-        resource_id = self._entid(resource)
 
-        response = self._make_entity_request('PUT',
-                                             resource_id,
-                                             add_path='tags',
-                                             json=data)
+        uniq_tags = set(tags)  # remove duplicates
+
+        if isinstance(resource, Sequence):
+            resource_ids = [self._entid(res) for res in resource]
+            response = self._make_request('PUT',
+                                          f'{self.endpoint_base}/tags',
+                                          json={'resource_ids': resource_ids,
+                                                'tags': list(uniq_tags)})
+        else:
+            resource_id = self._entid(resource)
+            response = self._make_entity_request('PUT',
+                                                 resource_id,
+                                                 add_path='tags',
+                                                 json={'tags': list(uniq_tags)})
         return response
 
     # def get_projects(self, resource: Resource) -> Sequence[Project]:
@@ -1018,7 +1127,7 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
     #     return [proj for proj in self.projects_api.get_all() if proj.id in proj_ids]
 
     def add_tags(self,
-                 resource: str | Resource,
+                 resource: str | Resource | Sequence[str | Resource],
                  tags: Sequence[str],
                  ):
         """
@@ -1029,8 +1138,26 @@ class ResourcesApi(CreatableEntityApi[Resource], DeletableEntityApi[Resource]):
         """
         if isinstance(resource, str):
             resource = self.get_by_id(resource)
+        elif isinstance(resource, Sequence):
+            # Transform every str to Resource first.
+            resources = [self.get_by_id(res) if isinstance(res, str) else res for res in resource]
+
+            # group resource having the exact same tags to minimize requests
+            tag_map: dict[tuple, list[Resource]] = defaultdict(list)
+            for res in resources:
+                old_tags = res.tags if res.tags is not None else []
+                # key = tuple(sorted(old_tags))
+                key = tuple(old_tags)  # keep order, assuming order matters for tags
+                tag_map[key].append(res)
+
+            # finally, set tags for each group
+            for old_tags_tuple, res_group in tag_map.items():
+                old_tags = list(old_tags_tuple)
+                self.set_tags(res_group, old_tags + list(tags))
+            return
+
         old_tags = resource.tags if resource.tags is not None else []
-        return self.set_tags(resource, old_tags + list(tags))
+        self.set_tags(resource, old_tags + list(tags))
 
     def bulk_delete(self, entities: Sequence[str | Resource]) -> None:
         """Delete multiple entities. Faster than deleting them one by one.