How to type a function that always return a list? #995

cdce8p · 2021-12-26T17:18:24Z

cdce8p
Dec 26, 2021

I'm trying to add type annotations to a function used for config validation.

def ensure_list(value):
    """Wrap value in list if it is not one."""
    if value is None:
        return []
    return value if isinstance(value, list) else [value]

At first glance it looks strait forward, but I'm struggling with some edge cases. This is what I've got so far

from typing import Any, TypeVar, overload

T = TypeVar("T")

@overload
def ensure_list(value: None) -> list[Any]:
    ...

@overload
def ensure_list(value: list[T]) -> list[T]:
    ...

@overload
def ensure_list(value: T) -> list[T]:
    ...

def ensure_list(value: T) -> list[T] | list[Any]:
    """Wrap value in list if it is not one."""
    if value is None:
        return []
    return value if isinstance(value, list) else [value]

My test cases

class X:
    a1: str
    a2: None
    a3: list[int]

    b1: list[int] | int
    b2: str | list[str]

    def __init__(self) -> None:
        reveal_type(ensure_list(X.a1))  # should be 'list[str]'
        reveal_type(ensure_list(X.a2))  # should be 'list[Any]'
        reveal_type(ensure_list(X.a3))  # should be 'list[int]'

        reveal_type(ensure_list(X.b1))  # should be 'list[int]'
        reveal_type(ensure_list(X.b2))  # should be 'list[str]'

a1 - a3 work fine, but what can I do for b1 and b2?

--
My initial idea was to add additional overloads with different TypeVars, but I can't seem to prevent the TypeVar from binding to the Union type instead. If I don't add them after the last overload but instead between the 2 and 3, at least pyright seems to get b1 and b2 right, although at the cost of a1.

T1 = TypeVar("T1")
T2 = TypeVar("T2")

@overload
def ensure_list(value: list[T1] | T2) -> list[T1] | list[T2]:
    ...

@overload
def ensure_list(value: T1 | list[T2]) -> list[T1] | list[T2]:
    ...

Mypy output

After the last one

error: Overloaded function signature 4 will never be matched: signature 3's parameter type(s) are the same or broader
error: Overloaded function signature 5 will never be matched: signature 3's parameter type(s) are the same or broader
error: Overloaded function signature 5 will never be matched: signature 4's parameter type(s) are the same or broader
note: Revealed type is "builtins.list[builtins.str*]"
note: Revealed type is "builtins.list[Any]"
note: Revealed type is "builtins.list[builtins.int*]"
note: Revealed type is "builtins.list[Union[builtins.list[builtins.int], builtins.int]]"
note: Revealed type is "builtins.list[Union[builtins.str, builtins.list[builtins.str]]]"

Between 2 and 3

error: Overloaded function signature 4 will never be matched: signature 3's parameter type(s) are the same or broader
error: Overloaded function signature 5 will never be matched: signature 3's parameter type(s) are the same or broader
error: Overloaded function signature 5 will never be matched: signature 4's parameter type(s) are the same or broader
note: Revealed type is "Union[builtins.list[<nothing>], builtins.list[builtins.str*]]"
note: Revealed type is "builtins.list[Any]"
note: Revealed type is "builtins.list[builtins.int*]"
note: Revealed type is "builtins.list[Union[builtins.list[builtins.int], builtins.int]]"
note: Revealed type is "builtins.list[Union[builtins.str, builtins.list[builtins.str]]]"

Pyright output

After the last one

error: Overload 3 for "ensure_list" will never be used because its parameters overlap overload 1 (reportOverlappingOverload)
error: Overload 4 for "ensure_list" will never be used because its parameters overlap overload 3 (reportOverlappingOverload)
error: Overload 5 for "ensure_list" will never be used because its parameters overlap overload 3 (reportOverlappingOverload)
info: Type of "ensure_list(X.a1)" is "list[str]"
info: Type of "ensure_list(X.a2)" is "list[Any]"
info: Type of "ensure_list(X.a3)" is "list[int]"
info: Type of "ensure_list(X.b1)" is "list[list[int] | int]"
info: Type of "ensure_list(X.b2)" is "list[str | list[str]]"

