A quick and dirty script to make stub files use variable annotations

convert_stubs.py

"""
convert_stubs.py

Modifies stub files to use variable annotations instead of type comments.

To run, pass in a path to the directory containing stubs. This script will
recursively traverse it and modify all .pyi files:

    $ python3 convert_stubs.py path/to/stubs

"""
from __future__ import annotations

import sys
from pathlib import Path
import re


ANNOTATION = re.compile(r"(\s*)(\w+) = \.\.\.\s+ # type: ([\w\[\], \.]+)")


def main() -> None:
    if len(sys.argv) != 2:
        print(__doc__)
        return

    root_folder = Path(sys.argv[1])
    for pyi_path in root_folder.glob("**/*.pyi"):
        pyi_path.write_text(modify_file(pyi_path.read_text()))


def modify_file(text: str) -> str:
    out = []
    for line in text.split('\n'):
        match = ANNOTATION.match(line)
        if not match:
            out.append(line)
            continue

        indent = match.group(1)
        variable = match.group(2)
        type_hint = match.group(3).strip()
        if type_hint == 'ignore':
            out.append(line)
            continue

        out.append('{}{}: {}'.format(indent, variable, type_hint))

    return '\n'.join(out)


if __name__ == '__main__':
    main()