From abd319ab80d476f14e96918d5615fa0ac5b41d75 Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Wed, 28 Nov 2018 19:20:44 -0800
Subject: [PATCH] Provide an "API" for invoking the daemon (#5972)

Like the existing "API" for invoking mypy, it is just a thin wrapper
around calling the main function.
---
 mypy/api.py   | 25 +++++++++++++++++++++----
 mypy/dmypy.py | 14 ++++++++++----
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/mypy/api.py b/mypy/api.py
index 95dc0447e95f..16ee25349521 100644
--- a/mypy/api.py
+++ b/mypy/api.py
@@ -17,6 +17,12 @@
 
 Any pretty formatting is left to the caller.
 
+The 'run_dmypy' function is similar, but instead mimics invocation of
+dmypy.
+
+Note that these APIs don't support incremental generation of error
+messages.
+
 Trivial example of code using this module:
 
 import sys
@@ -33,15 +39,15 @@
     print(result[1])  # stderr
 
 print ('\nExit status:', result[2])
+
 """
 
 import sys
 from io import StringIO
-from typing import List, Tuple
-from mypy.main import main
+from typing import List, Tuple, Callable
 
 
-def run(args: List[str]) -> Tuple[str, str, int]:
+def _run(f: Callable[[], None]) -> Tuple[str, str, int]:
     old_stdout = sys.stdout
     new_stdout = StringIO()
     sys.stdout = new_stdout
@@ -51,7 +57,7 @@ def run(args: List[str]) -> Tuple[str, str, int]:
     sys.stderr = new_stderr
 
     try:
-        main(None, args=args)
+        f()
         exit_status = 0
     except SystemExit as system_exit:
         exit_status = system_exit.code
@@ -60,3 +66,14 @@ def run(args: List[str]) -> Tuple[str, str, int]:
         sys.stderr = old_stderr
 
     return new_stdout.getvalue(), new_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))
+
+
+def run_dmypy(args: List[str]) -> Tuple[str, str, int]:
+    from mypy.dmypy import main
+    return _run(lambda: main(args))
diff --git a/mypy/dmypy.py b/mypy/dmypy.py
index 84af22ed65bd..7b0ee2ab86b4 100644
--- a/mypy/dmypy.py
+++ b/mypy/dmypy.py
@@ -15,8 +15,9 @@
 import subprocess
 import sys
 import time
+import traceback
 
-from typing import Any, Callable, Dict, Mapping, Optional, Tuple
+from typing import Any, Callable, Dict, Mapping, Optional, Tuple, List
 
 from mypy.dmypy_util import STATUS_FILE, receive
 from mypy.ipc import IPCClient, IPCException
@@ -114,9 +115,9 @@ class BadStatus(Exception):
     pass
 
 
-def main() -> None:
+def main(argv: List[str]) -> None:
     """The code is top-down."""
-    args = parser.parse_args()
+    args = parser.parse_args(argv)
     if not args.action:
         parser.print_usage()
     else:
@@ -124,6 +125,11 @@ def main() -> None:
             args.action(args)
         except BadStatus as err:
             sys.exit(err.args[0])
+        except Exception:
+            # We do this explicitly to avoid exceptions percolating up
+            # through mypy.api invocations
+            traceback.print_exc()
+            sys.exit(2)
 
 
 ActionFunction = Callable[[argparse.Namespace], None]
@@ -483,4 +489,4 @@ def is_running() -> bool:
 # Run main().
 
 if __name__ == '__main__':
-    main()
+    main(sys.argv[1:])