Vendoring logic cleanups + Move typing stub generation to dedic… (#7334)

Vendoring logic cleanups + Move typing stub generation to dedicated module

Author: pradyunsg

tools/automation/vendoring/__init__.py

@@ -3,7 +3,6 @@
 # The following comment should be removed at some point in the future.
 # mypy: disallow-untyped-defs=False
 
-import os
 import re
 import shutil
 import tarfile
@@ -13,7 +12,7 @@
 import invoke
 import requests
 
-TASK_NAME = 'update'
+from .typing import generate_stubs
 
 FILE_WHITE_LIST = (
     'Makefile',
@@ -49,7 +48,7 @@ def remove_all(paths):
 
 
 def log(msg):
-    print('[vendoring.%s] %s' % (TASK_NAME, msg))
+    print('[vendoring.update] ' + msg)
 
 
 def _get_vendor_dir(ctx):
@@ -270,36 +269,10 @@ def update_stubs(ctx):
     vendored_libs = detect_vendored_libs(vendor_dir)
 
     print("[vendoring.update_stubs] Add mypy stubs")
+    generate_stubs(vendor_dir, vendored_libs)
 
-    extra_stubs_needed = {
-        # Some projects need stubs other than a simple <name>.pyi
-        "six": [
-            "six.__init__",
-            "six.moves.__init__",
-            "six.moves.configparser",
-        ],
-        # Some projects should not have stubs coz they're single file modules
-        "appdirs": [],
-        "contextlib2": [],
-    }
 
-    for lib in vendored_libs:
-        if lib not in extra_stubs_needed:
-            (vendor_dir / (lib + ".pyi")).write_text("from %s import *" % lib)
-            continue
-
-        for selector in extra_stubs_needed[lib]:
-            fname = selector.replace(".", os.sep) + ".pyi"
-            if selector.endswith(".__init__"):
-                selector = selector[:-9]
-
-            f_path = vendor_dir / fname
-            if not f_path.parent.exists():
-                f_path.parent.mkdir()
-            f_path.write_text("from %s import *" % selector)
-
-
-@invoke.task(name=TASK_NAME, post=[update_stubs])
+@invoke.task(name="update", post=[update_stubs])
 def main(ctx):
     vendor_dir = _get_vendor_dir(ctx)
     log('Using vendor dir: %s' % vendor_dir)

tools/automation/vendoring/typing.py

@@ -0,0 +1,59 @@
+"""Logic for adding static typing related stubs of vendored dependencies.
+
+We autogenerate `.pyi` stub files for the vendored modules, when vendoring.
+These .pyi files are not distributed (thanks to MANIFEST.in). The stub files
+are merely `from ... import *` but they do what they're supposed to and mypy
+is able to find the correct declarations using these files.
+"""
+
+import os
+from pathlib import Path
+from typing import Dict, Iterable, List, Tuple
+
+EXTRA_STUBS_NEEDED = {
+    # Some projects need stubs other than a simple <name>.pyi
+    "six": [
+        "six.__init__",
+        "six.moves.__init__",
+        "six.moves.configparser",
+    ],
+    # Some projects should not have stubs because they're a single module
+    "appdirs": [],
+    "contextlib2": [],
+}  # type: Dict[str, List[str]]
+
+
+def determine_stub_files(lib):
+    # type: (str) -> Iterable[Tuple[str, str]]
+    # There's no special handling needed -- a <libname>.pyi file is good enough
+    if lib not in EXTRA_STUBS_NEEDED:
+        yield lib + ".pyi", lib
+        return
+
+    # Need to generate the given stubs, with the correct import names
+    for import_name in EXTRA_STUBS_NEEDED[lib]:
+        rel_location = import_name.replace(".", os.sep) + ".pyi"
+
+        # Writing an __init__.pyi file -> don't import from `pkg.__init__`
+        if import_name.endswith(".__init__"):
+            import_name = import_name[:-9]
+
+        yield rel_location, import_name
+
+
+def write_stub(destination, import_name):
+    # type: (Path, str) -> None
+    # Create the parent directories if needed.
+    if not destination.parent.exists():
+        destination.parent.mkdir()
+
+    # Write `from ... import *` in the stub file.
+    destination.write_text("from %s import *" % import_name)
+
+
+def generate_stubs(vendor_dir, libraries):
+    # type: (Path, List[str]) -> None
+    for lib in libraries:
+        for rel_location, import_name in determine_stub_files(lib):
+            destination = vendor_dir / rel_location
+            write_stub(destination, import_name)