rework a bunch of things

.gitignore

@@ -11,3 +11,6 @@
 .coverage
 .mypy_cache
 .pytest_cache
+build/
+dist/
+venv/

.pre-commit-config.yaml

@@ -14,43 +14,20 @@ repos:
           - id: check-yaml
           - id: debug-statements
           - id: end-of-file-fixer
-          - id: flake8
-            additional_dependencies:
-                - flake8>=3.6.0,<4
-                - flake8-bugbear
-                - flake8-builtins
-                - flake8-comprehensions
-                - flake8-commas
           - id: trailing-whitespace
-    - repo: https://github.com/asottile/yesqa
-      rev: v0.0.8
+    - repo: https://gitlab.com/pycqa/flake8
+      rev: 3.8.1
       hooks:
-          - id: yesqa
+          - id: flake8
             additional_dependencies:
-                - flake8>=3.6.0,<4
-                - flake8-bugbear
-                - flake8-builtins
-                - flake8-comprehensions
-                - flake8-commas
+            -   flake8-bugbear==18.8.0
+            -   flake8-comprehensions==1.4.1
+            -   flake8-tidy-imports==1.1.0
     - repo: https://github.com/asottile/pyupgrade
       rev: v1.11.0
       hooks:
         - id: pyupgrade
-    # Switch to standard pre-commit mypy when a version of mypy is released that has:
-    # - mypy.plugin.Plugin.lookup_fully_qualified
-    # - typeshed with https://github.com/ikonst/typeshed/tree/pynamodb-attr-nullable
-    - repo: local
+    - repo: https://github.com/pre-commit/mirrors-mypy
+      rev: v0.770
       hooks:
         - id: mypy
-          name: mypy
-          entry: mypy
-          language: python
-          'types': [python]
-          args: ["--ignore-missing-imports", "--scripts-are-modules", "--show-traceback"]
-          additional_dependencies: [
-            '-U', 'git+git://github.com/ikonst/mypy.git@5a8dffb5bd94bda615703f6994d39ea1c7c02ef5',
-          ]
-          exclude: >
-            (?x)^(
-              tests/.*|
-            )$

.travis.yml

@@ -1,7 +1,7 @@
 sudo: false
 language: python
 python:
-  - '3.6'
+  - '3.8'
 cache:
   directories:
     - $HOME/.cache/pip
@@ -18,8 +18,8 @@ jobs:
         provider: pypi
         user: ikonst
         password:
-          secure: "TlSgflmwfHh6pc03DuL9FOpieuDr2E2e4iaOtCnSa5kWdMldxOWVVUmjWJ0xDrL6nS+Hr5ZygAJzppVbiCvWiOC8curLu+Pl2eLc+Mf9WG9Mw1G2DiU8ci+uJdYCE6Dpak8OKG9lUu0O3XBDy+oZIEEaiWVdjBrBVplK0le8pFNB6fNu4Nru2JOGunHyQUfDmD+80m5tm8Mo32Y0Xryi30m7yRrLBlqn7eu2YnRcuhcY904edtuL36smSWNMl4rc6+IhxcQil11TlQD+DoyGlqxohYhaELgFs/2afYdJJJ/zq4qQ2ZJ8c6qwnKWxclLeoC8pGmDuhJo9zmyHuAJWeU+PZ1/ldyTjC1ttinJcbCZNlqZI7cz/slzxKwmcuTTXHRbnxS9VYv/n7csz8V0yoXAHww0ifXGXLl91sIo8jUrsc0ZEpcbLzESH7tR6wrOa3pBxV+yMrDwidsqDfnTGMyJ4PhC2zikM4HDwNuqsciUIYlY34fIcid4/1eep6AEMhHKtmfcLrBU6eg2WAaUnf1XaGmwK5J60zp0yX2fGo0XTxHm8ETKaTkgrnMhmaLvQ9BJ2Zx+a1sery0xYj/jIshMY4DZBMAjH5rvSV1T0SQslr0IBAGcyWV3jG9gwPTb5dtJmAX/WSYx1Sb52kbjYlUuBj7ffLADCjfVzKFWQxeQ="
+          secure: "d5zNL/guwonRQ7fI11ebzSTYLcrsTiVJROImJixYsr8G/NbM7/XEaRKr3nn/1br+LMd9TS89SeIY5MT1yAsCUYj8+zuACoOzi9lZslAwBi3sp2BhOve9T5LJhwHDA8WPehNq76jPaixN+0F9ytjpR4VE1xV6nlBjlXk5omqUUICFsbBjLAWS/rAFTsb2VmaECyF9dosdH26dXHb0oYJa3cpMbCHt/zbrAuFlaRG2qentgXv8PPCKBPxowGfTCUH/F25e/i7zzE1wbQBDcH4ZfkbvDFND6pMqb7Hfpumg1CZGrvgKEf1RmpKOCXHZWHDlm1C7H6hWaMACTw/AaHIEGNY2LHim3zEr19W6PYiEakVfGkfOXOo16pVSd65uxEp+C0+GB5KCimN57w7ZE9ucooEiaZMRfhIr8PX+4AGHfbLSj+6g0xDpg1x0woyPEvdzc1ZsZQJr4GY7OuI/4qqWrV48D+Hz5R2dVOxGFvaIYnYRGL3F/RmdPnO6B+IjXpQWHdm1FmhJM2jDG1wxXVqSTZmgEmUYKv8K5t0K/+f5SlJZ2tZOmbtDa3+Yo/1iTohCr0+OiDPojQY82EQmTiL8Ffr5aqpydwg8mQS6DX3m8w7vRzFvjlgoVQQ8ghP9BINhn2dvRMZdGfE/q8qqpL5qBIuLyHiU3Rnc6s1IxEbpw+Y="
         on:
           tags: true
           distributions: sdist
