RSDK-6684 - remove binarydata shadow type (viamrobotics#605)

stuqdog · May 8, 2024 · 76d24e9 · 76d24e9
1 parent 96ef235
commit 76d24e9
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 29 deletions.
diff --git a/src/viam/app/data_client.py b/src/viam/app/data_client.py
@@ -14,12 +14,12 @@
     AddBoundingBoxToImageByIDResponse,
     AddTagsToBinaryDataByFilterRequest,
     AddTagsToBinaryDataByIDsRequest,
+    BinaryData,
     BinaryDataByFilterRequest,
     BinaryDataByFilterResponse,
     BinaryDataByIDsRequest,
     BinaryDataByIDsResponse,
     BinaryID,
-    BinaryMetadata,
     BoundingBoxLabelsByFilterRequest,
     BoundingBoxLabelsByFilterResponse,
     CaptureMetadata,
@@ -140,25 +140,6 @@ def __eq__(self, other: object) -> bool:
                 return str(self) == str(other)
             return False
 
-    # TODO (RSDK-6684): Revisit if this shadow type is necessary
-    @dataclass
-    class BinaryData:
-        """Class representing a piece of binary data and associated metadata."""
-
-        data: bytes
-        """The request data"""
-
-        metadata: BinaryMetadata
-        """The metadata associated with the data"""
-
-        def __str__(self) -> str:
-            return f"{self.data}\n{self.metadata}"
-
-        def __eq__(self, other: object) -> bool:
-            if isinstance(other, DataClient.BinaryData):
-                return str(self) == str(other)
-            return False
-
     def __init__(self, channel: Channel, metadata: Mapping[str, str]):
         """Create a `DataClient` that maintains a connection to app.
 
@@ -335,7 +316,7 @@ async def binary_data_by_filter(
             dest (str): Optional filepath for writing retrieved data.
 
         Returns:
-            List[BinaryData]: The binary data.
+            List[viam.proto.app.data.BinaryData]: The binary data.
             int: The count (number of entries)
             str: The last-returned page ID.
         """
@@ -354,7 +335,7 @@ async def binary_data_by_filter(
             include_internal_data=include_internal_data,
         )
         response: BinaryDataByFilterResponse = await self._data_client.BinaryDataByFilter(request, metadata=self._metadata)
-        data = [DataClient.BinaryData(data.binary, data.metadata) for data in response.data]
+        data = list(response.data)
         if dest:
             try:
                 file = open(dest, "w")
@@ -401,7 +382,7 @@ async def binary_data_by_ids(
             GRPCError: If no `BinaryID` objects are provided.
 
         Returns:
-            List[BinaryData]: The binary data.
+            List[viam.proto.app.data.BinaryData]: The binary data.
         """
         request = BinaryDataByIDsRequest(binary_ids=binary_ids, include_binary=True)
         response: BinaryDataByIDsResponse = await self._data_client.BinaryDataByIDs(request, metadata=self._metadata)
@@ -412,7 +393,7 @@ async def binary_data_by_ids(
                 file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write binary data to file {dest}", exc_info=e)
-        return [DataClient.BinaryData(data.binary, data.metadata) for data in response.data]
+        return list(response.data)
 
     async def delete_tabular_data(self, organization_id: str, delete_older_than_days: int) -> int:
         """Delete tabular data older than a specified number of days.

diff --git a/tests/mocks/services.py b/tests/mocks/services.py
@@ -754,7 +754,7 @@ def __init__(
         self,
         tabular_response: List[DataClient.TabularData],
         tabular_query_response: List[Dict[str, ValueTypes]],
-        binary_response: List[DataClient.BinaryData],
+        binary_response: List[BinaryData],
         delete_remove_response: int,
         tags_response: List[str],
         bbox_labels_response: List[str],
@@ -817,7 +817,7 @@ async def BinaryDataByFilter(self, stream: Stream[BinaryDataByFilterRequest, Bin
         self.last = request.data_request.last
         await stream.send_message(
             BinaryDataByFilterResponse(
-                data=[BinaryData(binary=data.data, metadata=data.metadata) for data in self.binary_response],
+                data=self.binary_response,
                 count=len(self.binary_response),
                 last="LAST_BINARY_DATA_PAGE_ID",
             )
@@ -829,7 +829,7 @@ async def BinaryDataByIDs(self, stream: Stream[BinaryDataByIDsRequest, BinaryDat
         assert request is not None
         self.binary_ids = request.binary_ids
         await stream.send_message(
-            BinaryDataByIDsResponse(data=[BinaryData(binary=data.data, metadata=data.metadata) for data in self.binary_response])
+            BinaryDataByIDsResponse(data=self.binary_response)
         )
 
     async def DeleteTabularData(self, stream: Stream[DeleteTabularDataRequest, DeleteTabularDataResponse]) -> None:

diff --git a/tests/test_data_client.py b/tests/test_data_client.py
@@ -5,7 +5,7 @@
 from grpclib.testing import ChannelFor
 
 from viam.app.data_client import DataClient
-from viam.proto.app.data import Annotations, BinaryID, BinaryMetadata, BoundingBox, CaptureMetadata, Filter, Order
+from viam.proto.app.data import Annotations, BinaryData, BinaryID, BinaryMetadata, BoundingBox, CaptureMetadata, Filter, Order
 from viam.utils import create_filter
 
 from .mocks.services import MockData
@@ -103,7 +103,7 @@
 TABULAR_QUERY_RESPONSE = [
     {"key1": 1, "key2": "2", "key3": [1, 2, 3], "key4": {"key4sub1": 1}},
 ]
-BINARY_RESPONSE = [DataClient.BinaryData(BINARY_DATA, BINARY_METADATA)]
+BINARY_RESPONSE = [BinaryData(binary=BINARY_DATA, metadata=BINARY_METADATA)]
 DELETE_REMOVE_RESPONSE = 1
 TAGS_RESPONSE = ["tag"]
 HOSTNAME_RESPONSE = "host"