使用 Python 类型提示对数字强制使用单位

Question

有没有办法使用Python类型提示作为单位？类型提示文档显示了一些示例，表明可以使用NewType，但这些示例也表明，相同“新类型”的两个值相加不会给出“新类型”的结果，而是给出基本类型的结果。有没有一种方法可以丰富类型定义，以便您可以指定像单位一样工作的类型提示（不是它们转换的范围，而是当您获得不同单位时收到类型警告）？可以让我执行此操作或类似操作的东西：

Seconds = UnitType('Seconds', float)
Meters = UnitType('Meters', float)

time1 = Seconds(5)+ Seconds(8) # gives a value of type `Seconds`
bad_units1 = Seconds(1) + Meters(5) # gives a type hint error, but probably works at runtime
time2 = Seconds(1)*5 # equivalent to `Seconds(1*5)` 
# Multiplying units together of course get tricky, so I'm not concerned about that now.

我知道单元的运行时库存在，但我好奇的是 python 中的类型提示是否能够处理某些功能。

Answer 1

您可以通过创建类型存根文件来完成此操作，该文件定义__add__/__radd__方法（定义+运算符）和__sub__/__rsub__方法（定义运算-符）的可接受类型。当然，对于其他运算符还有许多类似的方法，但为了简洁起见，本示例仅使用这些方法。

units.py

这里我们将单位定义为 的简单别名int。这最大限度地减少了运行时成本，因为我们实际上并没有创建一个新类。

Seconds = int
Meters = int

units.pyi

这是一个类型存根文件。它告诉类型检查器 中定义的所有内容的类型units.py，而不是在代码中定义类型。类型检查器假定这是事实的来源，并且当它与 中实际定义的内容不同时不会引发错误units.py。

from typing import Generic, TypeVar

T = TypeVar("T")

class Unit(int, Generic[T]):
    def __add__(self, other: T) -> T: ...
    def __radd__(self, other: T) -> T: ...
    def __sub__(self, other: T) -> T: ...
    def __rsub__(self, other: T) -> T: ...
    def __mul__(self, other: int) -> T: ...
    def __rmul__(self, other: int) -> T: ...

class Seconds(Unit["Seconds"]): ...

class Meters(Unit["Meters"]): ...

这里我们定义Unit为继承自的泛型类型int，其中加/减采用并返回类型参数 的值T。Seconds然后将和Meters定义为 , 的子类Unit，分别T等于Seconds和Meters。

这样，类型检查器就知道 和 的加/减运算Seconds会返回 类型的其他值Seconds，对于 也类似Meters。

另外，我们将__mul__and __rmul__on定义Unit为采用 type 参数int并返回T- 所以Seconds(1) * 5应该具有 type Seconds。

main.py

这是你的代码。

from units import Seconds, Meters

time1 = Seconds(5) + Seconds(8)
# time1 has type Seconds, yay!

bad_units1 = Seconds(1) + Meters(5)
# I get a type checking error:
# Operator "+" not supported for types "Meters" and "Seconds"
# Yay!

time2 = Seconds(1) * 5
# time2 has type Seconds, yay!

meter_seconds = Seconds(1) * Meters(5)
# This is valid because `Meters` is a subclass of `int` (as far
# as the type checker is concerned). meter_seconds ends up being
# type Seconds though - as you say, multiplying gets tricky.

当然，所有这些都只是类型检查。您可以在运行时执行您喜欢的操作，并且该pyi文件甚至不会被加载。