Rework nullable attribute support

Author: ikonst

.gitignore

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

.pre-commit-config.yaml

@@ -14,43 +14,29 @@ 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/.*|
-            )$
+
+#    - repo: local
+#      hooks:
+#        - id: mypy
+#          args: ["--ignore-missing-imports", "--scripts-are-modules", "--show-traceback"]
+#          exclude: >
+#            (?x)^(
+#              tests/.*|
+#            )$

pynamodb_mypy/plugin.py

@@ -1,63 +1,68 @@
 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_attribute_hook(self, fullname: str
+                           ) -> Optional[Callable[[mypy.plugin.AttributeContext], mypy.types.Type]]:
+        sym = self.lookup_fully_qualified(fullname)
+        if sym and sym.type and _is_attribute_marked_nullable(sym.type):
+            return lambda ctx: mypy.types.UnionType([ctx.default_attr_type, mypy.types.NoneType()])
         return None
 
 
-def _get_attribute_underlying_type(attribute_class: TypeInfo) -> Optional[mypy.types.Type]:
-    """
-    For attribute classes, will return the underlying type.
-    e.g. for `class MyAttribute(Attribute[int])`, this will return `int`.
-    """
-    for base_instance in attribute_class.bases:
-        if base_instance.type.fullname() == ATTR_FULL_NAME:
-            return base_instance.args[0]
-    return None
+def _is_attribute_type_node(type_node: mypy.nodes.TypeInfo) -> bool:
+    return any(base.type.fullname == ATTR_FULL_NAME for base in type_node.bases)
+
+
+def _attribute_marked_as_nullable(t: mypy.types.Instance) -> mypy.types.Instance:
+    return t.copy_modified(args=t.args + [mypy.types.NoneType()])
 
 
-def _attribute_instantiation_hook(ctx: FunctionContext,
-                                  underlying_type: mypy.types.Type) -> mypy.types.Type:
+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(n: mypy.nodes.Node) -> Optional[bool]:
+    return {
+        'builtins.False': False,
+        'builtins.True': True,
+    }.get(n.fullname or '') if isinstance(n, mypy.nodes.NameExpr) else None
+
+
+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,30 @@
+[metadata]
+name = pynamodb-mypy
+version = 0.0.2
+description = mypy plugin for PynamoDB
+long_description = file: README.md
+long_description_content_type = text/markdown
+url = https://www.github.com/lyft/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 = ilya.konstantinov@gmail.com
+
+[options]
+packages = find:
+install_requires =
+    mypy>=0.770
+python_requires = >=3.6
+
+[options.packages.find]
+where = pynamodb_mypy
+exclude =
+    tests
+
 [flake8]
 format = pylint
 exclude = .svc,CVS,.bzr,.hg,.git,__pycache__,venv
@@ -21,3 +48,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='ilya.konstantinov@gmail.com',
-    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,53 @@
 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',
             '--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