-          repo: lyft/pynamodb-mypy
+          repo: pynamodb/pynamodb-mypy

pynamodb_mypy/plugin.py

@@ -1,63 +1,123 @@
 from typing import Callable
 from typing import Optional
 
+import mypy.nodes
+import mypy.plugin
 import mypy.types
-from mypy.nodes import NameExpr
-from mypy.nodes import TypeInfo
-from mypy.plugin import FunctionContext
-from mypy.plugin import Plugin
 
 ATTR_FULL_NAME = 'pynamodb.attributes.Attribute'
-NULL_ATTR_WRAPPER_FULL_NAME = 'pynamodb.attributes._NullableAttributeWrapper'
 
 
-class PynamodbPlugin(Plugin):
-    def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext], mypy.types.Type]]:
+class PynamodbPlugin(mypy.plugin.Plugin):
+    def get_function_hook(self, fullname: str) -> Optional[Callable[[mypy.plugin.FunctionContext], mypy.types.Type]]:
         sym = self.lookup_fully_qualified(fullname)
-        if sym and isinstance(sym.node, TypeInfo):
-            attr_underlying_type = _get_attribute_underlying_type(sym.node)
-            if attr_underlying_type:
-                _underlying_type = attr_underlying_type  # https://github.com/python/mypy/issues/4297
-                return lambda ctx: _attribute_instantiation_hook(ctx, _underlying_type)
+        if sym and isinstance(sym.node, mypy.nodes.TypeInfo) and _is_attribute_type_node(sym.node):
+            return _attribute_instantiation_hook
+        return None
 
+    def get_method_signature_hook(self, fullname: str
+                                  ) -> Optional[Callable[[mypy.plugin.MethodSigContext], mypy.types.CallableType]]:
+        class_name, method_name = fullname.rsplit('.', 1)
+        sym = self.lookup_fully_qualified(class_name)
+        if sym is not None and sym.node is not None and _is_attribute_type_node(sym.node):
+            if method_name == '__get__':
+                return _get_method_sig_hook
+            elif method_name == '__set__':
+                return _set_method_sig_hook
         return None
 
 
-def _get_attribute_underlying_type(attribute_class: TypeInfo) -> Optional[mypy.types.Type]:
+def _is_attribute_type_node(node: mypy.nodes.Node) -> bool:
+    return (
+        isinstance(node, mypy.nodes.TypeInfo) and
+        node.has_base(ATTR_FULL_NAME)
+    )
+
+
+def _attribute_marked_as_nullable(t: mypy.types.Instance) -> mypy.types.Instance:
+    return t.copy_modified(args=t.args + [mypy.types.NoneType()])
+
+
+def _is_attribute_marked_nullable(t: mypy.types.Type) -> bool:
+    return (
+            isinstance(t, mypy.types.Instance) and
+            _is_attribute_type_node(t.type) and
+            # In lieu of being able to attach metadata to an instance,
+            # having a None "fake" type argument is our way of marking the attribute as nullable
+            bool(t.args) and isinstance(t.args[-1], mypy.types.NoneType)
+    )
+
+
+def _get_bool_literal(node: mypy.nodes.Node) -> Optional[bool]:
+    return {
+        'builtins.False': False,
+        'builtins.True': True,
+    }.get(node.fullname or '') if isinstance(node, mypy.nodes.NameExpr) else None
+
+
+def _make_optional(t: mypy.types.Type) -> mypy.types.UnionType:
+    """Wraps a type in optionality"""
+    return mypy.types.UnionType([t, mypy.types.NoneType()])
+
+
+def _unwrap_optional(t: mypy.types.Type) -> mypy.types.Type:
+    """Unwraps a potentially optional type"""
+    if not isinstance(t, mypy.types.UnionType):  # pragma: no cover
+        return t
+    t = mypy.types.UnionType([item for item in t.items if not isinstance(item, mypy.types.NoneType)])
+    if len(t.items) == 0:  # pragma: no cover
+        return mypy.types.NoneType()
+    elif len(t.items) == 1:
+        return t.items[0]
+    else:
+        return t  # pragma: no cover
+
+
+def _get_method_sig_hook(ctx: mypy.plugin.MethodSigContext) -> mypy.types.CallableType:
     """
-    For attribute classes, will return the underlying type.
-    e.g. for `class MyAttribute(Attribute[int])`, this will return `int`.
+    Patches up the signature of Attribute.__get__ to respect attribute's nullability.
     """
