Fix sys.stdout overriding in mypy.api

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

Author: elkhadiy

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(f: 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()
+        f(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: TextIO = sys.stdout,
+          stderr: TextIO = sys.stderr,
           ) -> BuildResult:
     """Analyze a program.
 
@@ -161,7 +163,7 @@ def default_flush_errors(new_messages: List[str], is_serious: bool) -> None:
     flush_errors = flush_errors or default_flush_errors
 
     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 +182,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 +201,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 +216,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 +325,7 @@ 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 = sys.stdout) -> Tuple[Plugin, Dict[str, str]]:
     """Load all configured plugins.
 
     Return a plugin that encapsulates all plugins chained together. Always
@@ -383,7 +389,7 @@ 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 +404,7 @@ 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 +502,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 +566,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 +912,7 @@ 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 = sys.stdout) -> 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 +926,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 +1777,7 @@ 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 +2437,7 @@ 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 = sys.stdout) -> Graph:
     log_configuration(manager)
 
     t0 = time.time()
@@ -2454,11 +2462,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 +2488,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 +2536,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 +2570,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/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
@@ -586,7 +586,7 @@ def remove_path_prefix(path: str, prefix: Optional[str]) -> str:
 
 
 def report_internal_error(err: Exception, file: Optional[str], line: int,
-                          errors: Errors, options: Options) -> None:
+                          errors: Errors, options: Options, stdout: TextIO = sys.stdout, stderr: TextIO = sys.stderr) -> None:
     """Report internal error and exit.
 
     This optionally starts pdb or shows a traceback.
@@ -597,7 +597,7 @@ def report_internal_error(err: Exception, file: Optional[str], line: int,
         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 +612,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 +627,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.