Type hints and @singledispatch: how do I include `Union[...]` in an extensible way?

2.4k Views Asked by At

I'm refactoring a function that converts a variety of date formats (ie. ISO 8601 string, datetime.date, datetime.datetime, etc) to a Unix timestamp.

I want the new function to use @singledispatch instead of type inspection, but I can't figure out how to retain the previous function's type hinting:

Old function: using type inspection

import datetime
from typing import Union


MyDateTimeType = Union[int, str, datetime.datetime, datetime.date, None]


# How do I retain this functionality with @singledispatch?
#                    ⬇️⬇️⬇️⬇️⬇️⬇️⬇️
def to_unix_ts(date: MyDateTimeType = None) -> Union[int, None]:
    """Convert various date formats to Unix timestamp..."""
    if type(date) is int or date is None:
        return date

    if type(date) is str:
        # Handle string argument...

    elif type(date) is datetime.datetime:
        # Handle datetime argument...

    elif type(date) is datetime.date:
        # Handle date argument...

New function: using @singledispatch

import datetime
from functools import singledispatch
from typing import Union


@singledispatch
def to_unix_ts(date) -> Union[int, None]:
    """Handle generic case (probably string type)..."""

@to_unix_ts.register
def _(date: int) -> int:
    return date


@to_unix_ts.register
def _(date: None) -> None:
    return date


@to_unix_ts.register
def _(date: datetime.datetime) -> int:
    return int(date.replace(microsecond=0).timestamp())


# etc...

I've explored building the supported types like this:

supported_types = [type for type in to_unix_ts.registry.keys()]
MyDateTimeType = Union(supported_types)  # Example, doesn't work

...so that it's extensible with future @singledispatch registrations, but I can't get it to work.

How can I add Union[...] style type hints in a @singledispatch function in an extensible way?

1

There are 1 best solutions below

0
On

I think this is what you're looking for. Note that singledispatch requires register(type(None)) instead of register(None). It doesn't support register(Union[a, b]), but you can apply multiple register decorators to a function...

import datetime
from functools import singledispatch
from typing import Union

MyDateTimeType = Union[
    int,
    None,
    datetime.datetime,
    datetime.date
]

@singledispatch
def to_unix_ts(date: MyDateTimeType) -> Union[int, None]:
    raise NotImplementedError

@to_unix_ts.register(int)
@to_unix_ts.register(type(None))
def _(date: Union[int, None]) -> Union[int, None]:
    return date

@to_unix_ts.register(str)
def _(date: str):
    # Handle string argument...

@to_unix_ts.register(datetime.datetime)
def _(date: datetime.datetime):
    # Handle datetime argument...

@to_unix_ts.register(datetime.date)
def _(date: datetime.date):
    # Handle date argument...