From d594f4556cdc0c964c76daf5d9c891c32c5823bf Mon Sep 17 00:00:00 2001
From: Steve Kim <54872857+SteveKimSR@users.noreply.github.com>
Date: Fri, 22 Jan 2021 13:40:21 +0900
Subject: [PATCH] [mypy] Add/fix type annotations for similarity search in
 machine learning (#4088)

* [mypy] Add/fix type annotations for similarity search in machine learning

* fix annotation

* fix annotation (Union)

* isort
---
 machine_learning/similarity_search.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/machine_learning/similarity_search.py b/machine_learning/similarity_search.py
index 6bfb12ed8..af845c910 100644
--- a/machine_learning/similarity_search.py
+++ b/machine_learning/similarity_search.py
@@ -8,6 +8,7 @@ returns a list containing two data for each vector:
     2. distance between the vector and the nearest vector (float)
 """
 import math
+from typing import List, Union
 
 import numpy as np
 
@@ -30,7 +31,9 @@ def euclidean(input_a: np.ndarray, input_b: np.ndarray) -> float:
     return math.sqrt(sum(pow(a - b, 2) for a, b in zip(input_a, input_b)))
 
 
-def similarity_search(dataset: np.ndarray, value_array: np.ndarray) -> list:
+def similarity_search(
+    dataset: np.ndarray, value_array: np.ndarray
+) -> List[List[Union[List[float], float]]]:
     """
     :param dataset: Set containing the vectors. Should be ndarray.
     :param value_array: vector/vectors we want to know the nearest vector from dataset.