-    for base_instance in attribute_class.bases:
-        if base_instance.type.fullname() == ATTR_FULL_NAME:
-            return base_instance.args[0]
-    return None
+    sig = ctx.default_signature
+    if not _is_attribute_marked_nullable(ctx.type):
+        return sig
+    try:
+        (instance_type, owner_type) = sig.arg_types
+    except ValueError:  # pragma: no cover
+        return sig
+    if isinstance(instance_type, mypy.types.NoneType):  # class attribute access
+        return sig
+    return sig.copy_modified(ret_type=_make_optional(sig.ret_type))
 
 
-def _attribute_instantiation_hook(ctx: FunctionContext,
-                                  underlying_type: mypy.types.Type) -> mypy.types.Type:
+def _set_method_sig_hook(ctx: mypy.plugin.MethodSigContext) -> mypy.types.CallableType:
+    """
+    Patches up the signature of Attribute.__set__ to respect attribute's nullability.
+    """
+    sig = ctx.default_signature
+    if _is_attribute_marked_nullable(ctx.type):
+        return sig
+    try:
+        (instance_type, value_type) = sig.arg_types
+    except ValueError:  # pragma: no cover
+        return sig
+    return sig.copy_modified(arg_types=[instance_type, _unwrap_optional(value_type)])
+
+
+def _attribute_instantiation_hook(ctx: mypy.plugin.FunctionContext) -> mypy.types.Type:
     """
     Handles attribute instantiation, e.g. MyAttribute(null=True)
     """
     args = dict(zip(ctx.callee_arg_names, ctx.args))
 
-    # If initializer is passed null=True, wrap in _NullableAttribute
-    # to make the underlying type optional
+    # If initializer is passed null=True, mark attribute type instance as nullable
     null_arg_exprs = args.get('null')
+    nullable = False
     if null_arg_exprs and len(null_arg_exprs) == 1:
-        (null_arg_expr,) = null_arg_exprs
-        if (
-            not isinstance(null_arg_expr, NameExpr) or
-            null_arg_expr.fullname not in ('builtins.False', 'builtins.True')
-        ):
-            ctx.api.fail("'null' argument is not constant False or True, "
-                         "cannot deduce optionality", ctx.context)
-            return ctx.default_return_type
-
-        if null_arg_expr.fullname == 'builtins.True':
-            return ctx.api.named_generic_type(NULL_ATTR_WRAPPER_FULL_NAME, [
-                ctx.default_return_type,
-                underlying_type,
-            ])
-
-    return ctx.default_return_type
+        null_literal = _get_bool_literal(null_arg_exprs[0])
+        if null_literal is not None:
+            nullable = null_literal
+        else:
+            ctx.api.fail("'null' argument is not constant False or True, cannot deduce optionality", ctx.context)
+
+    assert isinstance(ctx.default_return_type, mypy.types.Instance)
+    return _attribute_marked_as_nullable(ctx.default_return_type) if nullable else ctx.default_return_type

setup.cfg

@@ -1,3 +1,29 @@
+[metadata]
+name = pynamodb-mypy
+version = 0.0.4
+description = mypy plugin for PynamoDB
+long_description = file: README.md
+long_description_content_type = text/markdown
+url = https://www.github.com/pynamodb/pynamodb-mypy
+classifiers =
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3 :: Only
+    Programming Language :: Python :: 3.6
+    Programming Language :: Python :: 3.7
+    Programming Language :: Python :: 3.8
+maintainer = Ilya Konstantinov
+maintainer_email = [email protected]
+
+[options]
+packages = find:
+install_requires =
+    mypy>=0.770
+python_requires = >=3.6
+
+[options.packages.find]
+exclude =
+    tests
+
 [flake8]
 format = pylint
 exclude = .svc,CVS,.bzr,.hg,.git,__pycache__,venv
