How to type annotate different return type depending on if a parameter is given like in case of next() function?

Question

next() function has a special property, that next(iterable) returns the element or raises an exception, but next(iterable, None) returns an element or None.

How to type annotate it? Consider the following, I am using pyright to check:

from typing import TypeVar, Union

R = TypeVar('R')

class SuperNone:
  pass
    
class MyDict:
   data = { "a": 1 }

   def get(
      self, key: str, default: R = SuperNone
   ) -> Union[R, Elem] if R == SuperNone else Elem:
      try:
         return self.data[key]
      except KeyError:
         if isinstance(default, SuperNone):
             raise
         else:
             return default

a: int = MyDict().get("a")  # Expression of type "SuperNone | int" cannot be assigned to declared type "int"
b: Union[int, str] = MyDict().get("a", "")

# vs next() function is fine:
c: int = next((x for x in [1]))
d: Union[int, str] = next((x for x in [1]), "")

This does not work, how to "dynamically" make the return typying value?

Answer 1

You achieve this typing using overload. This does not affect runtime, overload'ed signatures are for type checker exclusively. Here's how you could address this problem:

from typing import TypeVar, Union, Final, overload, Generic

_R = TypeVar('_R')
_T = TypeVar('_T')
  
_SENTINEL: Final = object()


class MyDict(Generic[_R]):
    data: dict[str, _R]
    
    def __init__(self, x: _R) ->None:
        ...    
    
    @overload
    def get(self, key: str) -> _R: ...
    @overload
    def get(self, key: str, default: _T) -> _T | _R: ...
    def get(
        self, key: str, default: object = _SENTINEL
    ) -> _T | object:
        try:
            return self.data[key]
        except KeyError:
            if default is _SENTINEL:
                raise
            else:
                return default

reveal_type(MyDict(1).get("a"))
reveal_type(MyDict(1).get("a", None))
reveal_type(MyDict(1).get("a", ""))

Here's a playground link (note that there's no need to build a separate class to create a unique sentinel)

If you don't know about overloads, you have an excellent thing - working example of next. You can just go and see how its type is defined in typeshed:

@overload
def next(__i: SupportsNext[_T]) -> _T: ...
@overload
def next(__i: SupportsNext[_T], __default: _VT) -> _T | _VT: ...