How to make type-annotation-only type assertions?

Question

I have two functions:

def get_foo(params) -> Optional[str]
def bar(foo: str)

And a function which chains these functions together:

def f(params):
    # other stuff up here
    foo = get_foo(params)
    return bar(foo)

I know based on the other things happening in my function that the result of get_foo will never be None.

When I run mypy against this file, I of course get errors:

error: Argument 1 of "bar" has incompatible type "Optional[str]"; expected "str"

which makes sense.

I could add an assert foo is not None statement, but this is hot-path code and in my tests it has measurable performance impact. I would like to make a type assertion for mypy only. How do I do that?

EDIT: I also tried adding a comment #type: str after the assignment statement, but this generated a similar error

Answer 1

You're not going to be happy about this. The officially designed way to assert to static type checkers that a value has a specific type is typing.cast, which is an actual function with a real runtime cost, I believe more expensive than the assert you want to replace. It just returns its second argument unchanged, but it's still got function call overhead. Python's type annotation system is not designed with a zero-overhead type assertion syntax.

As an alternative, you could use Any as an "escape hatch". If you annotate foo with type Any, mypy should allow the bar call. Local variable annotations have no runtime cost, so the only runtime cost is an extra local variable store and lookup:

from typing import Any

def f(params):
    foo: Any = get_foo(params)
    return bar(foo)

Aside from that, your best option may be to use the assert and run Python with the -O flag, which disables asserts.

Answer 2

You can use TYPE_CHECKING variable that is False during runtime but True during type checking. This would avoid performance hit of assert:

from typing import TYPE_CHECKING

def f(params):
    # other stuff up here
    foo = get_foo(params)
    
    if TYPE_CHECKING:
        assert foo is not None
    
    return bar(foo)