diff --git a/xrspatial/proximity.py b/xrspatial/proximity.py
index 78a9b7a0..5bea23d7 100644
--- a/xrspatial/proximity.py
+++ b/xrspatial/proximity.py
@@ -1517,7 +1517,7 @@ def allocation(
     target_values: list = [],
     max_distance: float = np.inf,
     distance_metric: str = "EUCLIDEAN",
-):
+) -> xr.DataArray:
     """
     Calculates, for all pixels in the input raster, the nearest source
     based on a set of target values and a distance metric.
@@ -1659,7 +1659,7 @@ def direction(
     target_values: list = [],
     max_distance: float = np.inf,
     distance_metric: str = "EUCLIDEAN",
-):
+) -> xr.DataArray:
     """
     Calculates, for all cells in the array, the downward slope direction
     Calculates, for all pixels in the input raster, the direction to
diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py
index cff04fa6..f3c41b8c 100644
--- a/xrspatial/tests/test_proximity.py
+++ b/xrspatial/tests/test_proximity.py
@@ -1,3 +1,4 @@
+import inspect
 from textwrap import dedent as textwrap_dedent
 from unittest.mock import patch
 
@@ -1091,6 +1092,13 @@ def test_no_scipy_dask_unbounded_memory_guard():
         prox_mod.cKDTree = original_ckdtree
 
 
+@pytest.mark.parametrize("func", [proximity, allocation, direction])
+def test_return_annotation_consistency(func):
+    # proximity, allocation, and direction share a signature and return
+    # type; their return annotations should match.
+    assert inspect.signature(func).return_annotation is xr.DataArray
+
+
 @pytest.mark.skipif(da is None, reason="dask is not installed")
 @pytest.mark.parametrize("func", [proximity, allocation, direction])
 def test_great_circle_dask_bounded_matches_numpy(func):