Add support to MyPy for first-party plugins (#10755)

Similar to with Pylint plugins, we introduce a new target type `mypy_source_plugin`, along with a new option `--mypy-source-plugins`. 

The new target type is not strictly necessary; you could use a `python_library`. But, the new target type a) provides clearer modeling of what's going on, e.g. that this is a "root" rather than a typical library; and b) allows us to add instructions via `./pants target-types --details=mypy_source_plugin`.
 
[ci skip-build-wheels]
[ci skip-rust]

build-support/mypy/mypy.ini

@@ -1,8 +1,7 @@
 [mypy]
-# TODO(#10131): Re-enable this plugin once we add support to v2 MyPy.
 # See: https://mypy.readthedocs.io/en/latest/extending_mypy.html#configuring-mypy-to-use-plugins
-# plugins =
-#   pants.mypy.plugins.total_ordering
+plugins =
+  mypy_plugins.total_ordering
 
 # Refer to https://mypy.readthedocs.io/en/latest/command_line.html for the definition of each
 # of these options. MyPy is frequently updated, so this file should be periodically reviewed

src/python/pants/mypy/BUILD → pants-plugins/mypy_plugins/BUILD

@@ -1,4 +1,7 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-python_library()
+mypy_source_plugin(
+  name="total_ordering",
+  sources=["total_ordering.py"],
+)

pants-plugins/mypy_plugins/total_ordering.py

@@ -0,0 +1,47 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+# See: https://mypy.readthedocs.io/en/latest/extending_mypy.html#high-level-overview
+
+from typing import Callable, Optional, Type
+
+from mypy.nodes import ARG_POS, Argument, TypeInfo, Var
+from mypy.plugin import ClassDefContext, Plugin
+from mypy.plugins.common import add_method
+
+
+class TotalOrderingPlugin(Plugin):
+    """This inserts method type stubs for the "missing" ordering methods that the `@total_ordering`
+    class decorator will fill in dynamically."""
+
+    def get_class_decorator_hook(
+        self, fullname: str
+    ) -> Optional[Callable[[ClassDefContext], None]]:
+        return adjust_class_def if fullname == "functools.total_ordering" else None
+
+
+def adjust_class_def(class_def_context: ClassDefContext) -> None:
+    api = class_def_context.api
+    ordering_other_type = api.named_type("__builtins__.object")
+    ordering_return_type = api.named_type("__builtins__.bool")
+    arg = Argument(
+        variable=Var(name="other", type=ordering_other_type),
+        type_annotation=ordering_other_type,
+        initializer=None,
+        kind=ARG_POS,
+    )
+
+    type_info: TypeInfo = class_def_context.cls.info
+    for ordering_method_name in "__lt__", "__le__", "__gt__", "__ge__":
+        existing_method = type_info.get(ordering_method_name)
+        if existing_method is None:
+            add_method(
+                ctx=class_def_context,
+                name=ordering_method_name,
+                args=[arg],
+                return_type=ordering_return_type,
+            )
+
+
+def plugin(_version: str) -> Type[Plugin]:
+    return TotalOrderingPlugin

pants.toml

@@ -75,6 +75,7 @@ config = ["pyproject.toml", "examples/.isort.cfg"]
 
 [mypy]
 config = "build-support/mypy/mypy.ini"
+source_plugins = ["pants-plugins/mypy_plugins:total_ordering"]
 
 [pants-releases]
 release_notes = """

src/python/pants/backend/python/lint/pylint/plugin_target_type.py

@@ -7,7 +7,6 @@
 
 class PylintPluginSources(PythonSources):
     required = True
-    expected_file_extensions = (".py",)
 
 
 # NB: We solely subclass this to change the docstring.
@@ -26,7 +25,7 @@ class PylintSourcePlugin(Target):
     To load a source plugin:
 
         1. Write your plugin. See http://pylint.pycqa.org/en/latest/how_tos/plugins.html.
-        2. Define a `pylint_source_plugin` target with the plugin's Python files included in the
+        2. Define a `pylint_source_plugin` target with the plugin's Python file(s) included in the
             `sources` field.
         3. Add the parent directory of your target to the `root_patterns` option in the `[source]`
             scope. For example, if your plugin is at `build-support/pylint/custom_plugin.py`, add
@@ -49,10 +48,7 @@ class PylintSourcePlugin(Target):
     third-party dependencies or are located in the same directory or a subdirectory.
 
     Other targets can depend on this target. This allows you to write a `python_tests` target for
-    this code.
-
-    You can define the `provides` field to release this plugin as a distribution
-    (https://www.pantsbuild.org/docs/python-setup-py-goal).
+    this code or a `python_distribution` target to distribute the plugin externally.
     """
 
     alias = "pylint_source_plugin"

