minor improvements to typing and ability to set mlflow project with a Project object

Author: Lucashsmello

datamint/api/endpoints/projects_api.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Literal, TYPE_CHECKING
+from typing import Sequence, Literal, TYPE_CHECKING, overload
 from ..entity_base_api import ApiConfig, CRUDEntityApi
 from datamint.entities.project import Project
 import httpx
@@ -38,13 +38,38 @@ def get_project_resources(self, project: Project | str) -> list[Resource]:
         resources = [self.resources_api._init_entity_obj(**item) for item in resources_data]
         return resources
 
+
+    @overload
+    def create(self,
+               name: str,
+               description: str,
+               resources_ids: list[str] | None = None,
+               is_active_learning: bool = False,
+               two_up_display: bool = False,
+               *,
+               return_entity: Literal[True] = True
+               ) -> Project: ...
+
+    @overload
+    def create(self,
+               name: str,
+               description: str,
+               resources_ids: list[str] | None = None,
+               is_active_learning: bool = False,
+               two_up_display: bool = False,
+               *,
+               return_entity: Literal[False]
+               ) -> str: ...
+
     def create(self,
                name: str,
                description: str,
                resources_ids: list[str] | None = None,
                is_active_learning: bool = False,
-               two_up_display: bool = False
-               ) -> str:
+               two_up_display: bool = False,
+               *,
+               return_entity: bool = True
+               ) -> str | Project:
         """Create a new project.
 
         Args:
@@ -53,6 +78,7 @@ def create(self,
             resources_ids: The list of resource ids to be included in the project.
             is_active_learning: Whether the project is an active learning project or not.
             two_up_display: Allow annotators to display multiple resources for annotation.
+            return_entity: Whether to return the created Project instance or just its ID.
 
         Returns:
             The id of the created project.
@@ -72,7 +98,7 @@ def create(self,
                         "require_review": False,
                         'description': description}
 
-        return self._create(project_data)
+        return self._create(project_data, return_entity=return_entity)
 
     def get_all(self, limit: int | None = None) -> Sequence[Project]:
         """Get all projects.

datamint/api/entity_base_api.py

@@ -1,4 +1,4 @@
-from typing import Any, TypeVar, Generic, Type, Sequence
+from typing import Any, Literal, TypeVar, Generic, Type, Sequence, AsyncGenerator, overload
 import logging
 import httpx
 from datamint.entities.base_entity import BaseEntity
@@ -7,7 +7,6 @@
 import asyncio
 from .base_api import ApiConfig, BaseApi
 import contextlib
-from typing import AsyncGenerator
 
 logger = logging.getLogger(__name__)
 T = TypeVar('T', bound=BaseEntity)
@@ -248,7 +247,16 @@ class CreatableEntityApi(EntityBaseApi[T]):
     This class adds methods to handle creation of new entities.
     """
 
-    def _create(self, entity_data: dict[str, Any]) -> str | list[str | dict]:
+    @overload
+    def _create(self, entity_data: dict[str, Any],
+                return_entity: Literal[True] = True) -> T | list[T]: ...
+
+    @overload
+    def _create(self, entity_data: dict[str, Any],
+                return_entity: Literal[False]) -> str | list: ...
+
+    def _create(self, entity_data: dict[str, Any],
+                return_entity: bool = False) -> str | T | list:
         """Create a new entity.
 
         Args:
@@ -263,14 +271,35 @@ def _create(self, entity_data: dict[str, Any]) -> str | list[str | dict]:
         response = self._make_request('POST', f'/{self.endpoint_base}', json=entity_data)
         respdata = response.json()
         if isinstance(respdata, str):
+            if return_entity:
+                return self.get_by_id(respdata)
             return respdata
         if isinstance(respdata, list):
+            if return_entity:
+                logger.warning("Current implementation is slow when returning entities on bulk create."
+                               " Try ``return_entity=False`` for better performance.")
+                return [self.get_by_id(item['id']) if isinstance(item, dict) and 'id' in item else self.get_by_id(item)
+                        for item in respdata]
             return respdata
         if isinstance(respdata, dict):
+            if return_entity:
+                try:
+                    return self._init_entity_obj(**respdata)
+                except:
+                    logger.debug("Failed to init entity obj on create response. Falling back to get_by_id.")
+                    return self.get_by_id(respdata.get('id'))
             return respdata.get('id')
         return respdata
 
-    def create(self, *args, **kwargs) -> str | T:
+    @overload
+    def create(self, *args, return_entity: Literal[True] = True, **kwargs) -> T: ...
+
+    @overload
+    def create(self, *args, return_entity: Literal[False], **kwargs) -> str: ...
+
+    def create(self, *args,
+               return_entity: bool = True,
+               **kwargs) -> str | T:
         raise NotImplementedError("Subclasses must implement the create method with their own custom parameters")
 
 

datamint/mlflow/tracking/fluent.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 import threading
 import logging
 from datamint import Api
@@ -7,6 +7,9 @@
 from datamint.mlflow.env_vars import EnvVars
 from datamint.mlflow.env_utils import ensure_mlflow_configured
 
+if TYPE_CHECKING:
+    from datamint.entities.project import Project
+
 _PROJECT_LOCK = threading.Lock()
 _LOGGER = logging.getLogger(__name__)
 
@@ -44,30 +47,35 @@ def _find_project_by_name(project_name: str):
     return project
 
 
-def set_project(project_name: Optional[str] = None, project_id: Optional[str] = None):
-    from mlflow.exceptions import MlflowException
+def _get_project_by_name_or_id(project_name_or_id: str) -> 'Project':
+    dt_client = Api(check_connection=False)
+    # If length >= 32, likely an ID
+    if len(project_name_or_id) >= 32 and ' ' not in project_name_or_id:
+        # Try to get by ID first
+        project = dt_client.projects.get_by_id(project_name_or_id)
+        if project is not None:
+            return project
+    project = dt_client.projects.get_by_name(project_name_or_id)
+    if project is None:
+        raise DatamintException(f"Project '{project_name_or_id}' does not exist.")
+    return project
+
+
+def set_project(project: 'Project | str'):
+    from datamint.entities.project import Project as ProjectEntity
     global _ACTIVE_PROJECT_ID
 
     # Ensure MLflow is properly configured before proceeding
     ensure_mlflow_configured()
 
-    if project_name is None and project_id is None:
-        raise MlflowException("You must specify either a project name or a project id")
-
-    if project_name is not None and project_id is not None:
-        raise MlflowException("You cannot specify both a project name and a project id")
-
     with _PROJECT_LOCK:
-        dt_client = Api(check_connection=False)
-        if project_id is None:
-            project = dt_client.projects.get_by_name(project_name)
-            if project is None:
-                raise DatamintException(f"Project with name '{project_name}' does not exist.")
+        if isinstance(project, str):
+            project_id = None
+            project = _get_project_by_name_or_id(project)
             project_id = project.id
         else:
-            project = dt_client.projects.get_by_id(project_id)
-            if project is None:
-                raise DatamintException(f"Project with id '{project_id}' does not exist.")
+            # It's a Project entity
+            project_id = project.id
 
         _ACTIVE_PROJECT_ID = project_id