albert 1.9.5__py3-none-any.whl → 1.10.0__py3-none-any.whl

albert/__init__.py

@@ -4,4 +4,4 @@ from albert.core.auth.sso import AlbertSSOClient
 
 __all__ = ["Albert", "AlbertClientCredentials", "AlbertSSOClient"]
 
-__version__ = "1.9.5"
+__version__ = "1.10.0"

albert/collections/attachments.py

@@ -9,7 +9,7 @@ from pydantic import validate_call
 from albert.collections.base import BaseCollection
 from albert.collections.files import FileCollection
 from albert.collections.notes import NotesCollection
-from albert.core.shared.identifiers import AttachmentId, InventoryId
+from albert.core.shared.identifiers import AttachmentId, DataColumnId, InventoryId
 from albert.core.shared.types import MetadataItem
 from albert.resources.attachments import Attachment, AttachmentCategory
 from albert.resources.files import FileCategory, FileNamespace
@@ -49,7 +49,9 @@ class AttachmentCollection(BaseCollection):
         response = self.session.get(url=f"{self.base_path}/{id}")
         return Attachment(**response.json())
 
-    def get_by_parent_ids(self, *, parent_ids: list[str]) -> dict[str, list[Attachment]]:
+    def get_by_parent_ids(
+        self, *, parent_ids: list[str], data_column_ids: list[DataColumnId] | None = None
+    ) -> dict[str, list[Attachment]]:
         """Retrieves attachments by their parent IDs.
 
         Note: This method returns a dictionary where the keys are parent IDs
@@ -69,7 +71,10 @@ class AttachmentCollection(BaseCollection):
         dict[str, list[Attachment]]
             A dictionary mapping parent IDs to lists of Attachment objects associated with each parent ID.
         """
