Add Config.cli_context() as a hook for custom CLI initialization and …

…cleanup logic. (#29)

CHANGES.rst

@@ -4,6 +4,7 @@ Changelog
 master
 ------
 
+* Add Config.cli_context() as a hook for custom CLI initialization and cleanup logic (#28; merge of #29).
 
 17.12.2
 -------

doc/configuration.rst

@@ -99,6 +99,34 @@ Subclassing ``Config`` or ``DefaultConfig``
 
     Defaults to ``False``.
 
+  .. method:: cli_context(command: str) -> Iterator[None]
+
+    A context manager which wraps the execution of the CLI command.
+
+    MonkeyType has to import your code in order to generate stubs for it. In
+    some cases, like if you're using Django, setup is required before your code
+    can be imported. Use this method to define the necessary setup or teardown
+    for your codebase.
+
+    This method must return a `context manager`_ instance. In most cases, the
+    simplest way to do this will be with the `contextlib.contextmanager`_
+    decorator. For example, if you run MonkeyType against a Django codebase,
+    you can setup Django before the command runs::
+
+      @contextmanager
+      def cli_context(self, command: str) -> Iterator[None]:
+          import django
+          django.setup()
+          yield
+
+    ``command`` is the name of the command passed to the monkeytype cli:
+    ``'run'``, ``'apply'``, etc.
+
+    The default implementation of this method returns a no-op context manager.
+
+    .. _context manager: https://docs.python.org/3/reference/datamodel.html#with-statement-context-managers
+    .. _contextlib.contextmanager: https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager
+
 .. class:: DefaultConfig()
 
   ``DefaultConfig`` is the config MonkeyType uses if you don't provide your own;

monkeytype/cli.py

@@ -184,7 +184,8 @@ def main(argv: List[str], stdout: IO, stderr: IO) -> int:
             "else monkeytype.config:DefaultConfig())"
         ),
     )
-    subparsers = parser.add_subparsers()
+
+    subparsers = parser.add_subparsers(title="commands", dest="command")
 
     run_parser = subparsers.add_parser(
         'run',
@@ -234,11 +235,15 @@ def main(argv: List[str], stdout: IO, stderr: IO) -> int:
 
     args = parser.parse_args(argv)
     update_args_from_config(args)
+
     handler = getattr(args, 'handler', None)
     if handler is None:
         parser.print_help(file=stderr)
         return 1
-    handler(args, stdout, stderr)
+
+    with args.config.cli_context(args.command):
+        handler(args, stdout, stderr)
+
     return 0
 
 

monkeytype/config.py

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+from contextlib import contextmanager
 import os
 import sys
 import sysconfig
@@ -12,7 +13,7 @@
     abstractmethod,
 )
 from types import CodeType
-from typing import Optional
+from typing import Optional, Iterator
 
 from monkeytype.db.base import (
     CallTraceStore,
@@ -39,6 +40,16 @@ def trace_store(self) -> CallTraceStore:
         """Return the CallTraceStore for storage/retrieval of call traces."""
         pass
 
+    @contextmanager
+    def cli_context(self, command: str) -> Iterator[None]:
+        """Lifecycle hook that is called once right after the CLI
+        starts.
+
+        `command` is the name of the command passed to monkeytype
+        ('run', 'apply', etc).
+        """
+        yield
+
     def trace_logger(self) -> CallTraceLogger:
         """Return the CallTraceLogger for logging call traces.
 

tests/test_cli.py

@@ -3,11 +3,13 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+from contextlib import contextmanager
 import io
 import os
 import pytest
 import sqlite3
 import tempfile
+from typing import Iterator
 
 from unittest import mock
 
@@ -30,6 +32,14 @@ def func2(a, b):
     pass
 
 
+class LoudContextConfig(DefaultConfig):
+    @contextmanager
+    def cli_context(self, command: str) -> Iterator[None]:
+        print(f"IN SETUP: {command}")
+        yield
+        print(f"IN TEARDOWN: {command}")
+
+
 @pytest.fixture
 def store_data():
     db_file = tempfile.NamedTemporaryFile(prefix='monkeytype_tests')
@@ -74,3 +84,11 @@ def test_no_traces(store_data, stdout, stderr):
     assert stderr.getvalue() == "No traces found\n"
     assert stdout.getvalue() == ''
     assert ret == 0
+
+
+def test_cli_context_manager_activated(capsys, stdout, stderr):
+    ret = cli.main(['-c', f'{__name__}:LoudContextConfig()', 'stub', 'some.module'], stdout, stderr)
+    out, err = capsys.readouterr()
+    assert out == "IN SETUP: stub\nIN TEARDOWN: stub\n"
+    assert err == ""
+    assert ret == 0