numpy против python: исправление функции раунда

Question

У меня есть файл, который определяет

def foo(x):
   return round(x)

Здесь round - это функция python built-in.

Теперь я хотел бы вызвать эту функцию с помощью numpy массив. Numpy также имеет функцию round. К сожалению, есть проблемы (например, # 11557 )

import numpy as np

foo(7.6)
foo(np.array([7.6]))       # works with python2 + np_1.14.0, not with 3.6.9 + np_1.14.5
foo(np.array([7.6, 8.9]))  # TypeError: type numpy.ndarray doesn't define __round__ method

Есть ли возможность заменить функцию round на np.round внутри функции foo? Какой-нибудь патч, например round = np.round до импорта или foo.round = np.round после импорта?

Редактировать: Я ищу решение без изменения файла.

Answer 1

Подкласс - это обходной путь (как прокомментировал @Anton Pomieshchenko)

class Array(np.ndarray):
    def __new__(cls, array, **kwargs):
        return np.asarray(array, **kwargs).view(cls)
    def __round__(self):
        return np.round(self.data)

foo(Array([7.6, 8.9]))

Работает для python 3, в то время как в python 2 я получаю TypeError: only size-1 arrays can be converted to Python scalars.

Я безуспешно пытался также np.__round__ = np.round (TypeError: только массивы размера 1 могут быть преобразованы в Python скаляры) или np.ndarray.__round__ = np.round (TypeError: невозможно установить атрибуты встроенного / расширенного типа 'numpy. ndarray ').

Answer 2

Да, конечно, вы можете сделать это. Просто проверьте класс x, используя isinstance

import numpy as np

def foo(x):
    if isinstance(x, np.ndarray):
        return x.round()
    return round(x)


print(foo(7.6))
print(foo(np.array([7.6])))
print(foo(np.array([7.6, 8.9])))