-        response = self.session.get(url=f"{self.base_path}/parents", params={"id": parent_ids})
+        response = self.session.get(
+            url=f"{self.base_path}/parents",
+            params={"id": parent_ids, "dataColumnId": data_column_ids},
+        )
         response_data = response.json()
         return {
             parent["parentId"]: [
@@ -125,7 +130,12 @@ class AttachmentCollection(BaseCollection):
         self.session.delete(f"{self.base_path}/{id}")
 
     def upload_and_attach_file_as_note(
-        self, parent_id: str, file_data: IO, note_text: str = "", file_name: str = ""
+        self,
+        parent_id: str,
+        file_data: IO,
+        note_text: str = "",
+        file_name: str = "",
+        upload_key: str | None = None,
     ) -> Note:
         """Uploads a file and attaches it to a new note. A user can be tagged in the note_text string by using f-string and the User.to_note_mention() method.
         This allows for easy tagging and referencing of users within notes. example: f"Hello {tagged_user.to_note_mention()}!"
@@ -140,24 +150,31 @@ class AttachmentCollection(BaseCollection):
             Any additional text to add to the note, by default ""
         file_name : str, optional
             The name of the file, by default ""
+        upload_key : str | None, optional
+            Override the storage key used when signing and uploading the file.
+            Defaults to the provided ``file_name``.
 
         Returns
         -------
         Note
             The created note.
         """
-        file_type = mimetypes.guess_type(file_name)[0]
+        upload_name = upload_key or file_name
+        if not upload_name:
+            raise ValueError("A file name or upload key must be provided for attachment upload.")
+
+        file_type = mimetypes.guess_type(file_name or upload_name)[0]
         file_collection = self._get_file_collection()
         note_collection = self._get_note_collection()
 
         file_collection.sign_and_upload_file(
             data=file_data,
-            name=file_name,
+            name=upload_name,
             namespace=FileNamespace.RESULT.value,
             content_type=file_type,
         )
         file_info = file_collection.get_by_name(
-            name=file_name, namespace=FileNamespace.RESULT.value
+            name=upload_name, namespace=FileNamespace.RESULT.value
         )
         note = Note(
             parent_id=parent_id,
@@ -166,7 +183,7 @@ class AttachmentCollection(BaseCollection):
         registered_note = note_collection.create(note=note)
         self.attach_file_to_note(
             note_id=registered_note.id,
-            file_name=file_name,
+            file_name=file_name or Path(upload_name).name,
             file_key=file_info.name,
         )
         return note_collection.get_by_id(id=registered_note.id)

albert/collections/custom_templates.py

@@ -1,5 +1,7 @@
 from collections.abc import Iterator
 
+from pydantic import validate_call
+
 from albert.collections.base import BaseCollection
 from albert.core.logging import logger
 from albert.core.pagination import AlbertPaginator
@@ -28,6 +30,7 @@ class CustomTemplatesCollection(BaseCollection):
         super().__init__(session=session)
         self.base_path = f"/api/{CustomTemplatesCollection._api_version}/customtemplates"
 
+    @validate_call
     def get_by_id(self, *, id: CustomTemplateId) -> CustomTemplate:
         """Get a Custom Template by ID
 

albert/collections/data_templates.py

@@ -8,7 +8,7 @@ from albert.core.logging import logger
 from albert.core.pagination import AlbertPaginator
 from albert.core.session import AlbertSession
 from albert.core.shared.enums import OrderBy, PaginationMode
-from albert.core.shared.identifiers import DataTemplateId, UserId
+from albert.core.shared.identifiers import DataColumnId, DataTemplateId, UserId
 from albert.core.shared.models.patch import (
     GeneralPatchDatum,
     GeneralPatchPayload,
@@ -17,13 +17,21 @@ from albert.core.shared.models.patch import (
 )
 from albert.exceptions import AlbertHTTPError
 from albert.resources.data_templates import (
+    CurveExample,
     DataColumnValue,
     DataTemplate,
     DataTemplateSearchItem,
+    ImageExample,
     ParameterValue,
 )
 from albert.resources.parameter_groups import DataType, EnumValidationValue, ValueValidation
 from albert.utils._patch import generate_data_template_patches
+from albert.utils.data_template import (
+    add_parameter_enums,
+    build_curve_example,
+    build_image_example,
+    get_target_data_column,
+)
 
 
 class DCPatchDatum(PGPatchPayload):
@@ -82,102 +90,6 @@ class DataTemplateCollection(BaseCollection):
         else:
             return self.add_parameters(data_template_id=dt.id, parameters=parameter_values)
 
-    def _add_param_enums(
-        self,
-        *,
-        data_template_id: DataTemplateId,
-        new_parameters: list[ParameterValue],
-    ) -> list[EnumValidationValue]:
-        """Adds enum values to a parameter."""
-
-        data_template = self.get_by_id(id=data_template_id)
-        existing_parameters = data_template.parameter_values
-        enums_by_sequence = {}
-        for parameter in new_parameters:
-            this_sequence = next(
-                (
-                    p.sequence
-                    for p in existing_parameters
-                    if p.id == parameter.id and p.short_name == parameter.short_name
-                ),
-                None,
-            )
-            enum_patches = []
-            if (
-                parameter.validation
-                and len(parameter.validation) > 0
-                and isinstance(parameter.validation[0].value, list)
-            ):
-                existing_validation = (
-                    [x for x in existing_parameters if x.sequence == parameter.sequence]
-                    if existing_parameters
-                    else []
-                )
-                existing_enums = (
-                    [
-                        x
-                        for x in existing_validation[0].validation[0].value
-                        if isinstance(x, EnumValidationValue) and x.id is not None
-                    ]
-                    if (
-                        existing_validation
-                        and len(existing_validation) > 0
-                        and existing_validation[0].validation
-                        and len(existing_validation[0].validation) > 0
-                        and existing_validation[0].validation[0].value
-                        and isinstance(existing_validation[0].validation[0].value, list)
-                    )
-                    else []
-                )
-                updated_enums = (
-                    [
-                        x
-                        for x in parameter.validation[0].value
-                        if isinstance(x, EnumValidationValue)
-                    ]
-                    if parameter.validation[0].value
-                    else []
-                )
-
-                deleted_enums = [
-                    x for x in existing_enums if x.id not in [y.id for y in updated_enums]
-                ]
-
-                new_enums = [
-                    x for x in updated_enums if x.id not in [y.id for y in existing_enums]
-                ]
-
-                matching_enums = [
-                    x for x in updated_enums if x.id in [y.id for y in existing_enums]
-                ]
-
-                for new_enum in new_enums:
-                    enum_patches.append({"operation": "add", "text": new_enum.text})
-                for deleted_enum in deleted_enums:
-                    enum_patches.append({"operation": "delete", "id": deleted_enum.id})
-                for matching_enum in matching_enums:
-                    if (
-                        matching_enum.text
-                        != [x for x in existing_enums if x.id == matching_enum.id][0].text
-                    ):
-                        enum_patches.append(
-                            {
-                                "operation": "update",
-                                "id": matching_enum.id,
-                                "text": matching_enum.text,
-                            }
-                        )
-
-                if len(enum_patches) > 0:
-                    enum_response = self.session.put(
-                        f"{self.base_path}/{data_template_id}/parameters/{this_sequence}/enums",
-                        json=enum_patches,
-                    )
-                    enums_by_sequence[this_sequence] = [
-                        EnumValidationValue(**x) for x in enum_response.json()
-                    ]
-        return enums_by_sequence
-
     @validate_call
     def get_by_id(self, *, id: DataTemplateId) -> DataTemplate:
         """Get a data template by its ID.
@@ -195,6 +107,7 @@ class DataTemplateCollection(BaseCollection):
         response = self.session.get(f"{self.base_path}/{id}")
         return DataTemplate(**response.json())
 
+    @validate_call
     def get_by_ids(self, *, ids: list[DataTemplateId]) -> list[DataTemplate]:
         """Get a list of data templates by their IDs.
 
@@ -234,6 +147,7 @@ class DataTemplateCollection(BaseCollection):
                 return t.hydrate()
         return None
 
+    @validate_call
     def add_data_columns(
         self, *, data_template_id: DataTemplateId, data_columns: list[DataColumnValue]
     ) -> DataTemplate:
@@ -277,6 +191,7 @@ class DataTemplateCollection(BaseCollection):
         )
         return self.get_by_id(id=data_template_id)
 
+    @validate_call
     def add_parameters(
         self, *, data_template_id: DataTemplateId, parameters: list[ParameterValue]
     ) -> DataTemplate:
@@ -325,7 +240,9 @@ class DataTemplateCollection(BaseCollection):
             if param.id in initial_enum_values:
                 param.validation[0].value = initial_enum_values[param.id]
                 param.validation[0].datatype = DataType.ENUM
-                self._add_param_enums(
+                add_parameter_enums(
+                    session=self.session,
+                    base_path=self.base_path,
                     data_template_id=data_template_id,
                     new_parameters=[param],
                 )
@@ -396,6 +313,12 @@ class DataTemplateCollection(BaseCollection):
         -------
         DataTemplate
             The Updated DataTemplate object.
+
+        Warnings
+        --------
+        Only scalar data column values (text, number, dropdown) can be updated using this function. Use
+        `set_curve_example` / `set_image_example` to set example values for other data column types.
+
         """
 
         existing = self.get_by_id(id=data_template.id)
@@ -471,7 +394,9 @@ class DataTemplateCollection(BaseCollection):
                     param.validation[0].datatype = DataType.ENUM  # Add this line
 
             # Add enum values to newly created parameters
-            self._add_param_enums(
+            add_parameter_enums(
+                session=self.session,
+                base_path=self.base_path,
                 data_template_id=existing.id,
                 new_parameters=returned_parameters,
             )
@@ -555,6 +480,7 @@ class DataTemplateCollection(BaseCollection):
             )
         return self.get_by_id(id=data_template.id)
 
+    @validate_call
     def delete(self, *, id: DataTemplateId) -> None:
         """Deletes a data template by its ID.
 
@@ -623,3 +549,99 @@ class DataTemplateCollection(BaseCollection):
                 yield from hydrated_templates
             except AlbertHTTPError as e:
                 logger.warning(f"Error hydrating batch {batch}: {e}")
+
+    @validate_call
+    def set_curve_example(
+        self,
+        *,
+        data_template_id: DataTemplateId,
+        data_column_id: DataColumnId | None = None,
+        data_column_name: str | None = None,
+        example: CurveExample,
+    ) -> DataTemplate:
+        """Set a curve example on a Curve data column.
+
+        Parameters
+        ----------
+        data_template_id : DataTemplateId
+            Target data template ID.
+        data_column_id : DataColumnId, optional
+            Target curve column ID (provide exactly one of id or name).
+        data_column_name : str, optional
+            Target curve column name (provide exactly one of id or name).
+        example : CurveExample
+            Curve example payload
+
+        Returns
+        -------
+        DataTemplate
+            The updated data template after the example is applied.
+        """
+        data_template = self.get_by_id(id=data_template_id)
+        target_column = get_target_data_column(
+            data_template=data_template,
+            data_template_id=data_template_id,
+            data_column_id=data_column_id,
+            data_column_name=data_column_name,
+        )
+        payload = build_curve_example(
+            session=self.session,
+            data_template_id=data_template_id,
+            example=example,
+            target_column=target_column,
+        )
+        if not payload.data:
+            return data_template
+        self.session.patch(
+            f"{self.base_path}/{data_template_id}",
+            json=payload.model_dump(mode="json", by_alias=True, exclude_none=True),
+        )
+        return self.get_by_id(id=data_template_id)
+
+    @validate_call
+    def set_image_example(
+        self,
+        *,
+        data_template_id: DataTemplateId,
+        data_column_id: DataColumnId | None = None,
+        data_column_name: str | None = None,
+        example: ImageExample,
+    ) -> DataTemplate:
+        """Set an image example on a Image data column.
+
+        Parameters
+        ----------
+        data_template_id : DataTemplateId
+            Target data template ID.
+        data_column_id : DataColumnId, optional
+            Target image column ID (provide exactly one of id or name).
+        data_column_name : str, optional
+            Target image column name (provide exactly one of id or name).
+        example : ImageExample
+            Image example payload
+
+        Returns
+        -------
+        DataTemplate
+            The updated data template after the example is applied.
+        """
+        data_template = self.get_by_id(id=data_template_id)
+        target_column = get_target_data_column(
+            data_template=data_template,
+            data_template_id=data_template_id,
+            data_column_id=data_column_id,
+            data_column_name=data_column_name,
+        )
+        payload = build_image_example(
+            session=self.session,
+            data_template_id=data_template_id,
+            example=example,
+            target_column=target_column,
+        )
+        if not payload.data:
+            return data_template
+        self.session.patch(
+            f"{self.base_path}/{data_template_id}",
+            json=payload.model_dump(mode="json", by_alias=True, exclude_none=True),
+        )
+        return self.get_by_id(id=data_template_id)

albert/collections/entity_types.py

@@ -1,5 +1,7 @@
 from collections.abc import Iterator
 
+from pydantic import validate_call
+
 from albert.collections.base import BaseCollection
 from albert.core.pagination import AlbertPaginator, PaginationMode
 from albert.core.session import AlbertSession
@@ -41,6 +43,7 @@ class EntityTypeCollection(BaseCollection):
         super().__init__(session=session)
         self.base_path = f"/api/{EntityTypeCollection._api_version}/entitytypes"
 
+    @validate_call
     def get_by_id(self, *, id: EntityTypeId) -> EntityType:
         """Get an entity type by its ID.
         Parameters
@@ -203,6 +206,7 @@ class EntityTypeCollection(BaseCollection):
 
         return patches
 
+    @validate_call
     def delete(self, *, id: EntityTypeId) -> None:
         """Delete an entity type.
         Parameters
@@ -212,6 +216,7 @@ class EntityTypeCollection(BaseCollection):
         """
         self.session.delete(f"{self.base_path}/{id}")
 
+    @validate_call
     def get_rules(self, *, id: EntityTypeId) -> list[EntityTypeRule]:
         """Get the rules for an entity type.
         Parameters
@@ -222,6 +227,7 @@ class EntityTypeCollection(BaseCollection):
         response = self.session.get(f"{self.base_path}/rules/{id}")
         return [EntityTypeRule(**rule) for rule in response.json()]
 
+    @validate_call
     def set_rules(self, *, id: EntityTypeId, rules: list[EntityTypeRule]) -> list[EntityTypeRule]:
         """Create or update the rules for an entity type.
         Parameters
@@ -241,6 +247,7 @@ class EntityTypeCollection(BaseCollection):
         )
         return [EntityTypeRule(**rule) for rule in response.json()]
 
+    @validate_call
     def delete_rules(self, *, id: EntityTypeId) -> None:
         """Delete the rules for an entity type.
         Parameters
@@ -263,12 +270,12 @@ class EntityTypeCollection(BaseCollection):
         ----------
         service : EntityServiceType | None, optional
             The service type the entity type is associated with, by default None
-        limit : int, optional
-            Maximum number of results to return, by default 100
         start_key : str | None, optional
             Key to start pagination from, by default None
         order : OrderBy | None, optional
             Sort order (ascending/descending), by default None
+        max_items : int | None, optional
+            Maximum number of items to return, by default None
         Yields
         ------
         Iterator[EntityType]

albert/collections/inventory.py

@@ -281,7 +281,7 @@ class InventoryCollection(BaseCollection):
         """Add inventory specs to the inventory item.
 
         An `InventorySpec` is a property that was not directly measured via a task,
-        but is a generic property of that inentory item.
+        but is a generic property of that inventory item.
 
         Parameters
         ----------

albert/collections/parameters.py

@@ -98,6 +98,7 @@ class ParameterCollection(BaseCollection):
         url = f"{self.base_path}/{id}"
         self.session.delete(url)
 
+    @validate_call
     def get_all(
         self,
         *,