[Numpy-discussion] What is the pythonic way to write a function that handles arrays and scalars?

Robert Kern robert.kern at gmail.com
Tue Dec 12 14:57:01 EST 2017


On Wed, Dec 13, 2017 at 4:50 AM, Joe <solarjoe at posteo.org> wrote:
>
> Hi,
>
> the best example I found was this one:
>
> https://stackoverflow.com/a/29319864/7919597
>
> def func_for_scalars_or_vectors(x):
>     x = np.asarray(x)
>     scalar_input = False
>     if x.ndim == 0:
>         x = x[None]  # Makes x 1D
>         scalar_input = True
>
>     # The magic happens here
>
>     if scalar_input:
>         return np.squeeze(ret)
>     return ret
>
> Is this as good as it gets or do you have other suggestions?

In general structure, yeah. I'd probably use `np.isscalar(x)` for the test
and `x = np.atleast_1d(x)` for the coercion for readability, but otherwise,
that's it.

--
Robert Kern
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20171213/66373c41/attachment.html>


More information about the NumPy-Discussion mailing list