“至少一个参数属于 X 类型”的类型注释

Question

我试图使用重载来使可变参数函数的返回类型以某种方式依赖于其参数的类型。具体来说，当且仅当其任何参数的类型为 X 时，我希望返回类型为X。

考虑以下最小示例：

from typing import overload

class Safe:
    pass

class Dangerous:
    pass

@overload
def combine(*args: Safe) -> Safe: ...

@overload
def combine(*args: Safe | Dangerous) -> Safe | Dangerous: ...

def combine(*args: Safe | Dangerous) -> Safe | Dangerous:
    if all(isinstance(arg, Safe) for arg in args):
        return Safe()
    else:
        return Dangerous()

reveal_type(combine())
reveal_type(combine(Safe()))
reveal_type(combine(Dangerous()))
reveal_type(combine(Safe(), Safe()))
reveal_type(combine(Safe(), Dangerous()))

这输出

example.py:21: note: Revealed type is "example.Safe"
example.py:22: note: Revealed type is "example.Safe"
example.py:23: note: Revealed type is "Union[example.Safe, example.Dangerous]"
example.py:24: note: Revealed type is "example.Safe"
example.py:25: note: Revealed type is "Union[example.Safe, example.Dangerous]"
Success: no issues found in 1 source file

我想进行设置，以便推断类型 和combine(Dangerous())，combine(Safe(), Dangerous())例如，而Dangerous不是Safe | Dangerous。将第二个重载的返回类型更改为仅Dangerous产生错误：

example.py:10: error: Overloaded function signatures 1 and 2 overlap with incompatible return types  [misc]
example.py:21: note: Revealed type is "example.Safe"
example.py:22: note: Revealed type is "example.Safe"
example.py:23: note: Revealed type is "example.Dangerous"
example.py:24: note: Revealed type is "example.Safe"
example.py:25: note: Revealed type is "example.Dangerous"
Found 1 error in 1 file (checked 1 source file)

因此，我似乎需要一种方法来注释第二个重载，以明确声明它的至少一个参数是Dangerous。有没有办法做到这一点？

我发现参数序列所需的类型是Sequence[Safe | Dangerous] - Sequence[Safe]，但我认为尚不支持类型减法。

Answer 1

使用 PyRight，您的解决方案确实有效。mypy 也可以通过简单的方式被迫接受它# type: ignore：

@overload
def combine(*args: Safe) -> Safe: ...  # type: ignore

@overload
def combine(*args: Safe | Dangerous) -> Dangerous: ...

def combine(*args: Safe | Dangerous) -> Safe | Dangerous:
    if all(isinstance(arg, Safe) for arg in args):
        return Safe()
    else:
        return Dangerous()


reveal_type(combine())  # Safe
reveal_type(combine(Safe()))  # Safe
reveal_type(combine(Dangerous()))  # Dangerous
reveal_type(combine(Safe(), Safe()))  # Safe
reveal_type(combine(Safe(), Dangerous()))  # Dangerous

但请注意，如果输入类型不完全正确，这可能会导致静态类型和动态类型之间的差异：

arg: Safe | Dangerous = random.choice([Safe(), Dangerous()])
result = combine(arg)
reveal_type(result)  # Dangerous
print(type(result))  # Can be Safe or Dangerous