RSDK-5936: change return type for get readings (#506)

purplenicole730 · web-flow · commit 63c03d80844b · 2023-12-27T11:58:45.000-05:00
diff --git a/docs/examples/example.ipynb b/docs/examples/example.ipynb
@@ -272,13 +272,14 @@
     "\n",
     "from viam.components.sensor import Sensor\n",
     "from viam.logging import getLogger\n",
+    "from viam.utils import SensorReading\n",
     "\n",
     "LOGGER = getLogger(__name__)\n",
     "\n",
     "\n",
     "class MySensor(Sensor):\n",
     "    # Subclass the Viam Sensor component and implement the required functions\n",
-    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:\n",
+    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:\n",
     "        with open(\"/proc/net/wireless\") as wifi_stats:\n",
     "            content = wifi_stats.readlines()\n",
     "        wifi_signal = [x for x in content[2].split(\" \") if x != \"\"]\n",
@@ -368,6 +369,8 @@
     "# ADD `Sequence` FROM `typing`.\n",
     "from typing import Sequence\n",
     "\n",
+    "from viam.utils import SensorReading\n",
+    "\n",
     "class MySensor(Sensor):\n",
     "    # ADD A VALIDATOR FUNCTION \n",
     "    @classmethod\n",
@@ -395,7 +398,7 @@
     "            multiplier = 1.0\n",
     "        self.multiplier = multiplier\n",
     "\n",
-    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:\n",
+    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:\n",
     "        with open(\"/proc/net/wireless\") as wifi_stats:\n",
     "            content = wifi_stats.readlines()\n",
     "        result = [x for x in content[2].split(\" \") if x != \"\"]\n",
@@ -1031,10 +1034,11 @@
     "\n",
     "from viam.components.sensor import Geometry, Sensor\n",
     "from viam.rpc.server import Server\n",
+    "from viam.utils import SensorReading\n",
     "\n",
     "\n",
     "class MySensor(Sensor):\n",
-    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:\n",
+    "    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:\n",
     "        with open(\"/proc/net/wireless\") as wifi_stats:\n",
     "            content = wifi_stats.readlines()\n",
     "        wifi_signal = [x for x in content[2].split(\" \") if x != \"\"]\n",
@@ -1485,7 +1489,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4 (main, Jun 20 2023, 17:23:00) [Clang 14.0.3 (clang-1403.0.22.14.1)]"
+   "version": "3.11.6 (main, Oct  2 2023, 20:46:14) [Clang 14.0.3 (clang-1403.0.22.14.1)]"
   },
   "vscode": {
    "interpreter": {
diff --git a/docs/examples/module_step2.py b/docs/examples/module_step2.py
@@ -11,6 +11,7 @@
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
+from viam.utils import SensorReading
 
 LOGGER = getLogger(__name__)
 
@@ -24,7 +25,7 @@ def new(cls, config: ComponentConfig, dependencies: Mapping[ResourceName, Resour
         sensor = cls(config.name)
         return sensor
 
-    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:
         with open("/proc/net/wireless") as wifi_stats:
             content = wifi_stats.readlines()
         wifi_signal = [x for x in content[2].split(" ") if x != ""]
diff --git a/docs/examples/module_step2_optional.py b/docs/examples/module_step2_optional.py
@@ -11,6 +11,7 @@
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
+from viam.utils import SensorReading
 
 LOGGER = getLogger(__name__)
 
@@ -35,7 +36,7 @@ def validate_config(cls, config: ComponentConfig) -> Sequence[str]:
                 raise Exception("Multiplier cannot be 0.")
         return []
 
-    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:
         with open("/proc/net/wireless") as wifi_stats:
             content = wifi_stats.readlines()
         result = [x for x in content[2].split(" ") if x != ""]
diff --git a/docs/examples/module_step3.py b/docs/examples/module_step3.py
@@ -12,6 +12,7 @@
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
+from viam.utils import SensorReading
 
 LOGGER = getLogger(__name__)
 
@@ -25,7 +26,7 @@ def new(cls, config: ComponentConfig, dependencies: Mapping[ResourceName, Resour
         sensor = cls(config.name)
         return sensor
 
-    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:
         with open("/proc/net/wireless") as wifi_stats:
             content = wifi_stats.readlines()
         wifi_signal = [x for x in content[2].split(" ") if x != ""]
diff --git a/examples/server/v1/components.py b/examples/server/v1/components.py
@@ -53,6 +53,7 @@
 from viam.proto.component.arm import JointPositions
 from viam.proto.component.audioinput import AudioChunk, AudioChunkInfo, SampleFormat
 from viam.proto.component.encoder import PositionType
+from viam.utils import SensorReading
 
 GEOMETRIES = [
     Geometry(center=Pose(x=1, y=2, z=3, o_x=2, o_y=3, o_z=4, theta=20), sphere=Sphere(radius_mm=2)),
@@ -645,7 +646,7 @@ def __init__(
         self.properties = properties
         self.accuracy = accuracy
 
-    async def get_readings(self, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, **kwargs) -> Mapping[str, SensorReading]:
         return {"abcdefghij"[idx]: random.random() for idx in range(self.num_readings)}
 
     async def get_position(self, **kwargs) -> Tuple[GeoPoint, float]:
@@ -693,7 +694,7 @@ def __init__(self, name: str):
         self.num_readings = random.randint(1, 10)
         super().__init__(name)
 
-    async def get_readings(self, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, **kwargs) -> Mapping[str, SensorReading]:
         return {"abcdefghij"[idx]: random.random() for idx in range(self.num_readings)}
 
     async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> List[Geometry]:
diff --git a/examples/simple_module/src/main.py b/examples/simple_module/src/main.py
@@ -11,7 +11,7 @@
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
-from viam.utils import ValueTypes
+from viam.utils import SensorReading, ValueTypes
 
 LOGGER = getLogger(__name__)
 
@@ -39,7 +39,7 @@ def validate_config(cls, config: ComponentConfig) -> Sequence[str]:
                 raise Exception("Multiplier cannot be 0.")
         return []
 
-    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Mapping[str, SensorReading]:
         return {"signal": 1 * self.multiplier}
 
     async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None, **kwargs) -> Mapping[str, ValueTypes]:
diff --git a/src/viam/components/movement_sensor/client.py b/src/viam/components/movement_sensor/client.py
@@ -24,7 +24,7 @@
     MovementSensorServiceStub,
 )
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, get_geometries, struct_to_dict
+from viam.utils import SensorReading, ValueTypes, dict_to_struct, get_geometries, sensor_readings_value_to_native, struct_to_dict
 
 from . import GeoPoint, Orientation, Vector3
 
@@ -93,12 +93,13 @@ async def get_accuracy(self, *, extra: Optional[Dict[str, Any]] = None, timeout:
         response: GetAccuracyResponse = await self.client.GetAccuracy(request, timeout=timeout)
         return response.accuracy
 
-    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, Any]:
+    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, SensorReading]:
         if extra is None:
             extra = {}
         request = GetReadingsRequest(name=self.name, extra=dict_to_struct(extra))
         response: GetReadingsResponse = await self.client.GetReadings(request, timeout=timeout)
-        return response.readings
+
+        return sensor_readings_value_to_native(response.readings)
 
     async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None) -> Mapping[str, ValueTypes]:
         request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
diff --git a/src/viam/components/movement_sensor/movement_sensor.py b/src/viam/components/movement_sensor/movement_sensor.py
@@ -7,6 +7,7 @@
 from viam.components.component_base import ComponentBase
 from viam.proto.component.movementsensor import GetPropertiesResponse
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
+from viam.utils import SensorReading
 
 from . import GeoPoint, Orientation, Vector3
 
@@ -129,7 +130,9 @@ async def get_accuracy(
         """
         ...
 
-    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Mapping[str, SensorReading]:
         """Obtain the measurements/data specific to this sensor.
         If a sensor is not configured to have a measurement or fails to read a piece of data, it will not appear in the readings dictionary.
 
diff --git a/src/viam/components/power_sensor/client.py b/src/viam/components/power_sensor/client.py
@@ -14,7 +14,7 @@
     PowerSensorServiceStub,
 )
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
+from viam.utils import SensorReading, ValueTypes, dict_to_struct, sensor_readings_value_to_native, struct_to_dict
 
 
 class PowerSensorClient(PowerSensor, ReconfigurableResourceRPCClientBase):
@@ -46,12 +46,12 @@ async def get_power(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Op
         response: GetPowerResponse = await self.client.GetPower(request, timeout=timeout)
         return response.watts
 
-    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, Any]:
+    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, SensorReading]:
         if extra is None:
             extra = {}
         request = GetReadingsRequest(name=self.name, extra=dict_to_struct(extra))
         response: GetReadingsResponse = await self.client.GetReadings(request, timeout=timeout)
-        return response.readings
+        return sensor_readings_value_to_native(response.readings)
 
     async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None) -> Mapping[str, ValueTypes]:
         request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
diff --git a/src/viam/components/power_sensor/power_sensor.py b/src/viam/components/power_sensor/power_sensor.py
@@ -3,6 +3,7 @@
 
 from viam.components.component_base import ComponentBase
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
+from viam.utils import SensorReading
 
 
 class PowerSensor(ComponentBase):
@@ -41,7 +42,9 @@ async def get_power(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Op
         """
         ...
 
-    async def get_readings(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> Mapping[str, Any]:
+    async def get_readings(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Mapping[str, SensorReading]:
         """Obtain the measurements/data specific to this sensor.
         If a sensor is not configured to have a measurement or fails to read a piece of data, it will not appear in the readings dictionary.
 
diff --git a/src/viam/components/sensor/client.py b/src/viam/components/sensor/client.py
@@ -5,7 +5,7 @@
 from viam.proto.common import DoCommandRequest, DoCommandResponse, Geometry, GetReadingsRequest, GetReadingsResponse
 from viam.proto.component.sensor import SensorServiceStub
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, get_geometries, sensor_readings_value_to_native, struct_to_dict
+from viam.utils import SensorReading, ValueTypes, dict_to_struct, get_geometries, sensor_readings_value_to_native, struct_to_dict
 
 from .sensor import Sensor
 
@@ -20,7 +20,9 @@ def __init__(self, name: str, channel: Channel):
         self.client = SensorServiceStub(channel)
         super().__init__(name)
 
-    async def get_readings(self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, Any]:
+    async def get_readings(
+        self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None
+    ) -> Mapping[str, SensorReading]:
         if extra is None:
             extra = {}
         request = GetReadingsRequest(name=self.name, extra=dict_to_struct(extra))
diff --git a/src/viam/components/sensor/sensor.py b/src/viam/components/sensor/sensor.py
@@ -2,6 +2,7 @@
 from typing import Any, Mapping, Optional
 
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
+from viam.utils import SensorReading
 
 from ..component_base import ComponentBase
 
@@ -20,7 +21,7 @@ class Sensor(ComponentBase):
     @abc.abstractmethod
     async def get_readings(
         self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None, **kwargs
-    ) -> Mapping[str, Any]:
+    ) -> Mapping[str, SensorReading]:
         """
         Obtain the measurements/data specific to this sensor.
 
diff --git a/src/viam/services/sensors/client.py b/src/viam/services/sensors/client.py
@@ -7,7 +7,7 @@
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_SERVICE, Subtype
 from viam.services.service_client_base import ServiceClientBase
-from viam.utils import ValueTypes, dict_to_struct, sensor_readings_value_to_native, struct_to_dict
+from viam.utils import SensorReading, ValueTypes, dict_to_struct, sensor_readings_value_to_native, struct_to_dict
 
 
 class SensorsClient(ServiceClientBase, ReconfigurableResourceRPCClientBase):
@@ -34,7 +34,7 @@ async def get_sensors(self, *, extra: Optional[Mapping[str, Any]] = None, timeou
 
     async def get_readings(
         self, sensors: List[ResourceName], *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None
-    ) -> Mapping[ResourceName, Mapping[str, Any]]:
+    ) -> Mapping[ResourceName, Mapping[str, SensorReading]]:
         """Get the readings from the specific sensors provided
 
         Args:
diff --git a/src/viam/utils.py b/src/viam/utils.py
@@ -31,6 +31,9 @@
 ValueTypes = Union[bool, SupportsBytes, SupportsFloat, List, Mapping, str, None]
 """Types that can be encoded into a protobuf `Value`"""
 
+SensorReading = Union[ValueTypes, Vector3, GeoPoint, Orientation]
+"""Types that can be returned from a sensor"""
+
 
 def primitive_to_value(v: ValueTypes) -> Value:
     """
@@ -172,7 +175,7 @@ async def get_geometries(
     return [geometry for geometry in response.geometries]
 
 
-def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str, Any]:
+def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str, Value]:
     prim_readings = dict(readings)
     for key, reading in readings.items():
         if isinstance(reading, Vector3):
@@ -190,7 +193,7 @@ def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str,
     return {key: primitive_to_value(value) for (key, value) in prim_readings.items()}
 
 
-def sensor_readings_value_to_native(readings: Mapping[str, Value]) -> Mapping[str, Any]:
+def sensor_readings_value_to_native(readings: Mapping[str, Value]) -> Mapping[str, SensorReading]:
     prim_readings: Dict[str, Any] = {key: value_to_primitive(value) for (key, value) in readings.items()}
     for key, reading in prim_readings.items():
         if isinstance(reading, Mapping):
diff --git a/tests/mocks/components.py b/tests/mocks/components.py
@@ -50,7 +50,7 @@
 from viam.proto.component.audioinput import AudioChunk, AudioChunkInfo, SampleFormat
 from viam.proto.component.board import PowerMode
 from viam.proto.component.encoder import PositionType
-from viam.utils import ValueTypes
+from viam.utils import SensorReading, ValueTypes
 
 GEOMETRIES = [
     Geometry(center=Pose(x=1, y=2, z=3, o_x=2, o_y=3, o_z=4, theta=20), sphere=Sphere(radius_mm=2)),
@@ -991,7 +991,7 @@ def __init__(self, name: str, result: Mapping[str, Any] = {"a": 0, "b": {"foo":
 
     async def get_readings(
         self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None, **kwargs
-    ) -> Mapping[str, Any]:
+    ) -> Mapping[str, SensorReading]:
         self.extra = extra
         self.timeout = timeout
         return self.readings
diff --git a/tests/test_board.py b/tests/test_board.py
@@ -175,15 +175,11 @@ async def test_get_digital_interrupt_value(self, board: MockBoard, service: Boar
         async with ChannelFor([service]) as channel:
             client = BoardServiceStub(channel)
 
-            request = GetDigitalInterruptValueRequest(
-                board_name=board.name, digital_interrupt_name="dne"
-            )
+            request = GetDigitalInterruptValueRequest(board_name=board.name, digital_interrupt_name="dne")
             with pytest.raises(GRPCError, match=r".*Status.NOT_FOUND.*"):
                 await client.GetDigitalInterruptValue(request)
 
-            request = GetDigitalInterruptValueRequest(
-                board_name=board.name, digital_interrupt_name="interrupt1"
-            )
+            request = GetDigitalInterruptValueRequest(board_name=board.name, digital_interrupt_name="interrupt1")
             response: GetDigitalInterruptValueResponse = await client.GetDigitalInterruptValue(request)
             assert response.value == 0
 
diff --git a/tests/test_movement_sensor.py b/tests/test_movement_sensor.py
@@ -34,7 +34,7 @@
     MovementSensorServiceStub,
 )
 from viam.resource.manager import ResourceManager
-from viam.utils import dict_to_struct, primitive_to_value, struct_to_dict
+from viam.utils import dict_to_struct, sensor_readings_value_to_native, struct_to_dict
 
 from . import loose_approx
 from .mocks.components import GEOMETRIES, MockMovementSensor
@@ -291,7 +291,7 @@ async def test_get_readings(self, movement_sensor: MockMovementSensor, service:
             request = GetReadingsRequest(name=movement_sensor.name, extra=dict_to_struct(EXTRA_PARAMS))
             assert movement_sensor.extra is None
             response: GetReadingsResponse = await client.GetReadings(request, timeout=8.90)
-            assert response.readings == {key: primitive_to_value(value) for (key, value) in READINGS.items()}
+            assert sensor_readings_value_to_native(response.readings) == READINGS
             assert movement_sensor.extra == EXTRA_PARAMS
             assert movement_sensor.timeout == loose_approx(8.90)
 
@@ -402,7 +402,7 @@ async def test_get_readings(self, movement_sensor: MockMovementSensor, service:
             client = MovementSensorClient(movement_sensor.name, channel)
             assert movement_sensor.extra is None
             value = await client.get_readings(extra=EXTRA_PARAMS, timeout=2.34)
-            assert value == {key: primitive_to_value(value) for (key, value) in READINGS.items()}
+            assert value == READINGS
             assert movement_sensor.extra == EXTRA_PARAMS
             assert movement_sensor.timeout == loose_approx(2.34)
 
diff --git a/tests/test_power_sensor.py b/tests/test_power_sensor.py
diff --git a/tests/test_sensor.py b/tests/test_sensor.py
