Fix setter

Author: ikonst

pynamodb_mypy/plugin.py

@@ -15,11 +15,15 @@ def get_function_hook(self, fullname: str) -> Optional[Callable[[mypy.plugin.Fun
             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()])
+    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 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
 
 
@@ -48,6 +52,48 @@ def _get_bool_literal(n: mypy.nodes.Node) -> Optional[bool]:
     }.get(n.fullname or '') if isinstance(n, mypy.nodes.NameExpr) else None
 
 
+def _make_optional(t: mypy.types.Type) -> mypy.types.UnionType:
+    return mypy.types.UnionType([t, mypy.types.NoneType()])
+
+
+def _unwrap_optional(t: mypy.types.Type) -> mypy.types.Type:
+    if not isinstance(t, mypy.types.UnionType):
+        return t
+    t = mypy.types.UnionType([item for item in t.items if not isinstance(item, mypy.types.NoneType)])
+    if len(t.items) == 0:
+        return mypy.types.NoneType()
+    elif len(t.items) == 1:
+        return t.items[0]
+    else:
+        return t
+
+
+def _get_method_sig_hook(ctx: mypy.plugin.MethodSigContext) -> mypy.types.CallableType:
+    sig = ctx.default_signature
+    if not _is_attribute_marked_nullable(ctx.type):
+        return sig
+    try:
+        (instance_type, owner_type) = sig.arg_types
+    except ValueError:
+        return sig
+    if not isinstance(instance_type, mypy.types.AnyType):  # instance attribute access
+        return sig
+    return sig.copy_modified(ret_type=_make_optional(sig.ret_type))
+
+
+def _set_method_sig_hook(ctx: mypy.plugin.MethodSigContext) -> mypy.types.CallableType:
+    sig = ctx.default_signature
+    if _is_attribute_marked_nullable(ctx.type):
+        return sig
+    try:
+        (instance_type, value_type) = sig.arg_types
+    except ValueError:
+        return sig
+    if not isinstance(instance_type, mypy.types.AnyType):  # instance attribute access
+        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)

tests/mypy_helpers.py

@@ -24,6 +24,8 @@ def _run_mypy(program: str, *, use_pdb: bool) -> Iterable[str]:
         mypy_args = [
             f.name,
             '--show-traceback',
+            '--raise-exceptions',
+            '--show-error-codes',
             '--config-file', config_file,
         ]
         if use_pdb:

tests/test_plugin.py

@@ -12,6 +12,14 @@ class MyModel(Model):
     reveal_type(MyModel().my_attr)  # N: Revealed type is 'builtins.float*'
     reveal_type(MyModel().my_nullable_attr)  # N: Revealed type is 'Union[builtins.float*, None]'
     reveal_type(MyModel().my_not_nullable_attr)  # N: Revealed type is 'builtins.float*'
+
+    my_model = MyModel()
+    my_model.my_attr = None
+    my_model.my_nullable_attr = None
+    my_model.my_not_nullable_attr = None
+    my_model.my_attr = 42
+    my_model.my_nullable_attr = 42
+    my_model.my_not_nullable_attr = 42
     """)
 
 
@@ -30,7 +38,7 @@ class MyModel(Model):
     reveal_type(MyModel().my_nullable_attr)  # N: Revealed type is 'Union[builtins.str*, None]'
 
     MyModel().my_attr.lower()
-    MyModel().my_nullable_attr.lower()  # E: Item "None" of "Optional[str]" has no attribute "lower"
+    MyModel().my_nullable_attr.lower()  # E: Item "None" of "Optional[str]" has no attribute "lower"  [union-attr]
     """)
 
 
@@ -64,7 +72,7 @@ def test_unexpected_number_of_nulls(assert_mypy_output):
     from pynamodb.models import Model
 
     class MyModel(Model):
-        my_attr = NumberAttribute(True, True, True, null=True)  # E: "NumberAttribute" gets multiple values for keyword argument "null"
+        my_attr = NumberAttribute(True, True, True, null=True)  # E: "NumberAttribute" gets multiple values for keyword argument "null"  [misc]
 
     reveal_type(MyModel().my_attr)  # N: Revealed type is 'builtins.float*'
     """)  # noqa: E501
@@ -76,7 +84,7 @@ def test_unexpected_value_of_null(assert_mypy_output):
     from pynamodb.models import Model
 
     class MyModel(Model):
-        my_attr = NumberAttribute(null=bool(5))  # E: 'null' argument is not constant False or True, cannot deduce optionality
+        my_attr = NumberAttribute(null=bool(5))  # E: 'null' argument is not constant False or True, cannot deduce optionality  [misc]
 
     reveal_type(MyModel().my_attr)  # N: Revealed type is 'builtins.float*'
     """)  # noqa: E501