Split Sample.{get,put} into {get,put}_{scalar,vector}

tests/features/test_sample.py

@@ -7,34 +7,6 @@ from typing import Any
 from miplearn.features.sample import MemorySample, Sample, Hdf5Sample
 
 
-def _test_sample(sample: Sample) -> None:
-    _assert_roundtrip(sample, "A")
-    _assert_roundtrip(sample, True)
-    _assert_roundtrip(sample, 1)
-    _assert_roundtrip(sample, 1.0)
-    _assert_roundtrip(sample, ["A", "BB", "CCC", "こんにちは"])
-    _assert_roundtrip(sample, [True, True, False])
-    _assert_roundtrip(sample, [1, 2, 3])
-    _assert_roundtrip(sample, [1.0, 2.0, 3.0])
-
-
-def _assert_roundtrip(sample: Sample, expected: Any) -> None:
-    sample.put("key", expected)
-    actual = sample.get("key")
-    assert actual == expected
-    assert actual is not None
-    if isinstance(actual, list):
-        assert isinstance(actual[0], expected[0].__class__), (
-            f"Expected class {expected[0].__class__}, "
-            f"found {actual[0].__class__} instead"
-        )
-    else:
-        assert isinstance(actual, expected.__class__), (
-            f"Expected class {expected.__class__}, "
-            f"found class {actual.__class__} instead"
-        )
-
-
 def test_memory_sample() -> None:
     _test_sample(MemorySample())
 
@@ -42,3 +14,51 @@ def test_memory_sample() -> None:
 def test_hdf5_sample() -> None:
     file = NamedTemporaryFile()
     _test_sample(Hdf5Sample(file.name))
+
+
+def _test_sample(sample: Sample) -> None:
+    # Scalar
+    _assert_roundtrip_scalar(sample, "A")
+    _assert_roundtrip_scalar(sample, True)
+    _assert_roundtrip_scalar(sample, 1)
+    _assert_roundtrip_scalar(sample, 1.0)
+
+    # Vector
+    _assert_roundtrip_vector(sample, ["A", "BB", "CCC", "こんにちは"])
+    _assert_roundtrip_vector(sample, [True, True, False])
+    _assert_roundtrip_vector(sample, [1, 2, 3])
+    _assert_roundtrip_vector(sample, [1.0, 2.0, 3.0])
+
+    # List[Optional[List[Primitive]]]
+    # _assert_roundtrip(
+    #     sample,
+    #     [
+    #         [1],
+    #         None,
+    #         [2, 2],
+    #         [3, 3, 3],
+    #     ],
+    # )
+
+
+def _assert_roundtrip_scalar(sample: Sample, expected: Any) -> None:
+    sample.put_scalar("key", expected)
+    actual = sample.get_scalar("key")
+    assert actual == expected
+    assert actual is not None
+    _assert_same_type(actual, expected)
+
+
+def _assert_roundtrip_vector(sample: Sample, expected: Any) -> None:
+    sample.put_vector("key", expected)
+    actual = sample.get_vector("key")
+    assert actual == expected
+    assert actual is not None
+    _assert_same_type(actual[0], expected[0])
+
+
+def _assert_same_type(actual: Any, expected: Any) -> None:
+    assert isinstance(actual, expected.__class__), (
+        f"Expected class {expected.__class__}, "
+        f"found class {actual.__class__} instead"
+    )