How to type hint functions that modify leaf type of nested structures.

[randolf-scholz]

Basically I have a function that acts on arbitrarily nested data, for simplicity consider the following which casts all int to float in a nested data structure.

def foo(x):
    match x:
        case int():
            return float(x)
        case list():
            return [foo(i) for i in x]
        case tuple():
            return tuple(foo(i) for i in x)
        case dict():
            return {k: foo(v) for k, v in x.items()}
    raise TypeError(f"Unsupported type: {type(x)}")

How can I type hint foo such that the following assertions hold

y: list[int] = [1]
z: dict[str, int] = {"a": 1}
w: list[int | tuple[int, ...]] = [1, (1, 2)]
assert_type(foo(y), list[float])
assert_type(foo(z), dict[str, float])
assert_type(foo(w), list[float | tuple[float, ...]])
# etc. etc.

[erictraut]

A recursive type alias should do the trick. The following works fine in pyright and mypy.

from typing import TypeAlias, TypeVar

T = TypeVar("T")
Nested: TypeAlias = T | list["Nested[T]"] | tuple["Nested[T]", ...] | dict[str, "Nested[T]"]

def foo(x: Nested[int]) -> Nested[float]:
    ...

[randolf-scholz]

That doesn't seem to work: https://mypy-play.net/?mypy=latest&python=3.11&gist=19d99727ccc6fc388b7519f769ca021b,
e.g. foo(y) will be considered Nested[float] instead of list[float]. Also, there is an issue with invariance.

I wonder if this is another case where one needs some typing construct not available in python/mypy. (Higher Kinded Types / Type Maps)

[erictraut]

It works fine if you use the correct type annotations.

y: list[Nested[int]] = [1]
z: dict[str, Nested[int]] = {"a": 1}
w: list[Nested[int]] = [1, (1, 2)]

assert_type(foo(y), Nested[float])
assert_type(foo(z), Nested[float])
assert_type(foo(w), Nested[float])

The type parameter for list is invariant, so list[int] is not type compatible with list[Nested[int]]. If you create a list object, it needs to "start life asalist[Nested[int]]. The same is true for dict`.

[randolf-scholz]

I wanted to know whether it is possible to type hint foo in such a way that these assertions hold, not change them. Having the result be Nested[float] is not helpful, because this widens too much for certain applications.

If the input is a list[int], I really do want the output to be considered list[float], since it may be passed into another function that only accepts list[float] and not Nested[float]! Any appropriate subtype of Nested[int] should result in the corresponding subtype of Nested[float].

One could add around a hundred @overloads to get this behavior for 2-3 levels of nesting, but can it be done generically?