Add mypy plugin for NewType

.travis.yml

@@ -5,7 +5,7 @@ dist: xenial
 install:
   - pip install --pre -e '.[dev]'
 script:
-  - python -m unittest
+  - pytest
 
 stages:
   - pre-commit hooks
@@ -53,4 +53,4 @@ matrix:
             secure: "WhUEch0YOZRtqr+r6BAd2KfgMY2IctHdoXNbl60qbyURjX3RzmICxB6GFHlAPkzu+b7GkY/WxKrAEa9NjjkHZ69M+x6ozORW4+YzicNVIgqxQkE65Su6l8zOCSexncMJrtmYceoxaRcX2Bsjy9r5fAv0gO6+WugOCAV6sZHkANqIjySTN82vAnla9/htWcesvw3JFckhdsPH/Lnu+o/hHKeM33vQUsJThQ73fatNuObyzglaoNu2L1PSjitI3wBA3jOpAL2o5v7bD/P7Krpu+YsGlaJNMHfHWkcNxypikZitCKcocoI7dWgQvSTqWA6TPV/MboRRF/2ZhnMKuw7ZY84GBx8Dyyw/len9s/1+m8kgu9ES1qJaIf03fvC60Uend2n2yk3sx6+rG1CCAvf21CJgMW+6XAqKpLI2tj0nPdBHvLGSdeRtTO5Ubpw8hDModlVai0RLhtYNkjLQkl+QrEehm2x/zhrP7gCudXQsKx8jvHHF70/pCN/WlvUhR42HEqWsPg7ZahkWi7YkFVY7e846ScxbWKjFL1s6vyCZso5U9DdUBvVEinshLoSZSbXotF4oHqdyKf17Hk0P6dPn/ocxGRB5rWWn5irhtEqcxZFKgawO8N1GYnTF64YewvO3O5gNqXrbpwCSScOrvDKDRosAWQ0yB/71ZvUEEcaYiFw="
           on:
             tags: true
-            all_branches: true
+            all_branches: true

README.md

@@ -216,6 +216,16 @@ from marshmallow_dataclass import NewType
 Email = NewType("Email", str, field=marshmallow.fields.Email)
 ```
 
+Note: if you are using `mypy`, you will notice that `mypy` throws an error if a variable defined with
+`NewType` is used in a type annotation. To resolve this, add the `marshmallow_dataclass.mypy` plugin 
+to your `mypy` configuration, e.g.:
+
+```ini
+[mypy]
+plugins = marshmallow_dataclass.mypy
+# ...
+```
+
 ### `Meta` options
 
 [`Meta` options](https://marshmallow.readthedocs.io/en/stable/api_reference.html#marshmallow.Schema.Meta) are set the same way as a marshmallow `Schema`.

marshmallow_dataclass/mypy.py

@@ -0,0 +1,65 @@
+import inspect
+from typing import Callable, Optional, Type
+
+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

setup.py

@@ -18,12 +18,21 @@
     ':python_version == "3.6"': ["dataclasses"],
     "lint": ["pre-commit~=1.18"],
     "docs": ["sphinx"],
+    "tests": [
+        "pytest",
+        # re: pypy: typed-ast (a dependency of mypy) fails to install on pypy
+        # https://github.com/python/typed_ast/issues/111
+        # re: win32: pytest-mypy-plugins depends on capturer, which isn't supported on
+        # windows
+        "pytest-mypy-plugins; implementation_name != 'pypy' and sys.platform != 'win32'",
+    ],
 }
 EXTRAS_REQUIRE["dev"] = (
     EXTRAS_REQUIRE["enum"]
     + EXTRAS_REQUIRE["union"]
     + EXTRAS_REQUIRE["lint"]
     + EXTRAS_REQUIRE["docs"]
+    + EXTRAS_REQUIRE["tests"]
 )
 
 setup(

tests/README.md

@@ -2,7 +2,7 @@
 
 ## How to run
  In the project root, install dependencies and then run:
- 
+
  ```bash
-python3 -m unittest
-```
+pytest
+```

tests/test_mypy.yml

@@ -0,0 +1,26 @@
+# Tests for marshmallow_dataclass.mypy, using pytest-mypy-plugins
+# NOTE: Since pytest-mypy-plugins is not installed on pypy (see setup.py), these tests
+#       will not even be collected by pytest
+- case: NewType_basic
+  mypy_config: |
+    follow_imports = silent
+    plugins = marshmallow_dataclass.mypy
+  main: |
+    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)  # N: Revealed type is 'builtins.str'
+    reveal_type(user.email)  # N: Revealed type is 'builtins.str'
+
+    User(id=42, email="[email protected]")  # E: Argument "id" to "User" has incompatible type "int"; expected "str"
+    User(id="a"*32, email=["not", "a", "string"])  # E: Argument "email" to "User" has incompatible type "List[str]"; expected "str"