ENH: override support for np.linspace and friends

Fixes gh-12379

1 files changed, 24 insertions, 5 deletions

diff --git a/numpy/core/function_base.py b/numpy/core/function_base.py
index b3dd313cf..0fc56e70e 100644
--- a/numpy/core/function_base.py
+++ b/numpy/core/function_base.py
@@ -1,17 +1,22 @@
 from __future__ import division, absolute_import, print_function
 
+import functools
 import warnings
 import operator
 
 from . import numeric as _nx
 from .numeric import (result_type, NaN, shares_memory, MAY_SHARE_BOUNDS,
-                      TooHardError,asanyarray)
+                      TooHardError, asanyarray)
 from numpy.core.multiarray import add_docstring
-from numpy.core.overrides import set_module
+from numpy.core import overrides
 
 __all__ = ['logspace', 'linspace', 'geomspace']
 
 
+array_function_dispatch = functools.partial(
+    overrides.array_function_dispatch, module='numpy')
+
+
 def _index_deprecate(i, stacklevel=2):
     try:
         i = operator.index(i)
@@ -24,7 +29,12 @@ def _index_deprecate(i, stacklevel=2):
     return i
 
 
-@set_module('numpy')
+def _linspace_dispatcher(
+        start, stop, num=None, endpoint=None, retstep=None, dtype=None):
+    return (start, stop)
+
+
+@array_function_dispatch(_linspace_dispatcher)
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
     """
     Return evenly spaced numbers over a specified interval.
@@ -156,7 +166,12 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
         return y.astype(dtype, copy=False)
 
 
-@set_module('numpy')
+def _logspace_dispatcher(
+        start, stop, num=None, endpoint=None, base=None, dtype=None):
+    return (start, stop)
+
+
+@array_function_dispatch(_logspace_dispatcher)
 def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     """
     Return numbers spaced evenly on a log scale.
@@ -241,7 +256,11 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
     return _nx.power(base, y).astype(dtype)
 
 
-@set_module('numpy')
+def _geomspace_dispatcher(start, stop, num=None, endpoint=None, dtype=None):
+    return (start, stop)
+
+
+@array_function_dispatch(_geomspace_dispatcher)
 def geomspace(start, stop, num=50, endpoint=True, dtype=None):
     """
     Return numbers spaced evenly on a log scale (a geometric progression).