Fix sys.stdout overriding in mypy.api (#6750)

Overriding sys.stdout and sys.stderr in mypy.api is not threadsafe.
This causes problems sometimes when using the api in pyls for example.

This started with the changes made by @elkhadiy to fix #6125.

Fixes #6125.

Author: AnOctopus

mypy/__main__.py

@@ -1,11 +1,12 @@
 """Mypy type checker command line tool."""
 
+import sys
 from mypy.main import main
 
 
 def console_entry() -> None:
-    main(None)
+    main(None, sys.stdout, sys.stderr)
 
 
 if __name__ == '__main__':
-    main(None)
+    main(None, sys.stdout, sys.stderr)

mypy/api.py

@@ -44,36 +44,31 @@
 
 import sys
 from io import StringIO
-from typing import List, Tuple, Callable
+from typing import List, Tuple, Union, TextIO, Callable
+from mypy_extensions import DefaultArg
 
 
-def _run(f: Callable[[], None]) -> Tuple[str, str, int]:
-    old_stdout = sys.stdout
-    new_stdout = StringIO()
-    sys.stdout = new_stdout
+def _run(main_wrapper: Callable[[TextIO, TextIO], None]) -> Tuple[str, str, int]:
 
-    old_stderr = sys.stderr
-    new_stderr = StringIO()
-    sys.stderr = new_stderr
+    stdout = StringIO()
+    stderr = StringIO()
 
     try:
-        f()
+        main_wrapper(stdout, stderr)
         exit_status = 0
     except SystemExit as system_exit:
         exit_status = system_exit.code
-    finally:
-        sys.stdout = old_stdout
-        sys.stderr = old_stderr
 
-    return new_stdout.getvalue(), new_stderr.getvalue(), exit_status
+    return stdout.getvalue(), stderr.getvalue(), exit_status
 
 
 def run(args: List[str]) -> Tuple[str, str, int]:
     # Lazy import to avoid needing to import all of mypy to call run_dmypy
     from mypy.main import main
-    return _run(lambda: main(None, args=args))
+    return _run(lambda stdout, stderr: main(None, args=args,
+                                            stdout=stdout, stderr=stderr))
 
 
 def run_dmypy(args: List[str]) -> Tuple[str, str, int]:
     from mypy.dmypy import main
-    return _run(lambda: main(args))
+    return _run(lambda stdout, stderr: main(args))

mypy/build.py

@@ -25,7 +25,7 @@
 import types
 
 from typing import (AbstractSet, Any, Dict, Iterable, Iterator, List,
-                    Mapping, NamedTuple, Optional, Set, Tuple, Union, Callable)
+                    Mapping, NamedTuple, Optional, Set, Tuple, Union, Callable, TextIO)
 MYPY = False
 if MYPY:
     from typing import ClassVar
@@ -128,6 +128,8 @@ def build(sources: List[BuildSource],
           alt_lib_path: Optional[str] = None,
           flush_errors: Optional[Callable[[List[str], bool], None]] = None,
           fscache: Optional[FileSystemCache] = None,
+          stdout: Optional[TextIO] = None,
+          stderr: Optional[TextIO] = None,
           ) -> BuildResult:
     """Analyze a program.
 
@@ -159,9 +161,11 @@ def default_flush_errors(new_messages: List[str], is_serious: bool) -> None:
         messages.extend(new_messages)
 
     flush_errors = flush_errors or default_flush_errors
+    stdout = stdout or sys.stdout
+    stderr = stderr or sys.stderr
 
     try:
-        result = _build(sources, options, alt_lib_path, flush_errors, fscache)
+        result = _build(sources, options, alt_lib_path, flush_errors, fscache, stdout, stderr)
         result.errors = messages
         return result
     except CompileError as e:
@@ -180,6 +184,8 @@ def _build(sources: List[BuildSource],
            alt_lib_path: Optional[str],
            flush_errors: Callable[[List[str], bool], None],
            fscache: Optional[FileSystemCache],
+           stdout: TextIO,
+           stderr: TextIO,
            ) -> BuildResult:
     # This seems the most reasonable place to tune garbage collection.
     gc.set_threshold(150 * 1000)
@@ -197,7 +203,7 @@ def _build(sources: List[BuildSource],
 
     source_set = BuildSourceSet(sources)
     errors = Errors(options.show_error_context, options.show_column_numbers)
-    plugin, snapshot = load_plugins(options, errors)
+    plugin, snapshot = load_plugins(options, errors, stdout)
 
     # Construct a build manager object to hold state during the build.
     #
@@ -212,12 +218,14 @@ def _build(sources: List[BuildSource],
                            plugins_snapshot=snapshot,
                            errors=errors,
                            flush_errors=flush_errors,
-                           fscache=fscache)
+                           fscache=fscache,
+                           stdout=stdout,
+                           stderr=stderr)
     manager.trace(repr(options))
 
     reset_global_state()
     try:
-        graph = dispatch(sources, manager)
+        graph = dispatch(sources, manager, stdout)
         if not options.fine_grained_incremental:
             TypeState.reset_all_subtype_caches()
         return BuildResult(manager, graph)
@@ -319,7 +327,10 @@ def import_priority(imp: ImportBase, toplevel_priority: int) -> int:
     return toplevel_priority
 
 
-def load_plugins(options: Options, errors: Errors) -> Tuple[Plugin, Dict[str, str]]:
+def load_plugins(options: Options,
+                 errors: Errors,
+                 stdout: TextIO,
+                 ) -> Tuple[Plugin, Dict[str, str]]:
     """Load all configured plugins.
 
     Return a plugin that encapsulates all plugins chained together. Always
