Allow use of entry-points like strings in mypy.ini to register plugins (#5358)

Fixes #3916

Author: chadrik

mypy/build.py

@@ -563,6 +563,7 @@ def load_plugins(options: Options, errors: Errors) -> Plugin:
     Return a plugin that encapsulates all plugins chained together. Always
     at least include the default plugin (it's last in the chain).
     """
+    import importlib
 
     default_plugin = DefaultPlugin(options)  # type: Plugin
     if not options.config_file:
@@ -579,34 +580,44 @@ def plugin_error(message: str) -> None:
     custom_plugins = []  # type: List[Plugin]
     errors.set_file(options.config_file, None)
     for plugin_path in options.plugins:
-        # Plugin paths are relative to the config file location.
-        plugin_path = os.path.join(os.path.dirname(options.config_file), plugin_path)
-
-        if not os.path.isfile(plugin_path):
-            plugin_error("Can't find plugin '{}'".format(plugin_path))
-        plugin_dir = os.path.dirname(plugin_path)
-        fnam = os.path.basename(plugin_path)
-        if not fnam.endswith('.py'):
+        func_name = 'plugin'
+        plugin_dir = None  # type: Optional[str]
+        if ':' in os.path.basename(plugin_path):
+            plugin_path, func_name = plugin_path.rsplit(':', 1)
+        if plugin_path.endswith('.py'):
+            # Plugin paths can be relative to the config file location.
+            plugin_path = os.path.join(os.path.dirname(options.config_file), plugin_path)
+            if not os.path.isfile(plugin_path):
+                plugin_error("Can't find plugin '{}'".format(plugin_path))
+            plugin_dir = os.path.dirname(plugin_path)
+            fnam = os.path.basename(plugin_path)
+            module_name = fnam[:-3]
+            sys.path.insert(0, plugin_dir)
+        elif re.search(r'[\\/]', plugin_path):
+            fnam = os.path.basename(plugin_path)
             plugin_error("Plugin '{}' does not have a .py extension".format(fnam))
-        module_name = fnam[:-3]
-        import importlib
-        sys.path.insert(0, plugin_dir)
+        else:
+            module_name = plugin_path
+
         try:
-            m = importlib.import_module(module_name)
+            module = importlib.import_module(module_name)
         except Exception:
-            print('Error importing plugin {}\n'.format(plugin_path))
-            raise  # Propagate to display traceback
+            plugin_error("Error importing plugin '{}'".format(plugin_path))
         finally:
-            assert sys.path[0] == plugin_dir
-            del sys.path[0]
-        if not hasattr(m, 'plugin'):
-            plugin_error('Plugin \'{}\' does not define entry point function "plugin"'.format(
-                plugin_path))
+            if plugin_dir is not None:
+                assert sys.path[0] == plugin_dir
+                del sys.path[0]
+
+        if not hasattr(module, func_name):
+            plugin_error('Plugin \'{}\' does not define entry point function "{}"'.format(
+                plugin_path, func_name))
+
         try:
-            plugin_type = getattr(m, 'plugin')(__version__)
+            plugin_type = getattr(module, func_name)(__version__)
         except Exception:
             print('Error calling the plugin(version) entry point of {}\n'.format(plugin_path))
             raise  # Propagate to display traceback
+
         if not isinstance(plugin_type, type):
             plugin_error(
                 'Type object expected as the return value of "plugin"; got {!r} (in {})'.format(

mypy/test/testcheck.py

@@ -8,7 +8,7 @@
 
 from mypy import build
 from mypy.build import BuildSource, Graph, SearchPaths
-from mypy.test.config import test_temp_dir
+from mypy.test.config import test_temp_dir, test_data_prefix
 from mypy.test.data import DataDrivenTestCase, DataSuite, FileOperation, UpdateFile
 from mypy.test.helpers import (
     assert_string_arrays_equal, normalize_error_messages, assert_module_equivalence,
@@ -152,6 +152,9 @@ def run_case_once(self, testcase: DataDrivenTestCase,
             sources.append(BuildSource(program_path, module_name,
                                        None if incremental_step else program_text))
 
+        plugin_dir = os.path.join(test_data_prefix, 'plugins')
+        sys.path.insert(0, plugin_dir)
+
         res = None
         try:
             res = build.build(sources=sources,
@@ -160,6 +163,10 @@ def run_case_once(self, testcase: DataDrivenTestCase,
             a = res.errors
         except CompileError as e:
             a = e.messages
+        finally:
+            assert sys.path[0] == plugin_dir
+            del sys.path[0]
+
         a = normalize_error_messages(a)
 
         # Make sure error messages match

test-data/unit/check-custom-plugin.test

@@ -3,14 +3,22 @@
 -- Note: Plugins used by tests live under test-data/unit/plugins. Defining
 --       plugin files in test cases does not work reliably.
 
-[case testFunctionPlugin]
+[case testFunctionPluginFile]
 # flags: --config-file tmp/mypy.ini
 def f() -> str: ...
 reveal_type(f())  # E: Revealed type is 'builtins.int'
 [file mypy.ini]
 [[mypy]
 plugins=<ROOT>/test-data/unit/plugins/fnplugin.py
 
+[case testFunctionPlugin]
+# flags: --config-file tmp/mypy.ini
+def f() -> str: ...
+reveal_type(f())  # E: Revealed type is 'builtins.int'
+[file mypy.ini]
+[[mypy]
+plugins=fnplugin
+
 [case testFunctionPluginFullnameIsNotNone]
 # flags: --config-file tmp/mypy.ini
 from typing import Callable, TypeVar
@@ -35,7 +43,19 @@ reveal_type(h())  # E: Revealed type is 'Any'
 plugins=<ROOT>/test-data/unit/plugins/fnplugin.py,
   <ROOT>/test-data/unit/plugins/plugin2.py
 
-[case testMissingPlugin]
+[case testTwoPluginsMixedType]
+# flags: --config-file tmp/mypy.ini
+def f(): ...
+def g(): ...
+def h(): ...
+reveal_type(f())  # E: Revealed type is 'builtins.int'
+reveal_type(g())  # E: Revealed type is 'builtins.str'
+reveal_type(h())  # E: Revealed type is 'Any'
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/fnplugin.py, plugin2
+
+[case testMissingPluginFile]
 # flags: --config-file tmp/mypy.ini
 [file mypy.ini]
 [[mypy]
@@ -44,6 +64,15 @@ plugins=missing.py
 tmp/mypy.ini:2: error: Can't find plugin 'tmp/missing.py'
 --' (work around syntax highlighting)
 
+[case testMissingPlugin]
+# flags: --config-file tmp/mypy.ini
+[file mypy.ini]
+[[mypy]
+plugins=missing
+[out]
+tmp/mypy.ini:2: error: Error importing plugin 'missing'
+--' (work around syntax highlighting)
+
 [case testMultipleSectionsDefinePlugin]
 # flags: --config-file tmp/mypy.ini
 [file mypy.ini]
@@ -74,6 +103,22 @@ tmp/mypy.ini:2: error: Plugin 'badext.pyi' does not have a .py extension
 [out]
 tmp/mypy.ini:2: error: Plugin '<ROOT>/test-data/unit/plugins/noentry.py' does not define entry point function "plugin"
 
+[case testCustomPluginEntryPointFile]
+# flags: --config-file tmp/mypy.ini
+def f() -> str: ...
+reveal_type(f())  # E: Revealed type is 'builtins.int'
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/customentry.py:register
+
+[case testCustomPluginEntryPoint]
+# flags: --config-file tmp/mypy.ini
+def f() -> str: ...
+reveal_type(f())  # E: Revealed type is 'builtins.int'
+[file mypy.ini]
+[[mypy]
+plugins=customentry:register
+
 [case testInvalidPluginEntryPointReturnValue]
 # flags: --config-file tmp/mypy.ini
 def f(): pass

test-data/unit/plugins/customentry.py

@@ -0,0 +1,14 @@
+from mypy.plugin import Plugin
+
+class MyPlugin(Plugin):
+    def get_function_hook(self, fullname):
+        if fullname == '__main__.f':
+            return my_hook
+        assert fullname is not None
+        return None
+
+def my_hook(ctx):
+    return ctx.api.named_generic_type('builtins.int', [])
+
+def register(version):
+    return MyPlugin