@@ -21,3 +47,6 @@ disallow_untyped_defs = True
 ignore_missing_imports = True
 strict_optional = True
 warn_no_return = True
+
+[mypy-tests.*]
+disallow_untyped_defs = False

setup.py

@@ -1,18 +1,2 @@
-from setuptools import find_packages
 from setuptools import setup
-
-setup(
-    name='pynamodb-mypy',
-    version='0.0.1',
-    description='mypy plugin for PynamoDB',
-    url='https://www.github.com/lyft/pynamodb-mypy',
-    maintainer='Ilya Konstantinov',
-    maintainer_email='[email protected]',
-    packages=find_packages(exclude=['tests/*']),
-    install_requires=[
-        'mypy>=0.660',
-        # TODO: update version after https://github.com/pynamodb/PynamoDB/pull/579 is released
-        'pynamodb',
-    ],
-    python_requires='>=3',
-)
+setup()

tests/conftest.py

@@ -0,0 +1,7 @@
+import pytest
+
+
+@pytest.fixture
+def assert_mypy_output(pytestconfig):
+    from .mypy_helpers import assert_mypy_output
+    return lambda program: assert_mypy_output(program, use_pdb=pytestconfig.getoption('usepdb'))

tests/mypy_helpers.py

@@ -8,46 +8,55 @@
 from typing import Dict
 from typing import Iterable
 from typing import List
+from typing import Tuple
 
 import mypy.api
 
 
-def _run_mypy(program: str) -> Iterable[str]:
+def _run_mypy(program: str, *, use_pdb: bool) -> Iterable[str]:
     with TemporaryDirectory() as tempdirname:
         with open('{}/__main__.py'.format(tempdirname), 'w') as f:
             f.write(program)
         config_file = tempdirname + '/mypy.ini'
         shutil.copyfile(os.path.dirname(__file__) + '/mypy.ini', config_file)
-        error_pattern = re.compile(r'^{}:(\d+): error: (.*)$'.format(re.escape(f.name)))
-        stdout, stderr, exit_status = mypy.api.run([
+        error_pattern = re.compile(fr'^{re.escape(f.name)}:'
+                                   r'(?P<line>\d+): (?P<level>note|warning|error): (?P<message>.*)$')
+        mypy_args = [
             f.name,
             '--show-traceback',
+            '--raise-exceptions',
+            '--show-error-codes',
             '--config-file', config_file,
-        ])
+        ]
+        if use_pdb:
+            mypy_args.append('--pdb')
+        stdout, stderr, exit_status = mypy.api.run(mypy_args)
         if stderr:
             print(stderr, file=sys.stderr)  # allow "printf debugging" of the plugin
 
         # Group errors by line
-        errors_by_line: Dict[int, List[str]] = defaultdict(list)
+        messages_by_line: Dict[int, List[Tuple[str, str]]] = defaultdict(list)
         for line in stdout.split('\n'):
             m = error_pattern.match(line)
             if m:
-                errors_by_line[int(m.group(1))].append(m.group(2))
+                messages_by_line[int(m.group('line'))].append((m.group('level'), m.group('message')))
             elif line:
-                print(line)  # allow "printf debugging"
+                # print(line)  # allow "printf debugging"
+                pass
 
         # Reconstruct the "actual" program with "error" comments
-        error_comment_pattern = re.compile(r'(\s+# E: .*)?$')
+        error_comment_pattern = re.compile(r'(\s+# (N|W|E): .*)?$')
         for line_no, line in enumerate(program.split('\n'), start=1):
             line = error_comment_pattern.sub('', line)
-            errors = errors_by_line.get(line_no)
-            if errors:
-                yield '{}{}'.format(line, ''.join('  # E: {}'.format(error) for error in errors))
+            messages = messages_by_line.get(line_no)
+            if messages:
+                messages_str = ''.join(f'  # {level[0].upper()}: {message}' for level, message in messages)
+                yield f'{line}{messages_str}'
             else:
                 yield line
 
 
-def assert_mypy_output(program: str) -> None:
+def assert_mypy_output(program: str, *, use_pdb: bool) -> None:
     expected = dedent(program).strip()
-    actual = '\n'.join(_run_mypy(expected))
+    actual = '\n'.join(_run_mypy(expected, use_pdb=use_pdb))
     assert actual == expected