Store temporary annotations when refinement is disallowed

shannonzhu · facebook-github-bot · commit b0a7de9d22b7 · 2021-08-19T11:08:28.000-07:00
Summary:
On assignment for globals, start storing temporary type data on otherwise non-refinable values.

Also start clearing all temporary type data whenever we see a call expression.

Future diffs will add temporary type data on attributes, asserts, etc. patterns that also can set local types.

Reviewed By: grievejia

Differential Revision: D30381727

fbshipit-source-id: 870cbc35e9b907510c0b9fb843c2a037540e37dc
diff --git a/source/analysis/resolution.ml b/source/analysis/resolution.ml
@@ -133,25 +133,30 @@ let partition_name resolution ~name =
 
 
 let set_local
+    ?(temporary = false)
     ({ annotation_store = { annotations; temporary_annotations }; _ } as resolution)
     ~reference
     ~annotation
   =
-  {
-    resolution with
-    annotation_store =
-      {
-        annotations =
-          Map.set
-            annotations
-            ~key:reference
-            ~data:(RefinementUnit.create () |> RefinementUnit.set_base ~base:annotation);
-        temporary_annotations;
-      };
-  }
+  let annotations, temporary_annotations =
+    if temporary then
+      ( annotations,
+        Map.set
+          temporary_annotations
+          ~key:reference
+          ~data:(RefinementUnit.create () |> RefinementUnit.set_base ~base:annotation) )
+    else
+      ( Map.set
+          annotations
+          ~key:reference
+          ~data:(RefinementUnit.create () |> RefinementUnit.set_base ~base:annotation),
+        temporary_annotations )
+  in
+  { resolution with annotation_store = { annotations; temporary_annotations } }
 
 
 let set_local_with_attributes
+    ?(temporary = false)
     ({ annotation_store = { annotations; temporary_annotations }; _ } as resolution)
     ~name
     ~annotation
@@ -162,30 +167,41 @@ let set_local_with_attributes
     | None, Some base -> RefinementUnit.set_base refinement_unit ~base
     | _ -> refinement_unit
   in
-  {
-    resolution with
-    annotation_store =
-      {
-        annotations =
-          Map.set
-            annotations
-            ~key:object_reference
-            ~data:
-              (Map.find annotations object_reference
-              |> Option.value ~default:(RefinementUnit.create ())
-              |> RefinementUnit.add_attribute_refinement ~reference:attribute_path ~base:annotation
-              |> set_base ~base);
-        temporary_annotations;
-      };
-  }
+  let annotations, temporary_annotations =
+    if temporary then
+      ( annotations,
+        Map.set
+          temporary_annotations
+          ~key:object_reference
+          ~data:
+            (Map.find temporary_annotations object_reference
+            |> (fun existing -> Option.first_some existing (Map.find annotations object_reference))
+            |> Option.value ~default:(RefinementUnit.create ())
+            |> RefinementUnit.add_attribute_refinement ~reference:attribute_path ~base:annotation
+            |> set_base ~base) )
+    else
+      ( Map.set
+          annotations
+          ~key:object_reference
+          ~data:
+            (Map.find annotations object_reference
+            |> Option.value ~default:(RefinementUnit.create ())
+            |> RefinementUnit.add_attribute_refinement ~reference:attribute_path ~base:annotation
+            |> set_base ~base),
+        temporary_annotations )
+  in
+  { resolution with annotation_store = { annotations; temporary_annotations } }
 
 
 let get_local
     ?(global_fallback = true)
     ~reference
-    ({ annotation_store = { annotations; _ }; global_resolution; _ } as resolution)
+    ({ annotation_store = { annotations; temporary_annotations }; global_resolution; _ } as
+    resolution)
   =
-  match Map.find annotations reference with
+  match
+    Option.first_some (Map.find temporary_annotations reference) (Map.find annotations reference)
+  with
   | Some result when global_fallback || not (is_global resolution ~reference) ->
       RefinementUnit.base result
   | _ when global_fallback ->
@@ -197,10 +213,15 @@ let get_local
 let get_local_with_attributes
     ?(global_fallback = true)
     ~name
-    ({ annotation_store = { annotations; _ }; global_resolution; _ } as resolution)
+    ({ annotation_store = { annotations; temporary_annotations }; global_resolution; _ } as
+    resolution)
   =
   let object_reference, attribute_path, _ = partition_name resolution ~name in
-  match Map.find annotations object_reference with
+  match
+    Option.first_some
+      (Map.find temporary_annotations object_reference)
+      (Map.find annotations object_reference)
+  with
   | Some result when global_fallback || not (is_global resolution ~reference:object_reference) ->
       RefinementUnit.annotation result ~reference:attribute_path
   | _ when global_fallback ->
