Implement {get,put}_vector_list

tests/features/test_sample.py

@@ -4,7 +4,7 @@
 from tempfile import NamedTemporaryFile
 from typing import Any
 
-from miplearn.features.sample import MemorySample, Sample, Hdf5Sample
+from miplearn.features.sample import MemorySample, Sample, Hdf5Sample, _pad, _crop
 
 
 def test_memory_sample() -> None:
@@ -29,16 +29,11 @@ def _test_sample(sample: Sample) -> None:
     _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],
-    #     ],
-    # )
+    # VectorList
+    _assert_roundtrip_vector_list(sample, [["A"], ["BB", "CCC"], None])
+    _assert_roundtrip_vector_list(sample, [[True], [False, False], None])
+    _assert_roundtrip_vector_list(sample, [[1], None, [2, 2], [3, 3, 3]])
+    _assert_roundtrip_vector_list(sample, [[1.0], None, [2.0, 2.0], [3.0, 3.0, 3.0]])
 
 
 def _assert_roundtrip_scalar(sample: Sample, expected: Any) -> None:
@@ -57,8 +52,76 @@ def _assert_roundtrip_vector(sample: Sample, expected: Any) -> None:
     _assert_same_type(actual[0], expected[0])
 
 
+def _assert_roundtrip_vector_list(sample: Sample, expected: Any) -> None:
+    sample.put_vector_list("key", expected)
+    actual = sample.get_vector_list("key")
+    assert actual == expected
+    assert actual is not None
+    _assert_same_type(actual[0][0], expected[0][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"
+    assert isinstance(
+        actual, expected.__class__
+    ), f"Expected {expected.__class__}, found {actual.__class__} instead"
+
+
+def test_pad_int() -> None:
+    _assert_roundtrip_pad(
+        original=[[1], [2, 2, 2], [], [3, 3], [4, 4, 4, 4], None],
+        expected_padded=[
+            [1, 0, 0, 0],
+            [2, 2, 2, 0],
+            [0, 0, 0, 0],
+            [3, 3, 0, 0],
+            [4, 4, 4, 4],
+            [0, 0, 0, 0],
+        ],
+        expected_lens=[1, 3, 0, 2, 4, -1],
+        dtype=int,
     )
+
+
+def test_pad_float() -> None:
+    _assert_roundtrip_pad(
+        original=[[1.0], [2.0, 2.0, 2.0], [3.0, 3.0], [4.0, 4.0, 4.0, 4.0], None],
+        expected_padded=[
+            [1.0, 0.0, 0.0, 0.0],
+            [2.0, 2.0, 2.0, 0.0],
+            [3.0, 3.0, 0.0, 0.0],
+            [4.0, 4.0, 4.0, 4.0],
+            [0.0, 0.0, 0.0, 0.0],
+        ],
+        expected_lens=[1, 3, 2, 4, -1],
+        dtype=float,
+    )
+
+
+def test_pad_str() -> None:
+    _assert_roundtrip_pad(
+        original=[["A"], ["B", "B", "B"], ["C", "C"]],
+        expected_padded=[["A", "", ""], ["B", "B", "B"], ["C", "C", ""]],
+        expected_lens=[1, 3, 2],
+        dtype=str,
+    )
+
+
+def _assert_roundtrip_pad(
+    original: Any,
+    expected_padded: Any,
+    expected_lens: Any,
+    dtype: Any,
+) -> None:
+    actual_padded, actual_lens = _pad(original)
+    assert actual_padded == expected_padded
+    assert actual_lens == expected_lens
+    for v in actual_padded:
+        for vi in v:  # type: ignore
+            assert isinstance(vi, dtype)
+    cropped = _crop(actual_padded, actual_lens)
+    assert cropped == original
+    for v in cropped:
+        if v is None:
+            continue
+        for vi in v:  # type: ignore
+            assert isinstance(vi, dtype)