add typehints

rdp/__init__.py

@@ -8,47 +8,43 @@
 :license: MIT, see LICENSE.txt for more details.
 
 """
-from math import sqrt
 from functools import partial
 import numpy as np
 import sys
+from typing import Callable, Literal, overload
 
 if sys.version_info[0] >= 3:
     xrange = range
 
 
-def pldist(point, start, end):
+DistanceFunction = Callable[[np.ndarray, np.ndarray, np.ndarray], float]
+
+def pldist(point: np.ndarray, start: np.ndarray, end: np.ndarray) -> float:
     """
     Calculates the distance from ``point`` to the line given
     by the points ``start`` and ``end``.
 
     :param point: a point
-    :type point: numpy array
     :param start: a point of the line
-    :type start: numpy array
     :param end: another point of the line
-    :type end: numpy array
     """
     if np.all(np.equal(start, end)):
-        return np.linalg.norm(point - start)
+        return float(np.linalg.norm(point - start))
 
     return np.divide(
             np.abs(np.linalg.norm(np.cross(end - start, start - point))),
             np.linalg.norm(end - start))
 
 
-def rdp_rec(M, epsilon, dist=pldist):
+def rdp_rec(M: np.ndarray, epsilon: float, dist: DistanceFunction = pldist) -> np.ndarray:
     """
     Simplifies a given array of points.
 
     Recursive version.
 
-    :param M: an array
-    :type M: numpy array
+    :param M: an array of points, see :func:`rdp.rdp`
     :param epsilon: epsilon in the rdp algorithm
-    :type epsilon: float
-    :param dist: distance function
-    :type dist: function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
+    :param dist: distance function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
     """
     dmax = 0.0
     index = -1
@@ -69,7 +65,7 @@ def rdp_rec(M, epsilon, dist=pldist):
         return np.vstack((M[0], M[-1]))
 
 
-def _rdp_iter(M, start_index, last_index, epsilon, dist=pldist):
+def _rdp_iter(M: np.ndarray, start_index: int, last_index: int, epsilon: float, dist: DistanceFunction = pldist):
     stk = []
     stk.append([start_index, last_index])
     global_start_index = start_index
@@ -98,20 +94,16 @@ def _rdp_iter(M, start_index, last_index, epsilon, dist=pldist):
     return indices
 
 
-def rdp_iter(M, epsilon, dist=pldist, return_mask=False):
+def rdp_iter(M: np.ndarray, epsilon: float, dist: DistanceFunction = pldist, return_mask=False) -> np.ndarray:
     """
     Simplifies a given array of points.
 
     Iterative version.
 
-    :param M: an array
-    :type M: numpy array
+    :param M: an array of points, see :func:`rdp.rdp`
     :param epsilon: epsilon in the rdp algorithm
-    :type epsilon: float
-    :param dist: distance function
-    :type dist: function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
+    :param dist: distance function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
     :param return_mask: return the mask of points to keep instead
-    :type return_mask: bool
     """
     mask = _rdp_iter(M, 0, len(M) - 1, epsilon, dist)
 
@@ -121,7 +113,7 @@ def rdp_iter(M, epsilon, dist=pldist, return_mask=False):
     return M[mask]
 
 
-def rdp(M, epsilon=0, dist=pldist, algo="iter", return_mask=False):
+def rdp(M: np.ndarray, epsilon: float = 0.0, dist: DistanceFunction = pldist, algo: Literal["iter", "rec"] = "iter", return_mask=False) -> np.ndarray:
     """
     Simplifies a given array of points using the Ramer-Douglas-Peucker
     algorithm.
@@ -157,26 +149,24 @@ def rdp(M, epsilon=0, dist=pldist, algo="iter", return_mask=False):
     array([[1, 1],
            [4, 4]])
 
-    :param M: a series of points
-    :type M: numpy array with shape ``(n,d)`` where ``n`` is the number of points and ``d`` their dimension
+    :param M: a series of points as numpy array with shape ``(n,d)`` where ``n`` is the number of points and ``d`` their dimension
     :param epsilon: epsilon in the rdp algorithm
-    :type epsilon: float
-    :param dist: distance function
-    :type dist: function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
+    :param dist: distance function with signature ``f(point, start, end)`` -- see :func:`rdp.pldist`
     :param algo: either ``iter`` for an iterative algorithm or ``rec`` for a recursive algorithm
-    :type algo: string
     :param return_mask: return mask instead of simplified array
-    :type return_mask: bool
     """
 
+    algo_function: Callable[[np.ndarray, float, DistanceFunction], np.ndarray]
+
     if algo == "iter":
-        algo = partial(rdp_iter, return_mask=return_mask)
+        algo_function = partial(rdp_iter, return_mask=return_mask)
     elif algo == "rec":
         if return_mask:
             raise NotImplementedError("return_mask=True not supported with algo=\"rec\"")
-        algo = rdp_rec
+        algo_function = rdp_rec
 
     if "numpy" in str(type(M)):
-        return algo(M, epsilon, dist)
+        return algo_function(M, epsilon, dist)
 
-    return algo(np.array(M), epsilon, dist).tolist()
+    # This is a 'secret feature': Use the rdp function with lists of data.
+    return algo_function(np.array(M), epsilon, dist).tolist()