mypy: place stronger requirements on a type without runtime cost

user_222189 0
Viewed 64

I am receiving messages from a remote party, that get decoded to classes that look like this:

class SomeMessage(MessageType):
    foo: Optional[int]
    bar: Optional[str]
    quux: Optional[AnotherMessageType]

Most fields are always optional. The wire message may or may not contain them.

However, my code does validation beforehand and then mostly assumes that the appropriate fields are present.

E.g.:

def process_all_foos(messages: List[SomeMessage]):
    messages_with_foo = [m for m in messages if m.foo is not None]
    foo_sum = sum(m.foo for m in messages_with_foo)

That code fails to type-check because mypy doesn't carry over the previous check.

Is there a possibility to make the above code type-check somehow without incurring runtime cost?

My idea was creating a protocol with a stronger requirement:

class SomeMessageWithFoo(Protocol):
    foo: int

def convert_foo(m: SomeMessage) -> SomeMessageWithFoo:
    assert m.foo is not None
    return m

# or, ideally:
def process_all_foos(messages: List[SomeMessage]):
    messages_with_foo: List[SomeMessageWithFoo] = [m for m in messages if m.foo is not None]
    foo_sum = sum(m.foo for m in messages_with_foo)

However, mypy rejects that:

test.py:16: error: Incompatible return value type (got "SomeMessage", expected "SomeMessageWithFoo")  [return-value]
test.py:16: note: Following member(s) of "SomeMessage" have conflicts:
test.py:16: note:     foo: expected "int", got "Optional[int]"
test.py:20: error: List comprehension has incompatible type List[SomeMessage]; expected List[SomeMessageWithFoo]  [misc]
test.py:20: note: Following member(s) of "SomeMessage" have conflicts:
test.py:20: note:     foo: expected "int", got "Optional[int]"

I considered using cast(SomeMessageWithFoo, m), or, more likely, # type: ignore because that does not have a runtime cost. (I'm in an embedded environment so I do care about the function call that cast does.) But that also lets me not check the field.

I.e., the following incorrectly typechecks:

def convert_foo(m: SomeMessage) -> SomeMessageWithFoo:
    # someone commented this out:
    #assert m.foo is not None
    return cast(SomeMessageWithFoo, m)

Is there a way to accomplish this that makes the code type-check, but still actually validates the requirements?

python-typing mypy python
1 Answers

I don't get it. Isn't this the obvious solution, since it a) makes your code faster, since you're doing only one iteration instead of two, and b) fixes the type checks:

def process_all_foos(messages: List[SomeMessage]):
    foo_sum = sum(m.foo for m in messages if m.foo is not None)
    # use foo_sum
user_98358720
Related