@@ -383,7 +394,8 @@ def plugin_error(message: str) -> None:
         try:
             plugin_type = getattr(module, func_name)(__version__)
         except Exception:
-            print('Error calling the plugin(version) entry point of {}\n'.format(plugin_path))
+            print('Error calling the plugin(version) entry point of {}\n'.format(plugin_path),
+                  file=stdout)
             raise  # Propagate to display traceback
 
         if not isinstance(plugin_type, type):
@@ -398,7 +410,8 @@ def plugin_error(message: str) -> None:
             custom_plugins.append(plugin_type(options))
             snapshot[module_name] = take_module_snapshot(module)
         except Exception:
-            print('Error constructing plugin instance of {}\n'.format(plugin_type.__name__))
+            print('Error constructing plugin instance of {}\n'.format(plugin_type.__name__),
+                  file=stdout)
             raise  # Propagate to display traceback
     # Custom plugins take precedence over the default plugin.
     return ChainedPlugin(options, custom_plugins + [default_plugin]), snapshot
@@ -496,8 +509,10 @@ def __init__(self, data_dir: str,
                  errors: Errors,
                  flush_errors: Callable[[List[str], bool], None],
                  fscache: FileSystemCache,
+                 stdout: TextIO,
+                 stderr: TextIO,
                  ) -> None:
-        super().__init__()
+        super().__init__(stdout, stderr)
         self.start_time = time.time()
         self.data_dir = data_dir
         self.errors = errors
