CODEWITHSUNDEEP
pythonnumpy
Aleksejs Fomins
Aleksejs Fomins

Reputation: 908

shorter way to change data type of scalars and numpy arrays

Can this function be written shorter? I end up doing this a lot in my code

def smartInt(x):
  if (type(x) == np.ndarray):
    return x.astype(int)
  else:
    return int(x)

Upvotes: 0

Views: 181

Answers (2)

jpp
jpp

Reputation: 164803

This isn't shorter in terms of number of lines, but try / except may be more efficient than explicit type checking:

def smartInt(x):
    try:
        return int(x)
    except TypeError:
        return x.astype(int)

If an array is the more likely input:

def smartInt(x):
    try:
        return x.astype(int)
    except AttributeError:
        return int(x)

Upvotes: 1

Paul Panzer
Paul Panzer

Reputation: 53089

If you are ok with numpy scalars (e.g. np.int_ instead of int, then you can do

np.asanyarray(x, int)[()]

or even

np.int_(x)

Upvotes: 2

Related Questions