Use type= to set the annotation

python · ilevkivskyi · Feb 13, 2018 · Jan 11, 2018 · Jan 11, 2018 · Jan 11, 2018
commit 405204d205e24faa990152f120c3872c2b88139d
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -571,6 +571,7 @@ def add_attribute(attr_name: str, attr_type: Optional[Type], default: bool, init
         for stmt in info.defn.defs.body:
             if isinstance(stmt, AssignmentStmt) and isinstance(stmt.lvalues[0], NameExpr):
                 lhs = stmt.lvalues[0]
+                assert isinstance(lhs, NameExpr)
                 name = lhs.name
                 typ = stmt.type
 
@@ -594,6 +595,10 @@ def add_attribute(attr_name: str, attr_type: Optional[Type], default: bool, init
                             ctx.api.fail('Invalid argument to type', type_arg)
                         else:
                             typ = ctx.api.anal_type(un_type)
+                            if typ and isinstance(lhs.node, Var) and not lhs.node.type:
+                                # If there is no annotation, add one.
+                                lhs.node.type = typ
+                                lhs.is_inferred_def = False
 
                     add_attribute(name, typ, attr_has_default,
                                   get_bool_argument(stmt.rvalue, attrib.init),

diff --git a/test-data/unit/check-attr.test b/test-data/unit/check-attr.test
@@ -140,11 +140,11 @@ class A:
 @attr.s
 class B(A):
     b: str = attr.ib()
-    x: int = attr.ib(default='value')
+    x: int = attr.ib(default=22)
 
 @attr.s
 class C(B):
-    c: bool = attr.ib()  # No error here because the x below overwrite the x above.
+    c: bool = attr.ib()  # No error here because the x below overwrites the x above.
 
 reveal_type(A)  # E: Revealed type is 'def (a: builtins.int, x: builtins.int) -> __main__.A'
@@ -158,7 +158,8 @@ import attr
 @attr.s
 class A:
     a = attr.ib(type=int)
-reveal_type(A) # E: Revealed type is 'def (a: builtins.int) -> __main__.A'
+    b = attr.ib(18, type=int)
+reveal_type(A) # E: Revealed type is 'def (a: builtins.int, b: builtins.int =) -> __main__.A'
 [builtins fixtures/bool.pyi]
 
 [case testAttrsFrozen]

diff --git a/test-data/unit/lib-stub/attr.pyi b/test-data/unit/lib-stub/attr.pyi
@@ -2,7 +2,10 @@ from typing import TypeVar, overload, Callable, Any, Type
 
 _T = TypeVar('_T')
 
-def attr(default: Any = ..., validator: Any = ..., type: Type[_T] = ...) -> Any: ...
+@overload
+def attr(default: _T = ..., validator: Any = ..., type: type = ...) -> _T: ...
+@overload
+def attr(default: None = ..., validator: Any = ..., type: type = ...) -> Any: ...
 
 @overload
 def attributes(maybe_cls: _T = ..., cmp: bool = ..., init: bool = ..., frozen: bool = ...) -> _T: ...