Add __setattr__ support (#3451)

* add __setattr__ support

* fix edge case of null signature, better tests for __setattr__

* modify test to be strict-optional friendly

* correct docs and test data

Author: emmatyping

docs/source/cheat_sheet.rst

@@ -149,6 +149,18 @@ When you're puzzled or when things are complicated
    reveal_type(c)  # -> error: Revealed type is 'builtins.list[builtins.str]'
    print(c)  # -> [4] the object is not cast
 
+   # if you want dynamic attributes on your class, have it override __setattr__ or __getattr__
+   # in a stub or in your source code.
+   # __setattr__ allows for dynamic assignment to names
+   # __getattr__ allows for dynamic access to names
+   class A:
+       # this will allow assignment to any A.x, if x is the same type as `value`
+       def __setattr__(self, name, value):
+           # type: (str, int) -> None
+           ...
+   a.foo = 42  # works
+   a.bar = 'Ex-parrot'  # fails type checking
+
    # TODO: explain "Need type annotation for variable" when
    # initializing with None or an empty container
 

docs/source/cheat_sheet_py3.rst

@@ -142,6 +142,19 @@ When you're puzzled or when things are complicated
    reveal_type(c)  # -> error: Revealed type is 'builtins.list[builtins.str]'
    print(c)  # -> [4] the object is not cast
 
+   # if you want dynamic attributes on your class, have it override __setattr__ or __getattr__
+   # in a stub or in your source code.
+   # __setattr__ allows for dynamic assignment to names
+   # __getattr__ allows for dynamic access to names
+   class A:
+       # this will allow assignment to any A.x, if x is the same type as `value`
+       def __setattr__(self, name: str, value: int) -> None: ...
+       # this will allow access to any A.x, if x is compatible with the return type
+       def __getattr__(self, name: str) -> int: ...
+   a.foo = 42  # works
+   a.bar = 'Ex-parrot'  # fails type checking
+
+
    # TODO: explain "Need type annotation for variable" when
    # initializing with None or an empty container
 

mypy/checker.py

@@ -616,7 +616,8 @@ def is_implicit_any(t: Type) -> bool:
                     self.check_reverse_op_method(item, typ, name)
                 elif name in ('__getattr__', '__getattribute__'):
                     self.check_getattr_method(typ, defn)
-
+                elif name == '__setattr__':
+                    self.check_setattr_method(typ, defn)
                 # Refuse contravariant return type variable
                 if isinstance(typ.ret_type, TypeVarType):
                     if typ.ret_type.variance == CONTRAVARIANT:
@@ -916,6 +917,15 @@ def check_getattr_method(self, typ: CallableType, context: Context) -> None:
         if not is_subtype(typ, method_type):
             self.msg.invalid_signature(typ, context)
 
+    def check_setattr_method(self, typ: CallableType, context: Context) -> None:
+        method_type = CallableType([AnyType(), self.named_type('builtins.str'), AnyType()],
+                                   [nodes.ARG_POS, nodes.ARG_POS, nodes.ARG_POS],
+                                   [None, None, None],
+                                   NoneTyp(),
+                                   self.named_type('builtins.function'))
+        if not is_subtype(typ, method_type):
+            self.msg.invalid_signature(typ, context)
+
     def expand_typevars(self, defn: FuncItem,
                         typ: CallableType) -> List[Tuple[FuncItem, CallableType]]:
         # TODO use generator

mypy/checkmember.py

@@ -245,6 +245,15 @@ def analyze_member_var_access(name: str, itype: Instance, info: TypeInfo,
                     getattr_type = expand_type_by_instance(bound_method, typ)
                     if isinstance(getattr_type, CallableType):
                         return getattr_type.ret_type
+        else:
+            setattr_meth = info.get_method('__setattr__')
+            if setattr_meth and setattr_meth.info.fullname() != 'builtins.object':
+                setattr_func = function_type(setattr_meth, builtin_type('builtins.function'))
+                bound_type = bind_self(setattr_func, original_type)
+                typ = map_instance_to_supertype(itype, setattr_meth.info)
+                setattr_type = expand_type_by_instance(bound_type, typ)
+                if isinstance(setattr_type, CallableType) and len(setattr_type.arg_types) > 0:
+                    return setattr_type.arg_types[-1]
 
     if itype.type.fallback_to_any:
         return AnyType()

test-data/unit/check-classes.test

@@ -1651,7 +1651,6 @@ b = a.bar
 [out]
 main:9: error: Incompatible types in assignment (expression has type "A", variable has type "B")
 
-
 [case testGetAttrSignature]
 class A:
     def __getattr__(self, x: str) -> A: pass
@@ -1665,6 +1664,85 @@ class D:
 main:4: error: Invalid signature "def (__main__.B, __main__.A) -> __main__.B"
 main:6: error: Invalid signature "def (__main__.C, builtins.str, builtins.str) -> __main__.C"
 
+[case testSetAttr]
+from typing import Union
+class A:
+    def __setattr__(self, name: str, value: Any) -> None: ...
+
+a = A()
+a.test = 'hello'
+
+class B:
+   def __setattr__(self, name: str, value: Union[int, str]) -> None: ...
+
+b = B()
+b.both = 1
+b.work = '2'
+
+class C:
+    def __setattr__(self, name: str, value: str) -> None: ...
+
+c = C()
+c.fail = 4  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+
+class D:
+    __setattr__ = 'hello'
+
+d = D()
+d.crash = 4  # E: "D" has no attribute "crash"
+
+class Ex:
+    def __setattr__(self, name: str, value: int) -> None:...
+    test = '42'  # type: str
+e = Ex()
+e.test = 'hello'
+e.t = 4
+
+class Super:
+    def __setattr__(self, name: str, value: int) -> None: ...
+
+class Sub(Super):
+    ...
+s = Sub()
+s.success = 4
+s.fail = 'fail'  # E: Incompatible types in assignment (expression has type "str", variable has type "int")
+
+[case testSetAttrSignature]
+class Test:
+    def __setattr__() -> None: ...  # E: Method must have at least one argument # E: Invalid signature "def ()"
+t = Test()
+t.crash = 'test'  # E: "Test" has no attribute "crash"
+
+class A:
+    def __setattr__(self): ...  # E: Invalid signature "def (self: Any) -> Any"
+a = A()
+a.test = 4  # E: "A" has no attribute "test"
+
+class B:
+    def __setattr__(self, name, value: int): ...
+b = B()
+b.integer = 5
+
+class C:
+    def __setattr__(self, name: int, value: int) -> None: ...  # E: Invalid signature "def (__main__.C, builtins.int, builtins.int)"
+c = C()
+c.check = 13
+
+[case testGetAttrAndSetattr]
+class A:
+    def __setattr__(self, name: str, value: Any) -> None: ...
+    def __getattr__(self, name: str) -> Any: ...
+a = A()
+a.test = 4
+t = a.test
+
+class B:
+    def __setattr__(self, name: str, value: int) -> None: ...
+    def __getattr__(self, name: str) -> str: ...
+integer = 0
+b = B()
+b.at = '3'  # E: Incompatible types in assignment (expression has type "str", variable has type "int")
+integer = b.at  # E: Incompatible types in assignment (expression has type "str", variable has type "int")
 
 -- CallableType objects
 -- ----------------