How to properly annotate a ContextManager in PyCharm?

Question

How can I annotate the yield type of a contextmanager in PyCharm so that it properly guesses the type of the value used in the with clauses - just as it guesses that the f created in with open(...) as f is a file?

For example, I have a context manager like this:

@contextlib.contextmanager
def temp_borders_file(geometry: GEOSGeometry, name='borders.json'):
    with TemporaryDirectory() as temp_dir:
        borders_file = Path(dir) / name
        with borders_file.open('w+') as f:
            f.write(geometry.json)
        yield borders_file

with temp_borders_file(my_geom) as borders_f:
    do_some_code_with(borders_f...)

How do I let PyCharm know that every borders_f created like this is a pathlib.Path (and thus enable the autocompletion for the Path methods on border_f)? Of course, I can make a comment like # type: Path after every with statement, but it seems that this can be done by properly annotating temp_border_file.

I tried Path, typing.Iterator[Path] and typing.Generator[Path, None, None] as the return type of temp_border_file, as well as adding # type: Path on borders_file within the context manager's code, but it seems like it doesn't help.

Answer 1

I believe you can use ContextManager from typing, e.g.:

import contextlib
from typing import ContextManager
from pathlib import Path


@contextlib.contextmanager
def temp_borders_file() -> ContextManager[Path]:
    pass


with temp_borders_file() as borders_f:
    borders_f  # has type Path here

Answer 2

This is a current PyCharm issue: PY-36444

A workaround for the issue is to rewrite an example code of:

from contextlib import contextmanager

@contextmanager
def generator_function():
    yield "some value"

with generator_function() as value:
    print(value.upper())  # no PyCharm autocompletion

to

from contextlib import contextmanager
from typing import ContextManager

def wrapper() -> ContextManager[str]:
    @contextmanager
    def generator_function():
        yield "some value"

    return generator_function()

with wrapper() as value:
    print(value.upper())  # PyCharm autocompletion works

There is also an easier workaround of annotating the return type with ContextManager[str] but there are multiple reasons against this:

• mypy will correctly emit errors for this annotation, as described in more detail in the PyCharm issue.
• this is not guaranteed to work in the future because PyCharm will hopefully fix the issue and therefore break this workaround