@@ -217,7 +238,18 @@ let unset_local
   =
   {
     resolution with
-    annotation_store = { annotations = Map.remove annotations reference; temporary_annotations };
+    annotation_store =
+      {
+        annotations = Map.remove annotations reference;
+        temporary_annotations = Map.remove temporary_annotations reference;
+      };
+  }
+
+
+let clear_temporary_annotations ({ annotation_store; _ } as resolution) =
+  {
+    resolution with
+    annotation_store = { annotation_store with temporary_annotations = Reference.Map.empty };
   }
 
 
diff --git a/source/analysis/resolution.mli b/source/analysis/resolution.mli
@@ -55,9 +55,14 @@ val resolve_attribute_access : t -> base_type:Type.t -> attribute:string -> Type
 
 val partition_name : t -> name:Expression.Name.t -> Reference.t * Reference.t * Annotation.t option
 
-val set_local : t -> reference:Reference.t -> annotation:Annotation.t -> t
+val set_local : ?temporary:bool -> t -> reference:Reference.t -> annotation:Annotation.t -> t
 
-val set_local_with_attributes : t -> name:Expression.Name.t -> annotation:Annotation.t -> t
+val set_local_with_attributes
+  :  ?temporary:bool ->
+  t ->
+  name:Expression.Name.t ->
+  annotation:Annotation.t ->
+  t
 
 val get_local : ?global_fallback:bool -> reference:Reference.t -> t -> Annotation.t option
 
@@ -69,6 +74,8 @@ val get_local_with_attributes
 
 val unset_local : t -> reference:Reference.t -> t
 
+val clear_temporary_annotations : t -> t
+
 val is_global : t -> reference:Reference.t -> bool
 
 val add_type_variable : t -> variable:Type.Variable.t -> t
diff --git a/source/analysis/test/integration/annotationTest.ml b/source/analysis/test/integration/annotationTest.ml
@@ -743,10 +743,7 @@ let test_check_immutable_annotations context =
         constant: str
         constant = "hi"
     |}
-    [
-      "Incompatible variable type [9]: constant is declared to have type `int` but is used as type \
-       `str`.";
-    ];
+    [];
   assert_type_errors
     {|
       import typing
diff --git a/source/analysis/test/integration/refinementTest.ml b/source/analysis/test/integration/refinementTest.ml
@@ -386,7 +386,10 @@ let test_check_global_refinement context =
         if x is not None:
           reveal_type(x)
     |}
-    ["Revealed type [-1]: Revealed type for `x` is `typing.Optional[int]`."]
+    [
+      "Revealed type [-1]: Revealed type for `x` is `typing.Optional[int]` (inferred: \
+       `typing_extensions.Literal[1]`).";
+    ]
 
 
 let test_check_local_refinement context =
@@ -1018,6 +1021,32 @@ let test_check_final_attribute_refinement context =
   ()
 
 
+let test_check_temporary_refinement context =
+  let assert_type_errors = assert_type_errors ~context in
+  assert_type_errors
+    {|
+      MY_GLOBAL = 1.0
+
+      def arbitrary_call() -> None:
+        pass
+
+      def test() -> None:
+        reveal_type(MY_GLOBAL)
+        global MY_GLOBAL
+        MY_GLOBAL = 1
+        reveal_type(MY_GLOBAL)
+        arbitrary_call()
+        reveal_type(MY_GLOBAL)
+    |}
+    [
+      "Revealed type [-1]: Revealed type for `MY_GLOBAL` is `float`.";
+      "Revealed type [-1]: Revealed type for `MY_GLOBAL` is `float` (inferred: \
+       `typing_extensions.Literal[1]`).";
+      "Revealed type [-1]: Revealed type for `MY_GLOBAL` is `float`.";
+    ];
+  ()
+
+
 let () =
   "refinement"
   >::: [
@@ -1029,5 +1058,6 @@ let () =
          "check_assert_contains_none" >:: test_assert_contains_none;
          "check_callable" >:: test_check_callable;
          "check_final_attribute_refinement" >:: test_check_final_attribute_refinement;
+         "check_temporary_refinement" >:: test_check_temporary_refinement;
        ]
   |> Test.run
diff --git a/source/analysis/test/integration/revealTypeTest.ml b/source/analysis/test/integration/revealTypeTest.ml
@@ -126,7 +126,10 @@ let test_reveal_type context =
          x = 1
          reveal_type(x)
      |}
-    ["Revealed type [-1]: Revealed type for `x` is `float`."];
+    [
+      "Revealed type [-1]: Revealed type for `x` is `float` (inferred: \
+       `typing_extensions.Literal[1]`).";
+    ];
   assert_type_errors
     {|
        import typing
diff --git a/source/analysis/typeCheck.ml b/source/analysis/typeCheck.ml
