自定义类型无法使用|联合类型语法的解决方案咨询

问题 #

我正在编写带生成器的Int、Str、Bool等自定义类，用来基于类型注解做函数模糊测试。目前一切正常，但用|表示联合类型时出错：

def test_handles_none_with_arged_types(
    x: Int[0, 10] | List[Int] | Str | Dict[NoneType, List[NoneType]]
):
    assert type(x) in [int, list, str, dict, list]
    if type(x) == int:
        assert x >= 0 and x <= 10
    if type(x) == list:
        assert all([el is None for el in x])
    if type(x) == dict:
        for k, v in x.items():
            assert k is None
            assert type(v) == list
            for el in v:
                assert el is None 

Python抛出错误：

TypeError: unsupported operand type(s) for |: 'Int' and 'List'

原因是Int[0,10]的类型是pybt.typing.Int而非type，但用typing.Union就能正常运行：

def test_handles_none_with_arged_types(
    x: Union[Int[0, 10], List[Int], Str, Dict[NoneType, List[NoneType]]]
):
    ... 

找不到延迟实例化Int等在__class_getitem__中被索引类型的方法，求解决办法。

附List类完整代码：

class List:
    def __init__(self, sub_type=_DEFAULT_SUB_TYPE, max_len=_DEFAULT_MAX_LEN):
        self.max_len: int = _DEFAULT_MAX_LEN
        self.sub_type = sub_type
        if max_len is not None:
            self.max_len = max_len

    def __str__(self):
        return "pybt.types.List"

    def __class_getitem__(cls, parameters):
        sub_type = None
        max_len = None
        if type(parameters) != tuple:
            parameters = (parameters,)

        if len(parameters) > 2:
            raise TypeError("Expected 2 arguments: List[sub_type, max_length]")
        if len(parameters):
            sub_type = parameters[0]
        if len(parameters) > 1:
            max_len = parameters[1]

        if max_len and max_len <= 0:
            raise TypeError(f"Max Length of {cls.max_len} is less than or equal to 0")

        return cls(sub_type, max_len)

解决方案 #

方法1：给自定义类实现__or__和__ror__方法 #

让Int、List等自定义类型支持|运算符，返回Union对象，直接兼容PEP 604的联合类型语法。以List类为例修改：

from typing import Union

class List:
    # 保留原有__init__、__str__、__class_getitem__方法
    
    def __or__(self, other):
        return Union[self, other]
    
    def __ror__(self, other):
        return Union[other, self]

这样Int[0,10] | List[Int]会直接返回Union[Int[0,10], List[Int]]，和使用Union的效果完全一致。

方法2：让自定义类继承自type #

如果你的自定义类型本质是模拟类型语义，可以让它们继承type（元类），这样天然支持|运算符（PEP 604对type实例原生支持该语法）。调整List类实现：

class List(type):
    def __new__(cls, sub_type=_DEFAULT_SUB_TYPE, max_len=_DEFAULT_MAX_LEN):
        instance = super().__new__(cls, 'List', (), {})
        instance.max_len = max_len if max_len is not None else _DEFAULT_MAX_LEN
        instance.sub_type = sub_type
        return instance
    
    def __class_getitem__(cls, parameters):
        sub_type = None
        max_len = None
        if not isinstance(parameters, tuple):
            parameters = (parameters,)

        if len(parameters) > 2:
            raise TypeError("Expected 2 arguments: List[sub_type, max_length]")
        if len(parameters):
            sub_type = parameters[0]
        if len(parameters) > 1:
            max_len = parameters[1]

        if max_len and max_len <= 0:
            raise TypeError(f"Max Length of {max_len} is less than or equal to 0")

        return cls(sub_type, max_len)

方法3：用包装器实现延迟实例化 #

如果需要严格延迟Int、List等类的实例化，可创建包装器类，在__class_getitem__中返回包装器而非直接实例化，直到需要使用时才完成实例化：

from typing import Union

class TypeWrapper:
    def __init__(self, cls, *args, **kwargs):
        self.cls = cls
        self.args = args
        self.kwargs = kwargs
        self._instance = None
    
    def get_instance(self):
        if self._instance is None:
            self._instance = self.cls(*self.args, **self.kwargs)
        return self._instance
    
    def __or__(self, other):
        left = self.get_instance()
        right = other.get_instance() if isinstance(other, TypeWrapper) else other
        return Union[left, right]
    
    def __ror__(self, other):
        left = other.get_instance() if isinstance(other, TypeWrapper) else other
        right = self.get_instance()
        return Union[left, right]

class List:
    # 保留原有__init__、__str__方法
    
    def __class_getitem__(cls, parameters):
        sub_type = None
        max_len = None
        if not isinstance(parameters, tuple):
            parameters = (parameters,)

        if len(parameters) > 2:
            raise TypeError("Expected 2 arguments: List[sub_type, max_length]")
        if len(parameters):
            sub_type = parameters[0]
        if len(parameters) > 1:
            max_len = parameters[1]

        if max_len and max_len <= 0:
            raise TypeError(f"Max Length of {max_len} is less than or equal to 0")

        return TypeWrapper(cls, sub_type, max_len)

内容的提问来源于stack exchange，提问作者dvr