fix MapAttribute support

Author: ikonst

pynamodb_mypy/plugin.py

@@ -30,7 +30,7 @@ def get_method_signature_hook(self, fullname: str
 def _is_attribute_type_node(node: mypy.nodes.Node) -> bool:
     return (
         isinstance(node, mypy.nodes.TypeInfo) and
-        any(base.type.fullname == ATTR_FULL_NAME for base in node.bases)
+        node.has_base(ATTR_FULL_NAME)
     )
 
 
@@ -84,7 +84,7 @@ def _get_method_sig_hook(ctx: mypy.plugin.MethodSigContext) -> mypy.types.Callab
         (instance_type, owner_type) = sig.arg_types
     except ValueError:
         return sig
-    if not isinstance(instance_type, mypy.types.AnyType):  # instance attribute access
+    if isinstance(instance_type, mypy.types.NoneType):  # class attribute access
         return sig
     return sig.copy_modified(ret_type=_make_optional(sig.ret_type))
 

tests/test_plugin.py

@@ -14,13 +14,13 @@ class MyModel(Model):
     reveal_type(MyModel().my_not_nullable_attr)  # N: Revealed type is 'builtins.float*'
 
     my_model = MyModel()
-    my_model.my_attr = None
+    my_model.my_attr = None  # E: Incompatible types in assignment (expression has type "None", variable has type "float")  [assignment]
     my_model.my_nullable_attr = None
-    my_model.my_not_nullable_attr = None
+    my_model.my_not_nullable_attr = None  # E: Incompatible types in assignment (expression has type "None", variable has type "float")  [assignment]
     my_model.my_attr = 42
     my_model.my_nullable_attr = 42
     my_model.my_not_nullable_attr = 42
-    """)
+    """)  # noqa: E501
 
 
 def test_unicode_attribute(assert_mypy_output):
@@ -42,6 +42,25 @@ class MyModel(Model):
     """)
 
 
+def test_map_attribute(assert_mypy_output):
+    assert_mypy_output("""
+    from pynamodb.attributes import MapAttribute, UnicodeAttribute
+    from pynamodb.models import Model
+
+    class MyMapAttribute(MapAttribute):
+        my_sub_attr = UnicodeAttribute()
+
+    class MyModel(Model):
+        my_attr = MyMapAttribute()
+        my_nullable_attr = MyMapAttribute(null=True)
+
+    reveal_type(MyModel.my_attr)  # N: Revealed type is '__main__.MyMapAttribute'
+    reveal_type(MyModel.my_nullable_attr)  # N: Revealed type is '__main__.MyMapAttribute[None]'
+    reveal_type(MyModel().my_attr)  # N: Revealed type is '__main__.MyMapAttribute'
+    reveal_type(MyModel().my_nullable_attr)  # N: Revealed type is 'Union[__main__.MyMapAttribute[None], None]'
+    """)
+
+
 def test_custom_attribute(assert_mypy_output):
     assert_mypy_output("""
     from pynamodb.attributes import Attribute