CODEWITHSUNDEEP
pythonmypypython-typing
PabloRdrRbl
PabloRdrRbl

Reputation: 247

Python optional types before assertion

I am using Python type hints and mypy with the following code:

from typing import Optional


def compute(
    use_wcog: bool = False,
    Xt: Optional[float] = None,
    Xd: Optional[float] = None,
    Xw: Optional[float] = None,
    Xs: Optional[float] = None,
):

    if use_wcog:
        reqs = (Xt, Xd, Xw)
        assert all([_ is not None for _ in reqs])
        res = ((Xt**2 + Xw**2) / Xd)**2
    else:
        reqs = (Xs, Xd)
        assert all([_ is not None for _ in reqs])
        res = (Xs**2 / Xd)**2

    return res

I get the following error: Unsuported operand types for ** (float and None) (mypy error)

What is the proper way of handling it?

Upvotes: 0

Views: 1492

Answers (1)

Samwise
Samwise

Reputation: 71424

Mypy isn't quite clever enough to narrow the type based on all and the list comprehension. Make it simple:

    if use_wcog:
        assert Xt and Xd and Xw
        res = ((Xt**2 + Xw**2) / Xd)**2
    else:
        assert Xs and Xd
        res = (Xs**2 / Xd)**2

FWIW, having these dependencies embedded in the body of the function makes typing largely useless IMO -- it's very easy to call this function in a way that will error at runtime, and the whole point of static typing is to make usage errors detectable statically. Given that the two implementations of the function have completely different contracts with respect to what inputs are expected and what result is returned, I'd just make it two functions.

def compute_wcog(Xt: float, Xd: float, Xw: float) -> float:
    return ((Xt**2 + Xw**2) / Xd)**2

def compute_no_wcog(Xd: float, Xs: float) -> float:
    return (Xs**2 / Xd)**2

If you really need to have a single compute function with effectively multiple interfaces, you can use the @overload decorator:

from typing import Literal, Optional, overload

@overload
def compute(use_wcog: Literal[True], Xt: float, Xd: float, Xw: float, Xs: None) -> float: ...

@overload
def compute(use_wcog: Literal[False], Xt: None, Xd: float, Xw: None, Xs: float) -> float: ...

def compute(
    use_wcog: bool = False,
    Xt: Optional[float] = None,
    Xd: Optional[float] = None,
    Xw: Optional[float] = None,
    Xs: Optional[float] = None,
):
    assert Xd is not None

    if use_wcog:
        assert Xt is not None and Xw is not None
        res = ((Xt**2 + Xw**2) / Xd)**2
    else:
        assert Xs is not None
        res = (Xs**2 / Xd)**2

    return res

compute(True, 1.0, 1.0, 1.0, None)    # works
compute(False, None, 1.0, None, 1.0)  # works
compute(True, None, 1.0, None, 1.0)   # error: No overload variant of "compute" matches argument types...

but this is obviously a lot more code for not a lot of gain. :)

Upvotes: 3

Related Questions