Call PartialUpdateElement but correct payload

Complements #222 (closed)

This endpoint helper was incomplete as many fields were missing and the payload sent to the server was not right. PartialUpdateElement should be called only with the fields to update. All others were overwritten.

arkindex_worker/worker/element.py

@@ -3,6 +3,7 @@
 ElementsWorker methods for elements and element types.
 """
 from typing import Dict, Iterable, List, NamedTuple, Optional, Union
+from uuid import UUID
 
 from peewee import IntegrityError
 
@@ -275,33 +276,39 @@ class ElementMixin(object):
 
         return created_ids
 
-    def update_element(
-        self,
-        element: Union[Element, CachedElement],
-        type: Optional[str] = None,
-        name: Optional[str] = None,
-        polygon: Optional[List[List[Union[int, float]]]] = None,
-        confidence: Optional[float] = None,
+    def partial_update_element(
+        self, element: Union[Element, CachedElement], **kwargs
     ) -> dict:
         """
-        Partially update an element through the API.
+        Partially updates an element through the API.
 
         :param element: The element to update.
-        :param type: Optional new slug type of the element.
-        :param name: Optional new name of the element.
-        :param polygon: Optional new polygon of the element.
-        :param confidence: Optional new confidence score, between 0.0 and 1.0.
+        :param **kwargs:
+
+            * *type* (``str``): Optional slug type of the element.
+            * *name* (``str``): Optional name of the element.
+            * *polygon* (``list``): Optional polygon for this element
+            * *confidence* (``float``): Optional confidence score of this element
+            * *rotation_angle* (``int``): Optional rotation angle of this element
+            * *mirrored* (``bool``): Optional mirror status of this element
+            * *image* (``UUID``): Optional ID of the image of this element
+
+
         :returns: A dict from the ``PartialUpdateElement`` API endpoint,
         """
         assert element and isinstance(
             element, (Element, CachedElement)
         ), "element shouldn't be null and should be an Element or CachedElement"
-        assert type is None or isinstance(type, str), "type should be None or a str"
-        assert name is None or isinstance(name, str), "name should be None or a str"
-        assert polygon is None or isinstance(
-            polygon, list
-        ), "polygon should be None or a list"
-        if polygon:
+
+        if "type" in kwargs:
+            assert isinstance(kwargs["type"], str), "type should be a str"
+
+        if "name" in kwargs:
+            assert isinstance(kwargs["name"], str), "name should be a str"
+
+        if "polygon" in kwargs:
+            polygon = kwargs["polygon"]
+            assert isinstance(polygon, list), "polygon should be a list"
             assert len(polygon) >= 3, "polygon should have at least three points"
             assert all(
                 isinstance(point, list) and len(point) == 2 for point in polygon
@@ -309,9 +316,27 @@ class ElementMixin(object):
             assert all(
                 isinstance(coord, (int, float)) for point in polygon for coord in point
             ), "polygon points should be lists of two numbers"
-        assert confidence is None or (
-            isinstance(confidence, float) and 0 <= confidence <= 1
-        ), "confidence should be None or a float in [0..1] range"
+
+        if "confidence" in kwargs:
+            confidence = kwargs["confidence"]
+            assert confidence is None or (
+                isinstance(confidence, float) and 0 <= confidence <= 1
+            ), "confidence should be None or a float in [0..1] range"
+
+        if "rotation_angle" in kwargs:
+            rotation_angle = kwargs["rotation_angle"]
+            assert (
+                isinstance(rotation_angle, int) and rotation_angle >= 0
+            ), "rotation_angle should be a positive integer"
+
+        if "mirrored" in kwargs:
+            assert isinstance(kwargs["mirrored"], bool), "mirrored should be a boolean"
+
+        if "image" in kwargs:
+            image = kwargs["image"]
+            assert isinstance(image, UUID), "image should be a UUID"
+            # Cast to string
+            kwargs["image"] = str(image)
 
         if self.is_read_only:
             logger.warning("Cannot update element as this worker is in read-only mode")
@@ -320,22 +345,24 @@ class ElementMixin(object):
         updated_element = self.request(
             "PartialUpdateElement",
             id=element.id,
-            body={
-                "type": type,
-                "name": name,
-                "polygon": polygon,
-                "confidence": confidence,
-            },
+            body=kwargs,
         )
 
         if self.use_cache:
-            CachedElement.update(
-                {
-                    CachedElement.type: type,
-                    CachedElement.polygon: str(polygon),
-                    CachedElement.confidence: confidence,
-                }
-            ).where(CachedElement.id == element.id).execute()
+            # Name is not present in CachedElement model
+            kwargs.pop("name", None)
+
+            # Stringify polygon if present
+            if "polygon" in kwargs:
+                kwargs["polygon"] = str(kwargs["polygon"])
+
+            # Retrieve the right image
+            if "image" in kwargs:
+                kwargs["image"] = CachedImage.get_by_id(kwargs["image"])
+
+            CachedElement.update(**kwargs).where(
+                CachedElement.id == element.id
+            ).execute()
 
         return updated_element