[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