Astropy: Как я могу проверить, что аргумент имеет правильную единицу

Question

В Python astropy, как я могу проверить, что аргумент функции не только имеет правильную единицу, но вообще имеет единицу?

Я знаком с is_equivalent(), таким образом, чтобы проверить, что M имеет единицы массы, я могу сказать

assert M.unit.is_equivalent(u.g)

, которая возвращает True, если, скажем, M = 1e12 * u.Msun. Но если M вообще не имеет каких-либо единиц - например, если я случайно передам M = 1e12 - я получу

AttributeError: 'Quantity' object has no '_normalize_equivalencies' member

Один из способов сделать это - сначала умножить M на, например, (u.m/u.m), что равно Unit(dimensionless) и не изменит M, если у него есть единица, в то время как оно конвертирует число типа 1e12 в <Quantity 1.e+12>, которое затем можно проверить без указания ошибка.

Хотя это не похоже на правильный путь. Есть ли лучший способ?

Answer 1

Как уже упоминалось в комментариях, если ваша цель состоит в том, чтобы конкретно проверить единицы аргументов функции, вы можете использовать служебный декоратор quantity_input, который специально предназначен для этой цели (и выдает либо ValueError или TypeError или UnitsError, в зависимости от ситуации, которые, как правило, являются более подходящими исключениями для недопустимых аргументов функции, чем голые AssertionError).

Мне особенно нравится синтаксис Python 3-speci c для этого, который использует аннотации функций. Вот несколько примеров:

>>> from astropy import units as u
>>> @u.quantity_input
... def my_function(M: u.g):
...     print(M)

Правильные единицы:

>>> my_function(1e12*u.Msun)                                                
1000000000000.0 solMass
>>> my_function(1e12*u.g)                                                   
1000000000000.0 g

Примеры с неверными аргументами:

>>> my_function(1e12)
Traceback (most recent call last)
...
TypeError: Argument 'M' to function 'my_function' has no 'unit' attribute. You may want to pass in an astropy Quantity instead.

>>> my_function(1e12*u.dimensionless_unscaled)
Traceback (most recent call last)
...
UnitsError: Argument 'M' to function 'my_function' must be in units convertible to 'g'.

Если у вас много функций, которые принимают M в качестве аргумента, который должен иметь массовые единицы, вы также можете сделать сокращение для него, например:

>>> M_mass = u.quantity_input(M=u.g)
>>> @M_mass
... def my_function(M):
...     print(M)

Тот же подход можно использовать для функций нескольких аргументов в разных единицах.