Store temporary annotations for defined attribute assignments

Summary:
Extend temporary annotation setting to attribute assignments. All temporary annotations are still cleared whenever we see a call.

This diff already makes sure that any property attributes are excluded from temporary refinement, but added a TODO to handle checking for the presence of getattr on the parent class to exclude those, too.

Will incrementally add temporary refinement + testing for asserts and other AST pieces that can trigger local type changes.

Reviewed By: grievejia

Differential Revision: D30400325

fbshipit-source-id: 4e3749c8bd0899c875985c0c3a0b73fb80e2553d

Author: shannonzhu

source/analysis/test/integration/annotationTest.ml

@@ -897,7 +897,6 @@ let test_check_immutable_annotations context =
     |}
     [
       "Missing attribute annotation [4]: Attribute `attribute` of class `Foo` has no type specified.";
-      "Incompatible return type [7]: Expected `int` but got `unknown`.";
     ];
   assert_type_errors
     {|

source/analysis/test/integration/attributeTest.ml

@@ -126,7 +126,7 @@ let test_check_attributes context =
     |}
     [
       "Undefined attribute [16]: `Foo` has no attribute `bar`.";
-      "Incompatible return type [7]: Expected `int` but got `unknown`.";
+      "Incompatible return type [7]: Expected `int` but got `str`.";
     ];
   assert_type_errors
     {|
@@ -248,7 +248,7 @@ let test_check_attributes context =
     |}
     [
       "Undefined attribute [16]: `Foo` has no attribute `bar`.";
-      "Incompatible return type [7]: Expected `int` but got `unknown`.";
+      "Incompatible return type [7]: Expected `int` but got `str`.";
     ];
   assert_type_errors
     {|
@@ -298,7 +298,6 @@ let test_check_attributes context =
     [
       "Uninitialized attribute [13]: Attribute `bar` is declared in class `Foo` "
       ^ "to have type `typing.Optional[int]` but is never initialized.";
-      "Incompatible return type [7]: Expected `int` but got `typing.Optional[int]`.";
     ];
   assert_type_errors
     {|
@@ -355,7 +354,6 @@ let test_check_attributes context =
       ^ "has type `int` but no type is specified.";
       "Missing attribute annotation [4]: Attribute `baz` of class `Foo` "
       ^ "has type `int` but no type is specified.";
-      "Incompatible return type [7]: Expected `int` but got `unknown`.";
     ];
   assert_type_errors
     {|
@@ -369,7 +367,7 @@ let test_check_attributes context =
     |}
     [
       "Undefined attribute [16]: `Foo` has no attribute `baz`.";
-      "Incompatible return type [7]: Expected `int` but got `unknown`.";
+      "Incompatible return type [7]: Expected `int` but got `typing.Optional[int]`.";
     ];
 
   (* Ensure synthetic attribute accesses don't mask errors on real ones. *)
@@ -617,7 +615,24 @@ let test_check_attributes context =
           self.attribute = not_annotated()
           a = self.attribute.something
     |}
-    ["Undefined attribute [16]: `int` has no attribute `something`."];
+    [];
+
+  assert_type_errors
+    {|
+      def returns_string() -> str:
+        return ""
+
+      class Foo:
+        attribute: int = 1
+        def foo(self) -> None:
+          self.attribute = returns_string()
+          a = self.attribute.something
+    |}
+    [
+      "Incompatible attribute type [8]: Attribute `attribute` declared in class `Foo` has type \
+       `int` but is used as type `str`.";
+      "Undefined attribute [16]: `int` has no attribute `something`.";
+    ];
 
   (* Do not resolve optional attributes to the optional type. *)
   assert_type_errors

source/analysis/test/integration/errorMessageTest.ml

@@ -126,8 +126,6 @@ let test_show_error_traces context =
        test.py:7:4.";
       "Unbound name [10]: Name `x` is used but not defined in the current scope. Did you forget to \
        import it or assign to it?";
-      "Incompatible return type [7]: Expected `str` but got `unknown`. Type `str` expected on line \
-       8, specified on line 5.";
     ];
   assert_type_errors
     {|

source/analysis/test/integration/refinementTest.ml

@@ -1044,6 +1044,71 @@ let test_check_temporary_refinement context =
        `typing_extensions.Literal[1]`).";
       "Revealed type [-1]: Revealed type for `MY_GLOBAL` is `float`.";
     ];
+  assert_type_errors
+    {|
+      from typing import Optional
+
+      class Foo:
+        attribute: Optional[int] = 1
+
+      def takes_non_optional_int(input: int) -> None:
+        pass
+
+      def test(foo: Foo) -> None:
+        reveal_type(foo.attribute)
+        foo.attribute = 1
+        reveal_type(foo.attribute)
+        takes_non_optional_int(foo.attribute)
+        reveal_type(foo.attribute)
+    |}
+    [
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]`.";
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]` (inferred: \
+       `typing_extensions.Literal[1]`).";
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]`.";
+    ];
+  assert_type_errors
+    {|
+      from typing import Optional
+
+      class Foo:
+        @property
+        def attribute(self) -> Optional[int]:
+          pass
+        @attribute.setter
+        def attribute(self, value: Optional[int]) -> None:
+          pass
+
+      def test(foo: Foo) -> None:
+        reveal_type(foo.attribute)
+        foo.attribute = 1
+        reveal_type(foo.attribute)
+    |}
+    [
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]`.";
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]`.";
+    ];
+  (* TODO(T70545381): Don't allow refinement if __getattr__ override exists. *)
+  assert_type_errors
+    {|
+      from typing import Optional
+
+      class Foo:
+        attribute: Optional[int] = 1
+
+        def __getattr__(self, value: str) -> int:
+          return 1
+
+      def test(foo: Foo) -> None:
+        reveal_type(foo.attribute)
+        foo.attribute = 1
+        reveal_type(foo.attribute)
+    |}
+    [
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]`.";
+      "Revealed type [-1]: Revealed type for `foo.attribute` is `Optional[int]` (inferred: \
+       `typing_extensions.Literal[1]`).";
+    ];
   ()
 
 

source/analysis/typeCheck.ml

@@ -4834,9 +4834,12 @@ module State (Context : Context) = struct
                         | Attribute _ as name when is_simple_name name -> (
                             match resolved_base, attribute with
                             | `Attribute (_, parent), Some (attribute, _)
-                              when not
-                                     (Annotated.Attribute.defined attribute
-                                     || is_undefined_attribute parent) ->
+                              when not (Annotated.Attribute.property attribute) ->
+                                let is_temporary_refinement =
+                                  is_temporary_refinement
+                                  || Annotated.Attribute.defined attribute
+                                  || is_undefined_attribute parent
+                                in
                                 Resolution.set_local_with_attributes
                                   ~temporary:is_temporary_refinement
                                   resolution