Compute: add result type to set attributes on blobs (#12051)

GitOrigin-RevId: 07677c4a670d8e9deb9a708c857aae179c3e6695

descarteslabs/core/catalog/catalog_base.py

@@ -12,23 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import wraps
-from types import MethodType
 import json
 import urllib.parse
+from functools import wraps
+from types import MethodType
 
 from descarteslabs.exceptions import NotFoundError
+
 from ..client.deprecation import deprecate
 from ..common.collection import Collection
 from .attributes import (
+    AttributeEqualityMixin,
     AttributeMeta,
     AttributeValidationError,
-    AttributeEqualityMixin,
     CatalogObjectReference,
     DocumentState,
-    Timestamp,
-    ListAttribute,
     ExtraPropertiesAttribute,
+    ListAttribute,
+    Timestamp,
     TypedAttribute,
 )
 from .catalog_client import CatalogClient, HttpRequestMethod
@@ -1026,7 +1027,7 @@ class CatalogObject(CatalogObjectBase):
     extra_properties = ExtraPropertiesAttribute(
         doc="""dict, optional: A dictionary of up to 50 key/value pairs.
 
-        The keys of this dictonary must be strings, and the values of this dictionary
+        The keys of this dictionary must be strings, and the values of this dictionary
         can be strings or numbers.  This allows for more structured custom metadata
         to be associated with objects.
         """

descarteslabs/core/compute/__init__.py

@@ -15,10 +15,11 @@
 from .compute_client import ComputeClient
 from .function import Function, FunctionStatus, Search
 from .job import Job, JobSearch, JobStatus
-from .result import Serializable
+from .result import ComputeResult, Serializable
 
 __all__ = [
     "ComputeClient",
+    "ComputeResult",
     "Function",
     "FunctionStatus",
     "Job",

descarteslabs/core/compute/result.py

@@ -12,6 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set, Union
+
+import shapely.geometry
+
+AnyDate = Union[datetime, str]
+AnyExtraProperties = Dict[str, Union[str, float]]
+AnyGeometry = Union[shapely.geometry.base.BaseGeometry, dict, str]
+AnyResult = Union[bytes, "Serializable", Any]
+AnyTags = Union[Set[str], List[str]]
+
 
 class Serializable:
     def serialize(self) -> bytes:
@@ -20,3 +32,72 @@ def serialize(self) -> bytes:
     @classmethod
     def deserialize(cls, data: bytes):
         raise NotImplementedError()
+
+
+class ComputeResult:
+    def __init__(
+        self,
+        value: Optional[AnyResult] = None,
+        description: Optional[str] = None,
+        expires: Optional[AnyDate] = None,
+        extra_properties: Optional[AnyExtraProperties] = None,
+        geometry: Optional[AnyGeometry] = None,
+        tags: Optional[AnyTags] = None,
+    ):
+        """Used to store the result of a compute function with additional attributes.
+
+        When returned from a compute function, the result will be serialized and stored
+        in the Storage service with the given attributes.
+
+        Parameters
+        ----------
+        value: bytes, Serializable, or Any
+            The resulting value of a compute function.
+            This can be any bytes, any JSON serializable type or any type implementing
+            the Serializable interface.
+        description: str or None
+            A description with further details on this result blob. The description can be
+            up to 80,000 characters and is used by :py:meth:`Search.find_text`.
+        expires: datetime, str, or None
+            The date the result should expire and be deleted from storage.
+            If a string is given, it must be in ISO 8601 format.
+        extra_properties: dict or None
+            A dictionary of up to 50 key/value pairs.
+
+            The keys of this dictionary must be strings, and the values of this dictionary
+            can be strings or numbers.  This allows for more structured custom metadata
+            to be associated with objects.
+        geometry: shapely.geometry.base.BaseGeometry, dict, str, or None
+            The geometry associated with the result if any.
+        tags: set, list, or None
+            The tags to set on the catalog object for the result.
+        """
+        type_ = type(value)
+
+        # If the result is already bytes
+        if isinstance(value, bytes):
+            value = value
+        # If the user implements serialize
+        elif callable(getattr(value, "serialize", None)):
+            value = value.serialize()
+
+            if not isinstance(value, bytes):
+                raise Exception(
+                    f"Serializer on {type_} must return bytes got {type(value)}"
+                )
+        # No specific serialize implementation try json
+        else:
+            try:
+                value = json.dumps(value).encode()
+            except Exception:
+                raise Exception(
+                    "Unable to serialize result. Return value must be"
+                    " JSON encodable or implement the Serializable interface"
+                ) from None
+
+        self.value = value
+        self.description = description
+        self.expires = expires
+        self.extra_properties = extra_properties
+        self.geometry = geometry
+        self.tags = tags