Python Type Annotations: Mark item in tuple as optional

user_6064820 0
Viewed 4220

I have a function that returns an 2-tuple but optionally a 3-tuple. How do I describe this with type annotations?

For example:

from typing import Tuple

def example(i):
    # type: (int) -> Tuple[int, int, <what to put here?>]
    if i < 10:
       return (i, i+1, 1)
    else:
       return (i, i+1)

I can use Union like below, but it seems quite messy.

# type: (int) -> Union[Tuple[int, int], Tuple[int, int, int]]
mypy static-typing python
1 Answers

As pointed out in the comments, the union is probably the better approach. If the signature looks messy, you can use type aliases like so:

from typing import Tuple, Union

MyType = Union[Tuple[int, int], Tuple[int, int, int]]

def example1(i):
    # type: (int) -> MyType
    ...snip...

An alternative approach would be to use the "indefinite length" Tuple type. You would basically give up on encoding the exact length of the tuple but in exchange can normalize the return type and avoid that union. (If your code does rely on the length, this probably isn't the best approach).

def example2(i):
    # type: (int) -> Tuple[int, ...]
    ...snip...

However, a somewhat more drastic approach would be to consider restructuring your code to avoid this type of situation.

After all, if you're returning these two distinct types, the caller of your function is probably going to need to check the length anyways, right?

In that case, one idea is to give up on returning a tuple and instead return either a NamedTuple or a custom class, both of which have an optional field. You can then convert your "length" check into a "is this field set to None" check.

I suppose, in a sense, the NamedTuple approach also fulfills your original request in a way, so long as you don't mind converting your Tuples/the added overhead.

from typing import NamedTuple, Optional

MyType2 = NamedTuple('MyType2', (
    ('x', int),
    ('y', int),
    ('z', Optional[int]),
))

class MyType3(object):
    def __init__(self, x, y, z):
        # type: (int, int, Optional[int]) -> None
        self.x = x
        self.y = y
        self.z = z

(The "custom class" approach will likely be more elegant once PEP 557 is accepted and integrated into the language).

The other approach is to split your function up into two if you know ahead of time which kind of tuple you're expecting. Then, you can just call the function that appropriate type.

user_6465430
Related