Between 2 and 3

error: Overload 4 for "ensure_list" will never be used because its parameters overlap overload 3 (reportOverlappingOverload)
error: Overload 5 for "ensure_list" will never be used because its parameters overlap overload 1 (reportOverlappingOverload)
info: Type of "ensure_list(X.a1)" is "list[T1@ensure_list] | list[str]"
info: Type of "ensure_list(X.a2)" is "list[Any]"
info: Type of "ensure_list(X.a3)" is "list[int]"
info: Type of "ensure_list(X.b1)" is "list[int]"
info: Type of "ensure_list(X.b2)" is "list[str]"

Answered by cdce8p

Jan 3, 2022

By coincidence I found an overload variation that works in both mypy and pyright.

@overload
def ensure_list(value: None) -> list[Any]: ...

@overload
def ensure_list(value: list[T]) -> list[T]: ...

@overload
def ensure_list(value: list[T] | T) -> list[T]: ...

def ensure_list(value: T | None) -> list[T] | list[Any]: ...

Extracting the concrete overload list[T] -> list[T] while also keeping the union one Eric suggested list[T] | T - > list[T] seems to work. Note the implementation type for value is only T | None. Unfortunately list[T] | T | None doesn't work in mypy although arguably a bit more fitting. T | None does work though in both mypy and pyright.

error: Overloaded function impleme…

View full answer

erictraut · 2021-12-28T09:08:35Z

erictraut
Dec 28, 2021
Collaborator