@@ -558,7 +573,7 @@ def __init__(self, data_dir: str,
         self.plugin = plugin
         self.plugins_snapshot = plugins_snapshot
         self.old_plugins_snapshot = read_plugins_snapshot(self)
-        self.quickstart_state = read_quickstart_file(options)
+        self.quickstart_state = read_quickstart_file(options, self.stdout)
 
     def dump_stats(self) -> None:
         self.log("Stats:")
@@ -904,7 +919,9 @@ def read_plugins_snapshot(manager: BuildManager) -> Optional[Dict[str, str]]:
     return snapshot
 
 
-def read_quickstart_file(options: Options) -> Optional[Dict[str, Tuple[float, int, str]]]:
+def read_quickstart_file(options: Options,
+                         stdout: TextIO,
+                         ) -> Optional[Dict[str, Tuple[float, int, str]]]:
     quickstart = None  # type: Optional[Dict[str, Tuple[float, int, str]]]
     if options.quickstart_file:
         # This is very "best effort". If the file is missing or malformed,
@@ -918,7 +935,7 @@ def read_quickstart_file(options: Options) -> Optional[Dict[str, Tuple[float, in
             for file, (x, y, z) in raw_quickstart.items():
                 quickstart[file] = (x, y, z)
         except Exception as e:
-            print("Warning: Failed to load quickstart file: {}\n".format(str(e)))
+            print("Warning: Failed to load quickstart file: {}\n".format(str(e)), file=stdout)
     return quickstart
 
 
@@ -1769,7 +1786,8 @@ def wrap_context(self, check_blockers: bool = True) -> Iterator[None]:
         except CompileError:
             raise
         except Exception as err:
-            report_internal_error(err, self.path, 0, self.manager.errors, self.options)
+            report_internal_error(err, self.path, 0, self.manager.errors,
+                                  self.options, self.manager.stdout, self.manager.stderr)
         self.manager.errors.set_import_context(save_import_context)
         # TODO: Move this away once we've removed the old semantic analyzer?
         if check_blockers:
@@ -2429,7 +2447,10 @@ def log_configuration(manager: BuildManager) -> None:
 # The driver
 
 
-def dispatch(sources: List[BuildSource], manager: BuildManager) -> Graph:
+def dispatch(sources: List[BuildSource],
+             manager: BuildManager,
+             stdout: TextIO,
+             ) -> Graph:
     log_configuration(manager)
 
     t0 = time.time()
@@ -2454,11 +2475,11 @@ def dispatch(sources: List[BuildSource], manager: BuildManager) -> Graph:
                       fm_cache_size=len(manager.find_module_cache.results),
                       )
     if not graph:
-        print("Nothing to do?!")
+        print("Nothing to do?!", file=stdout)
         return graph
     manager.log("Loaded graph with %d nodes (%.3f sec)" % (len(graph), t1 - t0))
     if manager.options.dump_graph:
-        dump_graph(graph)
+        dump_graph(graph, stdout)
         return graph
 
     # Fine grained dependencies that didn't have an associated module in the build
@@ -2480,7 +2501,7 @@ def dispatch(sources: List[BuildSource], manager: BuildManager) -> Graph:
             manager.log("Error reading fine-grained dependencies cache -- aborting cache load")
             manager.cache_enabled = False
             manager.log("Falling back to full run -- reloading graph...")
-            return dispatch(sources, manager)
+            return dispatch(sources, manager, stdout)
 
     # If we are loading a fine-grained incremental mode cache, we
     # don't want to do a real incremental reprocess of the
@@ -2528,7 +2549,7 @@ def dumps(self) -> str:
                                                      json.dumps(self.deps))
 
 
-def dump_graph(graph: Graph) -> None:
+def dump_graph(graph: Graph, stdout: TextIO = sys.stdout) -> None:
     """Dump the graph as a JSON string to stdout.
 
     This copies some of the work by process_graph()
@@ -2562,7 +2583,7 @@ def dump_graph(graph: Graph) -> None:
                         if (dep_id != node.node_id and
                                 (dep_id not in node.deps or pri < node.deps[dep_id])):
                             node.deps[dep_id] = pri
-    print("[" + ",\n ".join(node.dumps() for node in nodes) + "\n]")
+    print("[" + ",\n ".join(node.dumps() for node in nodes) + "\n]", file=stdout)
 
 
 def load_graph(sources: List[BuildSource], manager: BuildManager,

mypy/dmypy_server.py

@@ -308,7 +308,9 @@ def cmd_run(self, version: str, args: Sequence[str]) -> Dict[str, object]:
             if self.fine_grained_manager:
                 manager = self.fine_grained_manager.manager
                 start_plugins_snapshot = manager.plugins_snapshot
-                _, current_plugins_snapshot = mypy.build.load_plugins(options, manager.errors)
+                _, current_plugins_snapshot = mypy.build.load_plugins(options,
+                                                                      manager.errors,
+                                                                      sys.stdout)
                 if current_plugins_snapshot != start_plugins_snapshot:
                     return {'restart': 'plugins changed'}
         except InvalidSourceList as err:

mypy/errors.py

@@ -3,7 +3,7 @@
 import traceback
 from collections import OrderedDict, defaultdict
 
-from typing import Tuple, List, TypeVar, Set, Dict, Optional
+from typing import Tuple, List, TypeVar, Set, Dict, Optional, TextIO
 
 from mypy.scope import Scope
 from mypy.options import Options
@@ -585,19 +585,27 @@ def remove_path_prefix(path: str, prefix: Optional[str]) -> str:
         return path
 
 
-def report_internal_error(err: Exception, file: Optional[str], line: int,
-                          errors: Errors, options: Options) -> None:
+def report_internal_error(err: Exception,
+                          file: Optional[str],
+                          line: int,
+                          errors: Errors,
+                          options: Options,
+                          stdout: Optional[TextIO] = None,
+                          stderr: Optional[TextIO] = None,
+                          ) -> None:
     """Report internal error and exit.
 
     This optionally starts pdb or shows a traceback.
     """
+    stdout = (stdout or sys.stdout)
+    stderr = (stderr or sys.stderr)
     # Dump out errors so far, they often provide a clue.
     # But catch unexpected errors rendering them.
     try:
         for msg in errors.new_messages():
             print(msg)
     except Exception as e:
-        print("Failed to dump errors:", repr(e), file=sys.stderr)
+        print("Failed to dump errors:", repr(e), file=stderr)
 
     # Compute file:line prefix for official-looking error messages.
     if file:
@@ -612,11 +620,11 @@ def report_internal_error(err: Exception, file: Optional[str], line: int,
     print('{}error: INTERNAL ERROR --'.format(prefix),
           'please report a bug at https://github.com/python/mypy/issues',
           'version: {}'.format(mypy_version),
-          file=sys.stderr)
+          file=stderr)
 
     # If requested, drop into pdb. This overrides show_tb.
     if options.pdb:
-        print('Dropping into pdb', file=sys.stderr)
+        print('Dropping into pdb', file=stderr)
         import pdb
         pdb.post_mortem(sys.exc_info()[2])
 
@@ -627,15 +635,15 @@ def report_internal_error(err: Exception, file: Optional[str], line: int,
         if not options.pdb:
             print('{}: note: please use --show-traceback to print a traceback '
                   'when reporting a bug'.format(prefix),
-                  file=sys.stderr)
+                  file=stderr)
     else:
         tb = traceback.extract_stack()[:-2]
         tb2 = traceback.extract_tb(sys.exc_info()[2])
         print('Traceback (most recent call last):')
         for s in traceback.format_list(tb + tb2):
             print(s.rstrip('\n'))
-        print('{}: {}'.format(type(err).__name__, err))
-        print('{}: note: use --pdb to drop into pdb'.format(prefix), file=sys.stderr)
+        print('{}: {}'.format(type(err).__name__, err), file=stdout)
+        print('{}: note: use --pdb to drop into pdb'.format(prefix), file=stderr)
 
     # Exit.  The caller has nothing more to say.
     # We use exit code 2 to signal that this is no ordinary error.