python mypy: sorted on a list of objects on an optional member

stdklr 0
Viewed 22

I want to sort a list of objects with an optional member. If any of the list's objects have None for the optional member I will not sort. I realize this by an all() gate. Mypy doesn't like that...

This is the minimal example for my question:

from typing import Optional
from pydantic import BaseModel


class Test(BaseModel):
    member: Optional[int]

tests = [
    Test(member=1),
    Test(member=5),
    Test(member=2),
]

if all(test.member is not None for test in tests):
    tests = sorted(
        tests,
        key=lambda x: x.member,
        reverse=False,
    )
else:
    raise ValueError
    
print(tests)

This leads to mypy error

test.py:17: error: Argument "key" to "sorted" has incompatible type "Callable[[Test], Optional[int]]"; expected "Callable[[Test], Union[SupportsDunderLT, SupportsDunderGT]]"
test.py:17: error: Incompatible return value type (got "Optional[int]", expected "Union[SupportsDunderLT, SupportsDunderGT]")
Found 2 errors in 1 file (checked 1 source file)

If I adjust the lambda function

tests = sorted(
    tests,
    key=lambda x: x.member if x.member is not None else 0,
    reverse=False,
)

mypy is happy, but I don't find that really beautiful as the all() gate already takes care.

Do I miss something, is there a better solution? Shouldn't mypy understand the gate?

mypy sorting python
1 Answers

Protocol and TypeGuard may be helpful, but here you need to change the function sorted to the method list.sort:

import sys
from dataclasses import dataclass
from typing import Optional, Protocol

if sys.version_info >= (3, 10):
    from typing import TypeGuard
else:
    from typing_extensions import TypeGuard

@dataclass    # pydantic is not installed. Here I use dataclass instead.
class Test:
    member: Optional[int]


class TestProtocol(Protocol):
    member: int


def all_have_member(iterable: list[Test]) -> TypeGuard[list[TestProtocol]]:
    return all(test.member is not None for test in iterable)


tests = [
    Test(member=1),
    Test(member=5),
    Test(member=2),
]

if all_have_member(tests):
    tests.sort(
        key=lambda x: x.member,
        reverse=False,
    )
else:
    raise ValueError

print(tests)
Mechanic Pig0
Related