src/python/pants/backend/python/lint/pylint/rules.py

@@ -244,8 +244,7 @@ async def pylint_lint(
         return LintResults([], linter_name="Pylint")
 
     plugin_target_addresses = await MultiGet(
-        Get(Address, AddressInput, AddressInput.parse(plugin_addr))
-        for plugin_addr in pylint.source_plugins
+        Get(Address, AddressInput, plugin_addr) for plugin_addr in pylint.source_plugins
     )
     plugin_targets_request = Get(TransitiveTargets, Addresses(plugin_target_addresses))
     linted_targets_request = Get(

src/python/pants/backend/python/lint/pylint/rules_integration_test.py

@@ -333,6 +333,7 @@ def test_plugin(self):
 def test_source_plugin(rule_runner: RuleRunner) -> None:
     # NB: We make this source plugin fairly complex by having it use transitive dependencies.
     # This is to ensure that we can correctly support plugins with dependencies.
+    # The plugin bans `print()`.
     rule_runner.add_to_build_file(
         "",
         dedent(
@@ -350,53 +351,60 @@ def test_source_plugin(rule_runner: RuleRunner) -> None:
         ),
     )
     rule_runner.create_file(
-        "build-support/plugins/subdir/dep.py",
+        "pants-plugins/plugins/subdir/dep.py",
         dedent(
             """\
             from colors import red
 
             def is_print(node):
                 _ = red("Test that transitive deps are loaded.")
-                return node.func.name == "print"
+                return hasattr(node.func, "name") and node.func.name == "print"
             """
         ),
     )
     rule_runner.add_to_build_file(
-        "build-support/plugins/subdir", "python_library(dependencies=['//:colors'])"
+        "pants-plugins/plugins/subdir", "python_library(dependencies=['//:colors'])"
     )
     rule_runner.create_file(
-        "build-support/plugins/print_plugin.py",
+        "pants-plugins/plugins/print_plugin.py",
         dedent(
             """\
+            '''Docstring.'''
+
             from pylint.checkers import BaseChecker
             from pylint.interfaces import IAstroidChecker
 
             from subdir.dep import is_print
 
             class PrintChecker(BaseChecker):
+                '''Docstring.'''
+
                 __implements__ = IAstroidChecker
                 name = "print_plugin"
                 msgs = {
                     "C9871": ("`print` statements are banned", "print-statement-used", ""),
                 }
 
                 def visit_call(self, node):
+                    '''Docstring.'''
                     if is_print(node):
                         self.add_message("print-statement-used", node=node)
 
+
             def register(linter):
+                '''Docstring.'''
                 linter.register_checker(PrintChecker(linter))
             """
         ),
     )
     rule_runner.add_to_build_file(
-        "build-support/plugins",
+        "pants-plugins/plugins",
         dedent(
             """\
             pylint_source_plugin(
                 name='print_plugin',
                 sources=['print_plugin.py'],
-                dependencies=['//:pylint', 'build-support/plugins/subdir'],
+                dependencies=['//:pylint', 'pants-plugins/plugins/subdir'],
             )
             """
         ),
@@ -407,18 +415,31 @@ def register(linter):
         load-plugins=print_plugin
         """
     )
+
+    def run_pylint_with_plugin(tgt: Target) -> LintResult:
+        result = run_pylint(
+            rule_runner,
+            [tgt],
+            additional_args=[
+                "--pylint-source-plugins=['pants-plugins/plugins:print_plugin']",
+                f"--source-root-patterns=['pants-plugins/plugins', '{PACKAGE}']",
+            ],
+            config=config_content,
+        )
+        assert len(result) == 1
+        return result[0]
+
     target = make_target(
         rule_runner, [FileContent(f"{PACKAGE}/source_plugin.py", b"'''Docstring.'''\nprint()\n")]
     )
-    result = run_pylint(
-        rule_runner,
-        [target],
-        additional_args=[
-            "--pylint-source-plugins=['build-support/plugins:print_plugin']",
-            f"--source-root-patterns=['build-support/plugins', '{PACKAGE}']",
-        ],
-        config=config_content,
+    result = run_pylint_with_plugin(target)
+    assert result.exit_code == PYLINT_FAILURE_RETURN_CODE
+    assert f"{PACKAGE}/source_plugin.py:2:0: C9871" in result.stdout
+
+    # Ensure that running Pylint on the plugin itself still works.
+    plugin_tgt = rule_runner.get_target(
+        Address("pants-plugins/plugins", target_name="print_plugin")
     )
-    assert len(result) == 1
-    assert result[0].exit_code == PYLINT_FAILURE_RETURN_CODE
-    assert f"{PACKAGE}/source_plugin.py:2:0: C9871" in result[0].stdout
+    result = run_pylint_with_plugin(plugin_tgt)
+    assert result.exit_code == 0
+    assert "Your code has been rated at 10.00/10" in result.stdout

src/python/pants/backend/python/lint/pylint/subsystem.py

@@ -1,9 +1,10 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from typing import List, Optional, cast
+from typing import List, Optional, Tuple, cast
 
 from pants.backend.python.subsystems.python_tool_base import PythonToolBase
+from pants.engine.addresses import AddressInput
 from pants.option.custom_types import file_option, shell_str, target_option
 
 
@@ -66,5 +67,5 @@ def config(self) -> Optional[str]:
         return cast(Optional[str], self.options.config)
 
     @property
-    def source_plugins(self) -> List[str]:
-        return cast(List[str], self.options.source_plugins)
+    def source_plugins(self) -> Tuple[AddressInput, ...]:
+        return tuple(AddressInput.parse(v) for v in self.options.source_plugins)

src/python/pants/backend/python/typecheck/mypy/plugin_target_type.py

@@ -0,0 +1,43 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pants.backend.python.target_types import COMMON_PYTHON_FIELDS, PythonSources
+from pants.engine.target import Dependencies, Target
+
+
+class MyPyPluginSources(PythonSources):
+    required = True
+
+
+class MyPySourcePlugin(Target):
+    """A MyPy plugin loaded through source code.
+
+    To load a source plugin:
+
+        1. Write your plugin. See https://mypy.readthedocs.io/en/stable/extending_mypy.html.
+        2. Define a `mypy_source_plugin` target with the plugin's Python file(s) included in the
+            `sources` field.
+        3. Add `plugins = path.to.module` to your MyPy config file, using the name of the module
+            without source roots. For example, if your Python file is called
+            `pants-plugins/mypy_plugins/custom_plugin.py`, and you set `pants-plugins` as a source root,
+            then set `plugins = mypy_plugins.custom_plugin`. Set the `config`
+            option in the `[mypy]` scope to point to your MyPy config file.
+        5. Set the option `source_plugins` in the `[mypy]` scope to include this target's
+            address, e.g. `source_plugins = ["build-support/mypy_plugins:plugin"]`.
+
+    To instead load a third-party plugin, set the option `extra_requirements` in the `[mypy]`
+    scope (see https://www.pantsbuild.org/v2.0/docs/python-typecheck-goal). Set `plugins` in
+    your config file, like you'd do with a source plugin.
+
+    This target type is treated similarly to a `python_library` target. For example, Python linters
+    and formatters will run on this target.
+
+    You can include other targets in the `dependencies` field, including third-party requirements
+    and other source files (even if those source files live in a different directory).
+
+    Other targets can depend on this target. This allows you to write a `python_tests` target for
+    this code or a `python_distribution` target to distribute the plugin externally.
+    """
+
+    alias = "mypy_source_plugin"
+    core_fields = (*COMMON_PYTHON_FIELDS, Dependencies, MyPyPluginSources)

src/python/pants/backend/python/typecheck/mypy/register.py

@@ -8,6 +8,11 @@
 """
 
 from pants.backend.python.typecheck.mypy import rules as mypy_rules
+from pants.backend.python.typecheck.mypy.plugin_target_type import MyPySourcePlugin
+
+
+def target_types():
+    return [MyPySourcePlugin]
 
 
 def rules():