Conditionally specify type for TypeVar with typing.overload

[tomasgareau]

I'm wondering how to type the following situation. From what I gather this case might be type-checker dependent so I thought I'd ask this here but please let me know if this isn't the right place to be asking.

Setup

Let's say I had a function that read bytes from some stream, deserialized them into an object, then put that object in a queue. I'd type that like so:

from somewhere import Queue, read_bytes_from_somewhere
from typing import Callable, NoReturn, TypeVar


T = TypeVar("T")

def my_func(q: Queue[T], deserialize: Callable[[bytes], T]) -> NoReturn:
    while True:
        value = read_bytes_from_somewhere()
        obj = deserialize(value)
        q.put(obj)

to hint that the return type of the deserialize callable should match the generic type of the Queue. Works mint!

Now, I'd like to make deserialize an optional value: if omitted, the function will not perform any deserialization and will simply stick bytes into your queue. In this case, I'd like to hint that it will be a queue of bytes, not T. I'm not quite sure how to accomplish this -- my first crack is below:

from typing import Callable, Literal, NoReturn, TypeVar, overload

# if deserialize is None, q must be Queue[bytes]...
@overload
def my_func(q: Queue[bytes], deserialize: Literal[None]) -> NoReturn: ...

# if deserialize is not None, it must return T (to match Queue[T])
@overload
def my_func(q: Queue[T], deserialize: Callable[[bytes], T]) -> NoReturn: ...

Attempt 1: don't type the implementation

def my_func(q, deserialize)
    while True:
        value = read_bytes_from_somewhere()
        if deserialize:
            q.put(deserialize(value)
        else:
            q.put(value)

# doesn't quite work: when using my_func, I get the warning:
#     pyright: Type of "my_func" is partially unknown. Type of "my_func" is
#         (overload 1 signature...)
#         (overload 2 signature...)
#         (q: Unknown, deserialize: Unknown) -> Unknown

Attempt 2: type the implementation with Union?

# gives warnings in implementation (see inline)
def my_func(
        q: Union[Queue[T], Queue[bytes]],
        deserialize: Optional[Callable[[bytes], T]]
    ) -> NoReturn:
    while True:
        value = read_bytes_from_somewhere()
        if deserialize:
            q.put(deserialize(value)  # T@my_func cannot be assigned to "item" of type "bytes"
        else:
            q.put(value)  # bytes cannot be assigned to "item" of type "T@read_to"

Attempt 3: throw more typing.Any at the problem

def my_func(q: Any, deserialize: Any) -> NoReturn:  # same implementation as before

This seems to work the best: I get no warnings in the implementation or when using my_func, and get told off if I'm mixing the wrong types:

q: Queue[MyClass] = Queue()

# good -- can't pass None if queue doesn't take bytes
my_func(q, None)  # no overloads for "my_func" match the provided arguments

# good -- can't pass a callable returning the wrong type
my_func(q, lambda b: YourClass())  # no overloads for "my_func" match the provided arguments

# good -- even though the implementation used Any I don't actually accept Any
my_func(1, "1")  # no overloads for "my_func"... good, can't pass Any

If I hover over the function for help info, however, I see (q: Any, deserialize: Any) -> NoReturn as one of the signatures for that function, which makes me think there might be a better way to type what I'm trying.

[erictraut]

I don't see a perfect solution, but here's the best I can offer:

def my_func(q: Queue[Any], deserialize: Optional[Callable[[bytes], Any]]) -> NoReturn:

[tomasgareau]

That'll work -- just wanted to double-check I wasn't missing something basic. Thanks!

[hmc-cs-mdrissi]

Would having it be optional with a default callable that is the identity function let it work fully typed? Something like,

def id(x: T) -> T:
  return x

def my_func(q: Queue[T], deserialize: Callable[[bytes], T] = id) -> NoReturn:
   ...

[tomasgareau]

Could be a good approach in other cases! In our case though, we wanted to constrain T == bytes in the absence of a custom deserialize callable (the source we'll be deserializing will always start from bytes).

With the identity function as written above my understanding is that I could do something like:

q: Queue[SomeRandomType]
my_func(q)

without any errors, but I'd end up with bytes in my queue.