From d1c5929b32ff572c9bd479c9d880447e591e3951 Mon Sep 17 00:00:00 2001
From: Alexander Druz <alexandersdruz@gmail.com>
Date: Mon, 5 Feb 2024 11:01:56 +0100
Subject: [PATCH] Add type annotations

---
 .../integration/dataset/test_bounding_box_column.py  |  4 ++--
 tests/integration/dataset/test_embedding_column.py   | 12 +++++-------
 tests/integration/h5/data.py                         |  3 ++-
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/tests/integration/dataset/test_bounding_box_column.py b/tests/integration/dataset/test_bounding_box_column.py
index 234d9cb1..acec565e 100644
--- a/tests/integration/dataset/test_bounding_box_column.py
+++ b/tests/integration/dataset/test_bounding_box_column.py
@@ -26,7 +26,7 @@ def test_default(empty_dataset: spotlight.Dataset) -> None:
     empty_dataset.append_bounding_box_column("column")
     assert empty_dataset.get_dtype("column") == dtypes.bounding_box_dtype
 
-    valid_values = (
+    valid_values: tuple = (
         [0, 1, 2, 3],
         [0.0, 1.0, np.nan, np.inf],
         range(4),
@@ -84,7 +84,7 @@ def test_optional(empty_dataset: spotlight.Dataset) -> None:
     empty_dataset.append_bounding_box_column("column", optional=True)
     assert empty_dataset.get_dtype("column") == dtypes.bounding_box_dtype
 
-    valid_values = (
+    valid_values: tuple = (
         [0, 1, 2, 3],
         [0.0, 1.0, np.nan, np.inf],
         range(4),
diff --git a/tests/integration/dataset/test_embedding_column.py b/tests/integration/dataset/test_embedding_column.py
index 159d295f..24ca470d 100644
--- a/tests/integration/dataset/test_embedding_column.py
+++ b/tests/integration/dataset/test_embedding_column.py
@@ -1,11 +1,9 @@
-from typing import Tuple
 import numpy as np
 import pytest
 
 from renumics import spotlight
 from renumics.spotlight import dtypes
 from renumics.spotlight.dataset.exceptions import InvalidDTypeError, InvalidShapeError
-from renumics.spotlight.dataset.typing import EmbeddingColumnInputType
 
 
 @pytest.mark.parametrize("length", [1, 2, 8])
@@ -16,7 +14,7 @@ def test_default(empty_dataset: spotlight.Dataset, length: int) -> None:
     empty_dataset.append_embedding_column("embedding")
     assert empty_dataset.get_dtype("embedding") == dtypes.embedding_dtype
 
-    valid_values: Tuple[EmbeddingColumnInputType] = (
+    valid_values: tuple = (
         [0] * length,
         range(length),
         tuple(range(length)),
@@ -69,7 +67,7 @@ def test_default_with_values(empty_dataset: spotlight.Dataset, length: int) -> N
     Test default embedding column creation with given values and afterwards
     filling row-by-row.
     """
-    valid_values = (
+    valid_values: tuple = (
         [0] * length,
         range(length),
         tuple(range(length)),
@@ -194,7 +192,7 @@ def test_length(empty_dataset: spotlight.Dataset, length: int) -> None:
     empty_dataset.append_embedding_column("embedding", length=length)
     assert empty_dataset.get_dtype("embedding") == dtypes.EmbeddingDType(length=length)
 
-    valid_values = (
+    valid_values: tuple = (
         [0] * length,
         range(length),
         tuple(range(length)),
@@ -256,7 +254,7 @@ def test_dtype(
     empty_dataset.append_embedding_column("embedding", dtype=np_dtype)
     assert empty_dataset.get_dtype("embedding") == dtypes.embedding_dtype
 
-    valid_values = (
+    valid_values: tuple = (
         [0] * length,
         range(length),
         tuple(range(length)),
@@ -284,7 +282,7 @@ def test_generic(empty_dataset: spotlight.Dataset, length: int) -> None:
     empty_dataset.append_column("embedding", dtype)
     assert empty_dataset.get_dtype("embedding") == dtype
 
-    valid_values = (
+    valid_values: tuple = (
         [0] * length,
         range(length),
         tuple(range(length)),
diff --git a/tests/integration/h5/data.py b/tests/integration/h5/data.py
index 6c5a3f18..df118e71 100644
--- a/tests/integration/h5/data.py
+++ b/tests/integration/h5/data.py
@@ -3,12 +3,13 @@
 """
 
 import datetime
+from typing import Dict, Tuple
 
 import numpy as np
 from renumics import spotlight
 
 
-COLUMNS = {
+COLUMNS: Dict[str, Tuple[str, list]] = {
     "bool": ("bool", [True, False]),
     "int": ("int", [0, 1]),
     "float": ("float", [0.0, np.nan]),