Python Union Types with Objects

623 views Asked by At

I am writing my own classes for int, str, bool, etc. that have generators on them. I am using this to fuzz the function based on the type annotations. This is going fine, with the exception of the | notation for union type. If I type something like:

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 gives me the following error:

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

This seems to be because Int[0,10] has the type of its class pybt.typing.Int, not type. However, using typing.Union works just fine.

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

Is there a way to get around this? Unfortunately, I can't think of a way to hold off instantiating Int or other types that are indexed in their __class_getitem__ under.

EDIT:

Here is the full class (for 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

There are 1 answers

10
BoppreH On BEST ANSWER

You already know about dunder methods and type, so apologies if this answer restates the obvious to you.

For Int() | List() to work (or anything that creates instances, like your List[]), there must be either a Int.__or__ or a List.__ror__, or both.

This works with the classes, as in Int | List, because type itself implements both __or__ and __ror__:

>>> type.__or__
<slot wrapper '__or__' of 'type' objects>

So try adding the following methods to your classes:

from typing import Union

class List:
    ...
    def __or__(self, other):
        return Union[self, other]
    def __ror__(self, other):
        return Union[self, other]

Seems to play well with other types, even its own class:

>>> List() | List
typing.Union[<__main__.List object at 0x0000024A7F7F93D0>, __main__.List]