How to properly annotate a ContextManager in PyCharm?

问题

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.

回答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

回答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

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

来源：https://stackoverflow.com/questions/49335263/how-to-properly-annotate-a-contextmanager-in-pycharm

标签

pycharm

typing

contextmanager