Any time you attempt to use a type variable (especially when it's unbound and unconstrained) in a union with other types, you're likely to run into problems. This inevitably generates ambiguities for the constraint solver.

Take for example the annotation T | list[T]. If we attempt to match the type int | list[int] to this annotation, there are two legitimate solutions: T = int and T = int | list[int]. I've added many heuristics to pyright's constraint solver to try to give the "obvious" solution (which typically corresponds to the simplest solution) in most cases. It is possible to defeat these heuristics, as you have done in the example above by splitting list[T] and T across separate overloads. If you instead switch them to a single overload, pyright will give you the solutions you're looking for in your test cases.

@overload
def ensure_list(value: None) -> list[Any]:
    ...

@overload
def ensure_list(value: list[T] | T) -> list[T]:
    ...

def ensure_list(value: list[T] | T | None) -> list[T] | list[Any]:
    """Wrap value in list if it is not one."""
    if value is None:
        return []
    return cast(list[T], value) if isinstance(value, list) else [value]

Note that I needed to add a cast here. This is required because of another heuristic in the isinstance type narrowing logic.

The above sample mostly works for mypy, but it's emitting an error for X.a3 test case. I'm not sure why. It looks like it could be a bug.

2 replies

cdce8p Dec 29, 2021
Author

Thanks for the response! It does look like the best option.

Any time you attempt to use a type variable (especially when it's unbound and unconstrained) in a union with other types, you're likely to run into problems. This inevitably generates ambiguities for the constraint solver.

I would have liked something like: "Bind to everything except union types". Just don't think this is possible with TypeVars at the moment. Would it be worth opening an issue / feature request for it here?

The above sample mostly works for mypy, but it's emitting an error for X.a3 test case. I'm not sure why. It looks like it could be a bug.

I initially dismissed your solution due to the mypy issue. That's what lead me to split the overload cases further. You're right though that it looks like a bug.

--
There are also some other cases that produce interesting / unexpected results.

class Y:
    c1: list
    c2: list[Any]
    c3: list[None]

    def __init__(self) -> None:
        reveal_type(ensure_list(Y.c1))  # should be 'list[Any]' or 'list[Unknown]'
        reveal_type(ensure_list(Y.c2))  # should be 'list[Any]'
        reveal_type(ensure_list(Y.c3))  # should be 'list[None]'

# mypy
note: Revealed type is "builtins.list[Union[builtins.list[Any], Any]]"
note: Revealed type is "builtins.list[Union[builtins.list[Any], Any]]"
note: Revealed type is "builtins.list[<nothing>]"
error: Argument 1 to "ensure_list" has incompatible type "List[None]"; expected "List[<nothing>]"

# pyright
info: Type of "ensure_list(Y.c1)" is "list[list[Unknown]]"
info: Type of "ensure_list(Y.c2)" is "list[list[Any]]"
info: Type of "ensure_list(Y.c3)" is "list[list[None]]"

erictraut Dec 29, 2021
Collaborator

These are cases where the heuristics in pyright didn't produce the result I would have expected. I've tweaked the heuristics to do a better job in these cases. This will be included in the next release of pyright.

cdce8p · 2022-01-03T13:01:30Z

cdce8p
Jan 3, 2022
Author

By coincidence I found an overload variation that works in both mypy and pyright.

@overload
def ensure_list(value: None) -> list[Any]: ...

@overload
def ensure_list(value: list[T]) -> list[T]: ...

@overload
def ensure_list(value: list[T] | T) -> list[T]: ...

def ensure_list(value: T | None) -> list[T] | list[Any]: ...

Extracting the concrete overload list[T] -> list[T] while also keeping the union one Eric suggested list[T] | T - > list[T] seems to work. Note the implementation type for value is only T | None. Unfortunately list[T] | T | None doesn't work in mypy although arguably a bit more fitting. T | None does work though in both mypy and pyright.

error: Overloaded function implementation does not accept all possible arguments of signature 2  [misc]

Test cases

class X:
    a1: str
    a2: None
    a3: list[int]

    b1: list[int] | int
    b2: str | list[str]

    def __init__(self) -> None:
        reveal_type(ensure_list(X.a1))  # should be 'list[str]'
        reveal_type(ensure_list(X.a2))  # should be 'list[Any]'
        reveal_type(ensure_list(X.a3))  # should be 'list[int]'

        reveal_type(ensure_list(X.b1))  # should be 'list[int]'
        reveal_type(ensure_list(X.b2))  # should be 'list[str]'


class Y:
    c1: list
    c2: list[Any]
    c3: list[None]

    def __init__(self) -> None:
        reveal_type(ensure_list(Y.c1))  # should be 'list[Any]' or 'list[Unknown]'
        reveal_type(ensure_list(Y.c2))  # should be 'list[Any]'
        reveal_type(ensure_list(Y.c3))  # should be 'list[None]'

Mypy output

note: Revealed type is "builtins.list[builtins.str*]"
note: Revealed type is "builtins.list[Any]"
note: Revealed type is "builtins.list[builtins.int*]"
note: Revealed type is "builtins.list[builtins.int*]"
note: Revealed type is "builtins.list[builtins.str*]"

note: Revealed type is "builtins.list[Any]"
note: Revealed type is "builtins.list[Any]"
note: Revealed type is "builtins.list[None]"

Pyright output

info: Type of "ensure_list(X.a1)" is "list[str]"
info: Type of "ensure_list(X.a2)" is "list[Any]"
info: Type of "ensure_list(X.a3)" is "list[int]"
info: Type of "ensure_list(X.b1)" is "list[int]"
info: Type of "ensure_list(X.b2)" is "list[str]"

info: Type of "ensure_list(Y.c1)" is "list[Unknown]"
info: Type of "ensure_list(Y.c2)" is "list[Any]"
info: Type of "ensure_list(Y.c3)" is "list[None]"

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

How to type a function that always return a list? #995

{{title}}

Replies: 2 comments 2 replies

{{title}}

{{title}}

{{title}}

{{title}}

Select a reply

How to type a function that always return a list? #995

cdce8p Dec 26, 2021

Replies: 2 comments · 2 replies

erictraut Dec 28, 2021 Collaborator

cdce8p Dec 29, 2021 Author

erictraut Dec 29, 2021 Collaborator

cdce8p Jan 3, 2022 Author

cdce8p
Dec 26, 2021

Replies: 2 comments 2 replies

erictraut
Dec 28, 2021
Collaborator

cdce8p Dec 29, 2021
Author

erictraut Dec 29, 2021
Collaborator

cdce8p
Jan 3, 2022
Author