matplotlib

1
import ast
2
import os
3
import pathlib
4
import subprocess
5
import sys
6
import tempfile
7

8
root = pathlib.Path(__file__).parent.parent
9

10
lib = root / "lib"
11
mpl = lib / "matplotlib"
12

13

14
class Visitor(ast.NodeVisitor):
15
    def __init__(self, filepath, output):
16
        self.filepath = filepath
17
        self.context = list(filepath.with_suffix("").relative_to(lib).parts)
18
        self.output = output
19

20
    def visit_FunctionDef(self, node):
21
        # delete_parameter adds a private sentinel value that leaks
22
        # we do not want that sentinel value in the type hints but it breaks typing
23
        # Does not apply to variadic arguments (args/kwargs)
24
        for dec in node.decorator_list:
25
            if "delete_parameter" in ast.unparse(dec):
26
                deprecated_arg = dec.args[1].value
27
                if (
28
                    node.args.vararg is not None
29
                    and node.args.vararg.arg == deprecated_arg
30
                ):
31
                    continue
32
                if (
33
                    node.args.kwarg is not None
34
                    and node.args.kwarg.arg == deprecated_arg
35
                ):
36
                    continue
37

38
                parents = []
39
                if hasattr(node, "parent"):
40
                    parent = node.parent
41
                    while hasattr(parent, "parent") and not isinstance(
42
                        parent, ast.Module
43
                    ):
44
                        parents.insert(0, parent.name)
45
                        parent = parent.parent
46
                self.output.write(f"{'.'.join(self.context + parents)}.{node.name}\n")
47
                break
48

49
    def visit_ClassDef(self, node):
50
        for dec in node.decorator_list:
51
            if "define_aliases" in ast.unparse(dec):
52
                parents = []
53
                if hasattr(node, "parent"):
54
                    parent = node.parent
55
                    while hasattr(parent, "parent") and not isinstance(
56
                        parent, ast.Module
57
                    ):
58
                        parents.insert(0, parent.name)
59
                        parent = parent.parent
60
                aliases = ast.literal_eval(dec.args[0])
61
                # Written as a regex rather than two lines to avoid unused entries
62
                # for setters on items with only a getter
63
                for substitutions in aliases.values():
64
                    parts = self.context + parents + [node.name]
65
                    self.output.write(
66
                        "\n".join(
67
                            f"{'.'.join(parts)}.[gs]et_{a}\n" for a in substitutions
68
                        )
69
                    )
70
        for child in ast.iter_child_nodes(node):
71
            self.visit(child)
72

73

74
with tempfile.TemporaryDirectory() as d:
75
    p = pathlib.Path(d) / "allowlist.txt"
76
    with p.open("wt") as f:
77
        for path in mpl.glob("**/*.py"):
78
            v = Visitor(path, f)
79
            tree = ast.parse(path.read_text())
80

81
            # Assign parents to tree so they can be backtraced
82
            for node in ast.walk(tree):
83
                for child in ast.iter_child_nodes(node):
84
                    child.parent = node
85

86
            v.visit(tree)
87
    proc = subprocess.run(
88
        [
89
            "stubtest",
90
            "--mypy-config-file=pyproject.toml",
91
            "--allowlist=ci/mypy-stubtest-allowlist.txt",
92
            f"--allowlist={p}",
93
            "matplotlib",
94
        ],
95
        cwd=root,
96
        env=os.environ | {"MPLBACKEND": "agg"},
97
    )
98
    try:
99
        os.unlink(f.name)
100
    except OSError:
101
        pass
102

103
sys.exit(proc.returncode)
104