Add type overloads for get_value()

src/nipanel/_panel_value_accessor.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 from abc import ABC
+from typing import TypeVar, overload, cast
 
 import grpc
 from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
 from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
 
 from nipanel._panel_client import PanelClient
 
+_T = TypeVar("_T")
+
 
 class PanelValueAccessor(ABC):
     """This class allows you to access values for a panel's controls."""
@@ -41,7 +44,13 @@ def panel_id(self) -> str:
         """Read-only accessor for the panel ID."""
         return self._panel_id
 
-    def get_value(self, value_id: str, default_value: object = None) -> object:
+    @overload
+    def get_value(self, value_id: str) -> object: ...
+
+    @overload
+    def get_value(self, value_id: str, default_value: _T) -> _T: ...
+
+    def get_value(self, value_id: str, default_value: _T | None = None) -> _T:
         """Get the value for a control on the panel with an optional default value.
 
         Args:
@@ -52,7 +61,11 @@ def get_value(self, value_id: str, default_value: object = None) -> object:
             The value, or the default value if not set
         """
         try:
-            return self._panel_client.get_value(self._panel_id, value_id)
+            value = cast(_T, self._panel_client.get_value(self._panel_id, value_id))
+            if default_value is not None and not isinstance(value, type(default_value)):
+                raise TypeError("Value type does not match default value type.")
+            return value
+
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.NOT_FOUND and default_value is not None:
                 return default_value

tests/unit/test_streamlit_panel.py

@@ -1,5 +1,6 @@
 import grpc
 import pytest
+from typing_extensions import assert_type
 
 import tests.types as test_types
 from nipanel import StreamlitPanel, StreamlitPanelValueAccessor
@@ -158,6 +159,19 @@ def test___set_value___get_value_ignores_default(
     assert panel.get_value(value_id, "default") == string_value
 
 
+def test___set_value___get_value_with_wrong_type_raises_exception(
+    fake_panel_channel: grpc.Channel,
+) -> None:
+    panel = StreamlitPanel("my_panel", "path/to/script", grpc_channel=fake_panel_channel)
+
+    value_id = "test_id"
+    string_value = "test_value"
+    panel.set_value(value_id, string_value)
+
+    with pytest.raises(TypeError):
+        panel.get_value(value_id, 0)
+
+
 def test___get_value_returns_default_when_value_not_set(
     fake_panel_channel: grpc.Channel,
 ) -> None:
@@ -170,6 +184,60 @@ def test___get_value_returns_default_when_value_not_set(
     assert panel.get_value("missing_list", [1, 2, 3]) == [1, 2, 3]
 
 
+def test___set_value___get_value_returns_type_of_string_default_value(
+    fake_panel_channel: grpc.Channel,
+) -> None:
+    panel = StreamlitPanel("my_panel", "path/to/script", grpc_channel=fake_panel_channel)
+
+    value_id = "test_id"
+    string_value = "test_value"
+    panel.set_value(value_id, string_value)
+
+    value = panel.get_value(value_id, "")
+    assert_type(value, str)
+    assert value == string_value
+
+
+def test___set_value___get_value_returns_type_of_int_default_value(
+    fake_panel_channel: grpc.Channel,
+) -> None:
+    panel = StreamlitPanel("my_panel", "path/to/script", grpc_channel=fake_panel_channel)
+
+    value_id = "test_id"
+    int_value = 10
+    panel.set_value(value_id, int_value)
+
+    value = panel.get_value(value_id, 0)
+    assert_type(value, int)
+    assert value == int_value
+
+
+def test___set_value___get_value_returns_type_of_bool_default_value(
+    fake_panel_channel: grpc.Channel,
+) -> None:
+    panel = StreamlitPanel("my_panel", "path/to/script", grpc_channel=fake_panel_channel)
+
+    value_id = "test_id"
+    bool_value = True
+    panel.set_value(value_id, bool_value)
+
+    value = panel.get_value(value_id, False)
+    assert_type(value, bool)
+    assert value is bool_value
+
+
+def test___get_value_returns_correct_type_when_value_not_set(
+    fake_panel_channel: grpc.Channel,
+) -> None:
+    panel = StreamlitPanel("my_panel", "path/to/script", grpc_channel=fake_panel_channel)
+
+    assert_type(panel.get_value("missing_string", "default"), str)
+    assert_type(panel.get_value("missing_int", 123), int)
+    assert_type(panel.get_value("missing_float", 1.23), float)
+    assert_type(panel.get_value("missing_bool", True), bool)
+    assert_type(panel.get_value("missing_list", [1, 2, 3]), list[int])
+
+
 @pytest.mark.parametrize(
     "value_payload",
     [