Add --install-types option for installing missing stub packages

JukkaL · JukkaL · commit 0fd28c152e17 · 2020-06-26T17:21:54.000+01:00
It can used in a normal mypy run (missing stub packages will be installed
at the end). We also record the missing stub packages in a file so that
another "mypy --install-types" run (with no other arguments) can be used
to just install stub packages, without type checking anything.
diff --git a/mypy/build.py b/mypy/build.py
@@ -268,6 +268,8 @@ def _build(sources: List[BuildSource],
             reports.finish()
         if not cache_dir_existed and os.path.isdir(options.cache_dir):
             add_catch_all_gitignore(options.cache_dir)
+        if os.path.isdir(options.cache_dir):
+            record_missing_stub_packages(options.cache_dir, manager.missing_stub_packages)
 
 
 def default_data_dir() -> str:
@@ -640,6 +642,8 @@ def __init__(self, data_dir: str,
         # the semantic analyzer, used only for testing. Currently used only by the new
         # semantic analyzer.
         self.processed_targets = []  # type: List[str]
+        # Missing stub packages encountered.
+        self.missing_stub_packages = set()  # type: Set[str]
 
     def dump_stats(self) -> None:
         if self.options.dump_build_stats:
@@ -2498,6 +2502,8 @@ def module_not_found(manager: BuildManager, line: int, caller_state: State,
             if '{}' in note:
                 note = note.format(legacy_bundled_packages[target])
             errors.report(line, 0, note, severity='note', only_once=True, code=codes.IMPORT)
+        if reason is ModuleNotFoundReason.STUBS_NOT_INSTALLED:
+            manager.missing_stub_packages.add(legacy_bundled_packages[target])
     errors.set_import_context(save_import_context)
 
 
@@ -3223,3 +3229,23 @@ def topsort(data: Dict[AbstractSet[str],
                 for item, dep in data.items()
                 if item not in ready}
     assert not data, "A cyclic dependency exists amongst %r" % data
+
+
+def missing_stubs_file(cache_dir: str) -> str:
+    return os.path.join(cache_dir, 'missing_stubs')
+
+
+def record_missing_stub_packages(cache_dir: str, missing_stub_packages: Set[str]) -> None:
+    """Write a file containing missing stub packages.
+
+    This allows a subsequent "mypy --install-types" run (without other arguments)
+    to install missing stub packages.
+    """
+    fnam = missing_stubs_file(cache_dir)
+    if missing_stub_packages:
+        with open(fnam, 'w') as f:
+            for pkg in sorted(missing_stub_packages):
+                f.write('%s\n' % pkg)
+    else:
+        if os.path.isfile(fnam):
+            os.remove(fnam)
diff --git a/mypy/main.py b/mypy/main.py
@@ -66,6 +66,14 @@ def main(script_path: Optional[str],
     messages = []
     formatter = util.FancyFormatter(stdout, stderr, options.show_error_codes)
 
+    if options.install_types and (stdout is not sys.stdout or stderr is not sys.stderr):
+        # Since --install-types performs user input, we want regular stdout and stderr.
+        fail("--install-types not supported in this mode of running mypy", stderr, options)
+
+    if options.install_types and not sources:
+        install_types(options.cache_dir, formatter)
+        return
+
     def flush_errors(new_messages: List[str], serious: bool) -> None:
         if options.pretty:
             new_messages = formatter.fit_in_terminal(new_messages)
@@ -116,6 +124,11 @@ def flush_errors(new_messages: List[str], serious: bool) -> None:
             stdout.write(formatter.format_success(len(sources),
                                                   options.color_output) + '\n')
         stdout.flush()
+
+    if options.install_types:
+        install_types(options.cache_dir, formatter, after_run=True)
+        return
+
     if options.fast_exit:
         # Exit without freeing objects -- it's faster.
         #
@@ -722,6 +735,10 @@ def add_invertible_flag(flag: str,
         '--scripts-are-modules', action='store_true',
         help="Script x becomes module x instead of __main__")
 
+    add_invertible_flag('--install-types', default=False, strict_flag=False,
+                        help="Install detected missing library stub packages using pip",
+                        group=other_group)
+
     if server_options:
         # TODO: This flag is superfluous; remove after a short transition (2018-03-16)
         other_group.add_argument(
@@ -849,7 +866,7 @@ def set_strict_flags() -> None:
         code_methods = sum(bool(c) for c in [special_opts.modules + special_opts.packages,
                                              special_opts.command,
                                              special_opts.files])
-        if code_methods == 0:
+        if code_methods == 0 and not options.install_types:
             parser.error("Missing target module, package, files, or command.")
         elif code_methods > 1:
             parser.error("May only specify one of: module/package, files, or command.")
@@ -996,3 +1013,31 @@ def fail(msg: str, stderr: TextIO, options: Options) -> None:
     stderr.write('%s\n' % msg)
     maybe_write_junit_xml(0.0, serious=True, messages=[msg], options=options)
     sys.exit(2)
+
+
+def install_types(cache_dir: str,
+                  formatter: util.FancyFormatter,
+                  after_run: bool = False) -> None:
+    """Install stub packages using pip if some missing stubs were detected."""
+    if not os.path.isdir(cache_dir):
+        sys.stderr.write(
+            "Error: no mypy cache directory (you must enable incremental mode)\n")
+        sys.exit(2)
+    fnam = build.missing_stubs_file(cache_dir)
+    if not os.path.isfile(fnam):
+        # If there are no missing stubs, generate no output.
+        return
+    with open(fnam) as f:
+        packages = [line.strip() for line in f.readlines()]
+    if after_run:
+        print()
+    print('Installing missing stub packages:')
+    cmd = ['python3', '-m', 'pip', 'install'] + packages
+    print(formatter.style(' '.join(cmd), 'none', bold=True))
+    print()
+    x = input('Install? [yN] ')
+    if not x.strip() or not x.lower().startswith('y'):
+        print(formatter.style('mypy: Skipping installation', 'red', bold=True))
+        sys.exit(2)
+    print()
+    subprocess.run(cmd)
diff --git a/mypy/modulefinder.py b/mypy/modulefinder.py
@@ -62,7 +62,9 @@ def error_message_templates(self) -> Tuple[str, List[str]]:
             notes = [doc_link]
         elif self is ModuleNotFoundReason.STUBS_NOT_INSTALLED:
             msg = 'Library stubs not installed for "{}"'
-            notes = ['Hint: "python3 -m pip install {}"', doc_link]
+            notes = ['Hint: "python3 -m pip install {}"',
+                     '(or run "mypy --install-types" to install all missing stub packages)',
+                     doc_link]
         else:
             assert False
         return msg, notes
diff --git a/mypy/options.py b/mypy/options.py
@@ -268,6 +268,8 @@ def __init__(self) -> None:
         self.transform_source = None  # type: Optional[Callable[[Any], Any]]
         # Print full path to each file in the report.
         self.show_absolute_path = False  # type: bool
+        # Install missing stub packages if True
+        self.install_types = False
 
     # To avoid breaking plugin compatibility, keep providing new_semantic_analyzer
     @property
