Add mypy plugin for NewType

Author: selimb

.gitignore

@@ -1,3 +1,6 @@
+# Auto-generated from tests
+test-output
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

marshmallow_dataclass/mypy.py

@@ -0,0 +1,65 @@
+import inspect
+from typing import Any, Callable, Optional, Type, Optional
+
+from mypy import nodes
+from mypy.plugin import DynamicClassDefContext, Plugin
+
+import marshmallow_dataclass
+
+_NEW_TYPE_SIG = inspect.signature(marshmallow_dataclass.NewType)
+
+
+def plugin(version: str) -> Type[Plugin]:
+    return MarshmallowDataclassPlugin
+
+
+class MarshmallowDataclassPlugin(Plugin):
+    def get_dynamic_class_hook(
+        self, fullname: str
+    ) -> Optional[Callable[[DynamicClassDefContext], None]]:
+        if fullname == "marshmallow_dataclass.NewType":
+            return new_type_hook
+        return None
+
+
+def new_type_hook(ctx: DynamicClassDefContext) -> None:
+    """
+    Dynamic class hook for :func:`marshmallow_dataclass.NewType`.
+
+    Uses the type of the ``typ`` argument.
+    """
+    typ = _get_arg_by_name(ctx.call, "typ", _NEW_TYPE_SIG)
+    if not isinstance(typ, nodes.RefExpr):
+        return
+    info = typ.node
+    if not isinstance(info, nodes.TypeInfo):
+        return
+    ctx.api.add_symbol_table_node(ctx.name, nodes.SymbolTableNode(nodes.GDEF, info))
+
+
+def _get_arg_by_name(
+    call: nodes.CallExpr, name: str, sig: inspect.Signature
+) -> Optional[nodes.Expression]:
+    """
+    Get value of argument from a call.
+
+    :return: The argument value, or ``None`` if it cannot be found.
+
+    .. warning::
+        This probably doesn't yet work for calls with ``*args`` and/or ``*kwargs``.
+    """
+    args = []
+    kwargs = {}
+    for arg_name, arg_value in zip(call.arg_names, call.args):
+        if arg_name is None:
+            args.append(arg_value)
+        else:
+            kwargs[arg_name] = arg_value
+    try:
+        bound_args = sig.bind(*args, **kwargs)
+    except TypeError:
+        return None
+    try:
+        return bound_args.arguments[name]
+    except KeyError:
+        return None

tests/test_mypy.py

@@ -0,0 +1,85 @@
+import shutil
+import textwrap
+import os
+import unittest
+
+import mypy.api
+
+HERE = os.path.dirname(__file__)
+TEST_OUTPUT_DIR = os.path.join(HERE, "test-output")
+MYPY_INI = """\
+[mypy]
+follow_imports = silent
+plugins = marshmallow_dataclass.mypy
+"""
+
+
+class TestMypyPlugin(unittest.TestCase):
+    maxDiff = None
+
+    def setUp(self):
+        """
+        Prepare a clean test directory at tests/test-output/test_mypy-{testname}.
+        Also cd into it for the duration of the test to get simple filenames in mypy output.
+        """
+        testname = self.id().split(".")[-1]
+        self.testdir = os.path.join(TEST_OUTPUT_DIR, f"test_mypy-{testname}")
+        if os.path.exists(self.testdir):
+            shutil.rmtree(self.testdir)
+        os.makedirs(self.testdir)
+        self.old_cwd = os.getcwd()
+        os.chdir(self.testdir)
+
+    def tearDown(self):
+        os.chdir(self.old_cwd)
+
+    def mypy_test(self, contents: str, expected: str):
+        """
+        Run mypy and assert output matches ``expected``.
+
+        The file with ``contents`` is always named ``main.py``.
+        """
+        config_path = os.path.join(self.testdir, "mypy.ini")
+        script_path = os.path.join(self.testdir, "main.py")
+        with open(config_path, "w") as f:
+            f.write(MYPY_INI)
+        with open(script_path, "w") as f:
+            f.write(textwrap.dedent(contents).strip())
+        command = [script_path, "--config-file", config_path, "--no-error-summary"]
+        out, err, returncode = mypy.api.run(command)
+        err_msg = "\n".join(["", f"returncode: {returncode}", "stdout:", out, "stderr:", err])
+        self.assertEqual(out.strip(), textwrap.dedent(expected).strip(), err_msg)
+
+    def test_basic(self):
+        self.mypy_test(
+            """
+            from dataclasses import dataclass
+            import marshmallow as ma
+            from marshmallow_dataclass import NewType
+
+            Email = NewType("Email", str, validate=ma.validate.Email)
+            UserID = NewType("UserID", validate=ma.validate.Length(equal=32), typ=str)
+
+            @dataclass
+            class User:
+                id: UserID
+                email: Email
+
+            user = User(id="a"*32, email="[email protected]")
+            reveal_type(user.id)
+            reveal_type(user.email)
+
+            User(id=42, email="[email protected]")
+            User(id="a"*32, email=["not", "a", "string"])
+            """,
+            """
+            main.py:14: note: Revealed type is 'builtins.str'
+            main.py:15: note: Revealed type is 'builtins.str'
+            main.py:17: error: Argument "id" to "User" has incompatible type "int"; expected "str"
+            main.py:18: error: Argument "email" to "User" has incompatible type "List[str]"; expected "str"
+            """,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()