Proper way to add type hints to a JSON object hook

Viewed 3885

Taking the example from the json module docs:

>>> def as_complex(dct):
...     if '__complex__' in dct:
...         return complex(dct['real'], dct['imag'])
...     return dct

What would be the proper way to add type hints here? My naive approach with using a TypedDict and overloads fails:

from typing import Any, Dict, TypedDict, TypeVar, Union, overload


_T = TypeVar('_T', bound=Dict[str, Any])


class JSONDict(TypedDict):
    __complex__: Any
    real: float
    imag: float


@overload
def as_complex(dct: JSONDict) -> complex: ...
@overload
def as_complex(dct: _T) -> _T: ...


def as_complex(dct: _T) -> Union[complex, _T]:
    if '__complex__' in dct:
        return complex(dct['real'], dct['imag'])
    return dct

as mypy objects with:

main.py:14: error: Overloaded function signatures 1 and 2 overlap with incompatible return types
main.py:19: error: Overloaded function implementation cannot satisfy signature 1 due to inconsistencies in how they use type variables

Playground gist, should you want to try the sample out in the browser.

1 Answers

Obviously the mypy treats JSONDict and _T as overlapped types, most likely due to Any in the TypeVar definition. This reminds me of the following situations (it's just example):

class A: ...
class B: ...
class C: ...
class D: ...

@overload
def f(x: Union[A, B]) -> int: ...  # E: Overloaded function signatures 1 and 2 overlap with incompatible return types
@overload
def f(x: Union[B, C]) -> str: ...
def f(x): ...

We can either leave them overlapping or try to separate them. Currently, I can suggest two, rather ugly solutions.

from typing import Any, Dict, TypedDict, TypeVar, Union, overload, cast


_T = TypeVar('_T', bound=Dict[str, Any])


class JSONDict(TypedDict):
    __complex__: Any
    real: float
    imag: float


@overload
def as_complex(dct: JSONDict) -> complex: ...
@overload
def as_complex(dct: _T) -> Union[_T, complex]: ...


def as_complex(dct: Union[_T, JSONDict]) -> Union[complex, _T]:
    if '__complex__' in dct:
        return complex(dct['real'], dct['imag'])
    return cast(_T, dct)

Or try to separate

from typing import Any, Dict, TypedDict, TypeVar, Union, overload, cast

# narrowing the list of types excluding float
_T = TypeVar('_T', bound=Dict[str, Union[int, str, list, tuple, dict, set]])  
 

class JSONDict(TypedDict):
    __complex__: Any
    real: float
    imag: float


@overload
def as_complex(dct: JSONDict) -> complex: ...
@overload
def as_complex(dct: _T) -> _T: ...


def as_complex(dct: Union[_T, JSONDict]) -> Union[complex, _T]:
    if '__complex__' in dct:
        dct = cast(JSONDict, dct)
        return complex(dct['real'], dct['imag'])
    return cast(_T, dct)

Hopefully, someone will suggest a more elegant solution

Related