pytorch

1
import re
2
import sys
3
from pathlib import Path
4

5
from mypy.plugin import Plugin
6

7

8
def get_correct_mypy_version():
9
    # there's probably a more elegant way to do this
10
    (match,) = re.finditer(
11
        r"mypy==(\d+(?:\.\d+)*)",
12
        (
13
            Path(__file__).parent.parent / ".ci" / "docker" / "requirements-ci.txt"
14
        ).read_text(),
15
    )
16
    (version,) = match.groups()
17
    return version
18

19

20
def plugin(version: str):
21
    correct_version = get_correct_mypy_version()
22
    if version != correct_version:
23
        print(
24
            f"""\
25
You are using mypy version {version}, which is not supported
26
in the PyTorch repo. Please switch to mypy version {correct_version}.
27

28
For example, if you installed mypy via pip, run this:
29

30
    pip install mypy=={correct_version}
31

32
Or if you installed mypy via conda, run this:
33

34
    conda install -c conda-forge mypy={correct_version}
35
""",
36
            file=sys.stderr,
37
        )
38
    return Plugin
39