How to deal with Optional[int] in Protocol classes

Question

I have created the following Protocol class:

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional, Protocol, runtime_checkable

@runtime_checkable
class PipelineProcess(Protocol):

    def update(self):
        ...

    def set_inputs(self, inputs: Dict[str, int]):
        ...

    def get_outputs(self) -> Dict[str, int] | None:
        ...

I use it for this example process and others which are built the same way:

@dataclass
class ExampleProcess:
    in_value: Optional[int | None] = None
    out_value: Optional[int | None] = None

    def update(self):
        assert(self.in_value is not None)
        self.out_value = self.in_value * 2

    def set_inputs(self, values: Dict[str, int]):
        assert "value" in values.keys()
        self.in_value = values["value"]

    def get_outputs(self) -> Dict[str, int]:
        return {"value": self.out_value}

If I analyze this with mypy I got the following error:

error: Dict entry 0 has incompatible type "str": "Optional[int]"; expected "str": "int"

To solve this I used dict[str, Optional[int]] as the return type of the get_outputs method, which seems a bit weird, since I used Optional for the dataclass attribute.

Or is the Union for Optional redundant and I have to use Optional[int] always instead of int because it implicitly says None is also valid?

Answer 1

Since out_value is declared to be of the union type int | None, setting it as value of a dictionary expected to be of type dict[str, int] is incorrect. So you need to assure the type checker (and yourself) that by the time you put into it the dictionary out_value is in fact an int. This can be easily accomplished with an assert inside get_outputs:

class ExampleProcess:
    in_value: int | None = None
    out_value: int | None = None

    def update(self) -> None:
        assert (self.in_value is not None)
        self.out_value = self.in_value * 2

    def set_inputs(self, values: dict[str, int]) -> None:
        assert "value" in values.keys()
        self.in_value = values["value"]

    def get_outputs(self) -> dict[str, int]:
        assert self.out_value is not None
        return {"value": self.out_value}

---

Side note: Contrasting this with the other two assert statements in your code, I wonder why you put them there. While get_outputs would still work without the additional assert, the other two methods would still fail more or less transparently, if you removed the assert statements.

update would raise TypeError: unsupported operand type(s) for *: 'NoneType' and 'int' and point to the offending line self.out_value = self.in_value * 2.

set_inputs would raise KeyError: 'value' pointing to self.in_value = values["value"].

But I guess if you wanted to be very explicit or even substitute a custom exception class, you could do it like this:

class ProcessError(Exception):
    pass

class ExampleProcess:
    in_value: int | None = None
    out_value: int | None = None

    def update(self) -> None:
        if self.in_value is None:
            raise ProcessError(
                "Cannot update while `in_value` is None"
            )
        self.out_value = self.in_value * 2

    def set_inputs(self, values: dict[str, int]) -> None:
        try:
            self.in_value = values["value"]
        except KeyError:
            raise ProcessError(
                "Setting inputs requires a dictionary with a `value` key."
            ) from None

    def get_outputs(self) -> dict[str, int]:
        assert self.out_value is not None, "The `out_value` must be set first!"
        return {"value": self.out_value}

EDIT: Just realized that the first assert inside update actually does serve the same purpose as the one I suggested for get_outputs. Only the second one in set_inputs is technically unnecessary.

But there you see that you can assure the static type checker in different ways. One is an if-branch based on the type distinction and one is just a straight up assertion.