make extension of the plugin more ergonomic (#4)

pytest.ini

@@ -1,4 +1,5 @@
 [pytest]
 python_files = test_*.py
 addopts =
-    -s
+    -s
+    --mypy-extension-hook pytest_mypy.tests.reveal_type_hook.hook

pytest_mypy/collect.py

@@ -1,12 +1,10 @@
-from typing import Any, Dict, List, TYPE_CHECKING, Type
+from typing import Any, Dict, List
 
 import pytest
 import yaml
 from _pytest.config.argparsing import Parser
 
 from pytest_mypy import utils
-if TYPE_CHECKING:
-    from pytest_mypy.item import YamlTestItem
 
 
 class File:
@@ -44,11 +42,9 @@ def construct_mapping(self, node, deep=False):
 
 
 class YamlTestFile(pytest.File):
-    def get_test_class(self) -> 'Type[YamlTestItem]':
+    def collect(self):
         from pytest_mypy.item import YamlTestItem
-        return YamlTestItem
 
-    def collect(self):
         parsed_file = yaml.load(stream=self.fspath.read_text('utf8'), Loader=SafeLineLoader)
         if parsed_file is None:
             return
@@ -76,15 +72,15 @@ def collect(self):
             disable_cache = raw_test.get('disable_cache', False)
             expected_output_lines = raw_test.get('out', '').split('\n')
 
-            yield self.get_test_class()(name=test_name,
-                                        collector=self,
-                                        config=self.config,
-                                        files=test_files,
-                                        starting_lineno=starting_lineno,
-                                        environment_variables=extra_environment_variables,
-                                        disable_cache=disable_cache,
-                                        expected_output_lines=output_from_comments + expected_output_lines,
-                                        parsed_test_data=raw_test)
+            yield YamlTestItem(name=test_name,
+                               collector=self,
+                               config=self.config,
+                               files=test_files,
+                               starting_lineno=starting_lineno,
+                               environment_variables=extra_environment_variables,
+                               disable_cache=disable_cache,
+                               expected_output_lines=output_from_comments + expected_output_lines,
+                               parsed_test_data=raw_test)
 
 
 def pytest_collect_file(path, parent):
@@ -100,3 +96,6 @@ def pytest_addoption(parser: Parser) -> None:
                     help='Which .ini file to use as a default config for tests')
     group.addoption('--mypy-same-process', action='store_true',
                     help='Run in the same process. Useful for debugging, will create problems with import cache')
+    group.addoption('--mypy-extension-hook', type=str,
+                    help='Fully qualifield path to the extension hook function, in case you need custom yaml keys. '
+                         'Has to be top-level.')

pytest_mypy/item.py

@@ -1,9 +1,10 @@
+import importlib
 import os
 import subprocess
 import sys
 import tempfile
 from pathlib import Path
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, Callable
 
 import capturer
 import pytest
@@ -152,9 +153,6 @@ def find_dependent_paths(self, path: Path) -> List[Path]:
                     dependants.append(path.with_suffix('').with_suffix(''))
         return dependants
 
-    def custom_init(self):
-        pass
-
     def typecheck_in_new_subprocess(self, execution_path: Path, mypy_cmd_options: List[Any]) -> Tuple[int, str]:
         import distutils.spawn
         mypy_executable = distutils.spawn.find_executable('mypy')
@@ -163,10 +161,12 @@ def typecheck_in_new_subprocess(self, execution_path: Path, mypy_cmd_options: Li
         # add current directory to path
         self.environment_variables['PYTHONPATH'] = str(execution_path)
         completed = subprocess.run([mypy_executable, *mypy_cmd_options],
-                                   stdout=subprocess.PIPE, cwd=os.getcwd(),
+                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                                   cwd=os.getcwd(),
                                    env=self.environment_variables)
         captured_stdout = completed.stdout.decode()
-        return completed.returncode, captured_stdout
+        captured_stderr = completed.stderr.decode()
+        return completed.returncode, captured_stdout + captured_stderr
 
     def typecheck_in_same_process(self, execution_path: Path, mypy_cmd_options: List[Any]) -> Tuple[int, str]:
             with utils.temp_environ(), utils.temp_path(), utils.temp_sys_modules():
@@ -182,6 +182,13 @@ def typecheck_in_same_process(self, execution_path: Path, mypy_cmd_options: List
 
                 return return_code, captured_std_streams.get_text()
 
+    def execute_extension_hook(self) -> None:
+        extension_hook_fqname = self.config.option.mypy_extension_hook
+        module_name, func_name = extension_hook_fqname.rsplit('.', maxsplit=1)
+        module = importlib.import_module(module_name)
+        extension_hook = getattr(module, func_name)
+        extension_hook(self)
+
     def runtest(self):
         try:
             temp_dir = tempfile.TemporaryDirectory(prefix='pytest-mypy-', dir=self.root_directory)
@@ -196,7 +203,8 @@ def runtest(self):
                 mypy_cmd_options.append(main_file)
 
                 # extension point for derived packages
-                self.custom_init()
+                if hasattr(self.config.option, 'mypy_extension_hook'):
+                    self.execute_extension_hook()
 
                 # make files
                 self.make_test_files_in_current_directory()

pytest_mypy/tests/reveal_type_hook.py

@@ -0,0 +1,10 @@
+from pytest_mypy.item import YamlTestItem
+
+
+def hook(item: YamlTestItem) -> None:
+    parsed_test_data = item.parsed_test_data
+    obj_to_reveal = parsed_test_data.get('reveal_type')
+    if obj_to_reveal:
+        for file in item.files:
+            if file.path.endswith('main.py'):
+                file.content = f'reveal_type({obj_to_reveal})'

pytest_mypy/tests/test-extension.yml

@@ -0,0 +1,7 @@
+-   case: reveal_type_extension_is_loaded
+    main: |
+        # if hook works, main should contain 'reveal_type(1)'
+    reveal_type: 1
+    out: |
+        main:1: note: Revealed type is 'builtins.int'
+