Add origins to await expressions

Summary: We also create implicit await expressions during preprocessing and control flow graph building, thus those need origins.

Reviewed By: tianhan0

Differential Revision: D74737704

fbshipit-source-id: 8ff7d68d1f0d011884c8748d3e2b70d05e7684a2

Author: arthaud

source/analysis/attributeResolution.ml

@@ -3425,8 +3425,8 @@ class base ~queries:(Queries.{ controls; get_class_summary; class_hierarchy; _ }
       in
       let order = self#full_order ~cycle_detections in
       match Node.value expression with
-      | Expression.Await expression ->
-          self#resolve_literal ~cycle_detections expression ~scoped_type_variables
+      | Expression.Await { Await.operand; origin = _ } ->
+          self#resolve_literal ~cycle_detections operand ~scoped_type_variables
           |> Type.awaitable_value
           |> Option.value ~default:Type.Any
       | BooleanOperator { BooleanOperator.left; right; _ } ->

source/analysis/decoratorPreprocessing.ml

@@ -374,21 +374,22 @@ let convert_parameter_to_argument ~location { Node.value = { Parameter.name; _ }
 
 
 let create_function_call ~should_await ~location ~callee_name arguments =
+  let origin = Some (Origin.create ~location Origin.DecoratorInlining) in
   let call =
     Expression.Call
       {
         callee =
           Expression.Name
-            (create_name_from_reference
-               ~location
-               ~create_origin:(fun _ -> Some (Origin.create ~location Origin.DecoratorInlining))
-               callee_name)
+            (create_name_from_reference ~location ~create_origin:(fun _ -> origin) callee_name)
           |> Node.create ~location;
         arguments;
-        origin = Some (Origin.create ~location Origin.DecoratorInlining);
+        origin;
       }
   in
-  if should_await then Expression.Await (Node.create ~location call) else call
+  if should_await then
+    Expression.Await { Await.operand = Node.create ~location call; origin }
+  else
+    call
 
 
 let create_function_call_to ~location ~callee_name { Define.Signature.parameters; async; _ } =

source/analysis/globalLeakCheck.ml

@@ -318,7 +318,7 @@ module State (Context : Context) = struct
     | Expression.Constant _
     | Yield None ->
         empty_result
-    | Await expression
+    | Await { Await.operand = expression; origin = _ }
     | Yield (Some expression)
     | YieldFrom expression
     | UnaryOperator { operand = expression; _ }

source/analysis/preprocessing.ml

@@ -587,8 +587,9 @@ module Qualify = struct
         scope, List.rev reversed_generators
       in
       match value with
-      | Expression.Await expression ->
-          Expression.Await (qualify_expression ~qualify_strings ~scope expression)
+      | Expression.Await { Await.operand; origin } ->
+          Expression.Await
+            { Await.operand = qualify_expression ~qualify_strings ~scope operand; origin }
       | BinaryOperator { BinaryOperator.left; operator; right; origin } ->
           BinaryOperator
             {
@@ -3242,7 +3243,7 @@ module AccessCollector = struct
         (* For attribute access, only count the base *)
         from_expression collected base
     (* The rest is boilerplates to make sure that expressions are visited recursively *)
-    | Await await -> from_expression collected await
+    | Await { Await.operand; origin = _ } -> from_expression collected operand
     | BinaryOperator { BinaryOperator.left; right; _ }
     | BooleanOperator { BooleanOperator.left; right; _ }
     | ComparisonOperator { ComparisonOperator.left; right; _ } ->

source/analysis/scope.ml

@@ -80,7 +80,7 @@ module Binding = struct
         let sofar = of_expression sofar value in
         of_unannotated_target ~kind:Kind.WalrusTarget sofar target
     (* Boilerplates to make sure all expressions are visited. *)
-    | Expression.Await expression
+    | Expression.Await { Await.operand = expression; origin = _ }
     | Expression.Name (Name.Attribute { Name.Attribute.base = expression; _ })
     | Expression.UnaryOperator { UnaryOperator.operand = expression; _ }
     | Expression.Starred (Starred.Once expression | Starred.Twice expression)

source/analysis/typeCheck.ml

@@ -2088,7 +2088,7 @@ module State (Context : Context) = struct
     in
     let resolved =
       match value with
-      | Await expression -> (
+      | Await { Await.operand = expression; origin = _ } -> (
           let { Resolved.resolution; resolved; errors; _ } =
             forward_expression ~resolution expression
           in

source/analysis/unawaitedAwaitableCheck.ml

@@ -665,12 +665,13 @@ module State (Context : Context) = struct
       Return the new state along with nested awaitable expressions. *)
   and forward_expression ~resolution ~(state : t) ~expression:{ Node.value; location } =
     match value with
-    | Await ({ Node.value = Name name; _ } as expression) when is_simple_name name ->
+    | Await { Await.operand = { Node.value = Name name; _ } as expression; origin = _ }
+      when is_simple_name name ->
         let { state; nested_awaitable_expressions } =
           forward_expression ~resolution ~state ~expression
         in
         { state = mark_name_as_awaited state ~name; nested_awaitable_expressions }
-    | Await ({ Node.location; _ } as expression) ->
+    | Await { Await.operand = { Node.location; _ } as expression; origin = _ } ->
         let { state; nested_awaitable_expressions } =
           forward_expression ~resolution ~state ~expression
         in

source/analysis/uninitializedLocalCheck.ml

@@ -63,7 +63,7 @@ module AccessCollector = struct
         | Name (Name.Identifier _) -> collected
         | _ -> from_expression collected base)
     (* The rest is boilerplates to make sure that expressions are visited recursively *)
-    | Await await -> from_expression collected await
+    | Await { Await.operand; _ } -> from_expression collected operand
     | BinaryOperator { BinaryOperator.left; right; _ }
     | BooleanOperator { BooleanOperator.left; right; _ }
     | ComparisonOperator { ComparisonOperator.left; right; _ } ->

source/ast/expression.ml

@@ -262,6 +262,28 @@ end = struct
     | _ -> List.compare Argument.location_insensitive_compare left.arguments right.arguments
 end
 
+and Await : sig
+  type t = {
+    operand: Expression.t;
+    (* If this AST node was created from lowering down another AST node (for instance, `a + b` is
+       turned into `a.__add__(b)`), `origin` stores the location of the original AST node and the
+       reason for lowering it. *)
+    origin: Origin.t option;
+  }
+  [@@deriving equal, compare, sexp, show, hash, to_yojson]
+
+  val location_insensitive_compare : t -> t -> int
+end = struct
+  type t = {
+    operand: Expression.t;
+    origin: Origin.t option;
+  }
+  [@@deriving equal, compare, sexp, show, hash, to_yojson]
+
+  let location_insensitive_compare left right =
+    Expression.location_insensitive_compare left.operand right.operand
+end
+
 and ComparisonOperator : sig
   type operator =
     | Equals
@@ -1147,8 +1169,10 @@ and Origin : sig
     | SubscriptGetItem (* `d[a]` is turned into `d.__getitem__(a)` *)
     | ForIter (* `for e in l:` is turned into `l.__iter__().__next__()` *)
     | ForNext (* `for e in l:` is turned into `l.__iter__().__next__()` *)
+    | ForAwait (* `for e in l:` might be turned into `await l.__iter__().__next__()` *)
     | GeneratorIter (* `(e for e in l)` is turned into `l.__iter__().__next__()` *)
     | GeneratorNext (* `(e for e in l)` is turned into `l.__iter__().__next__()` *)
+    | GeneratorAwait (* `(e for e in l)` might be turned into `await l.__iter__().__next__()` *)
     | With (* `with e1 as e2` is turned into `e2 = e1.__enter__()` *)
     | InContains (* `e in l` can be turned into `l.__contains__(e)` *)
     | InIter (* `e in l` can be turned into `l.__iter__().__next__().__eq__(e)` *)
@@ -1253,8 +1277,10 @@ end = struct
     | SubscriptGetItem
     | ForIter
     | ForNext
+    | ForAwait
     | GeneratorIter
     | GeneratorNext
+    | GeneratorAwait
     | With
     | InContains
     | InIter
@@ -1368,7 +1394,7 @@ end
 
 and Expression : sig
   type expression =
-    | Await of t
+    | Await of Await.t
     | BinaryOperator of BinaryOperator.t
     | BooleanOperator of BooleanOperator.t
     | Call of Call.t
@@ -1407,7 +1433,7 @@ and Expression : sig
   val pp_type_param_list : Format.formatter -> TypeParam.t list -> unit
 end = struct
   type expression =
-    | Await of t
+    | Await of Await.t
     | BinaryOperator of BinaryOperator.t
     | BooleanOperator of BooleanOperator.t
     | Call of Call.t
@@ -1439,7 +1465,7 @@ end = struct
 
   let rec location_insensitive_compare_expression left right =
     match left, right with
-    | Await left, Await right -> location_insensitive_compare left right
+    | Await left, Await right -> Await.location_insensitive_compare left right
     | BinaryOperator left, BinaryOperator right ->
         BinaryOperator.location_insensitive_compare left right
     | BooleanOperator left, BooleanOperator right ->
@@ -1700,7 +1726,7 @@ end = struct
 
     and pp_expression formatter expression =
       match expression with
-      | Await expression -> Format.fprintf formatter "await %a" pp_expression_t expression
+      | Await { Await.operand; _ } -> Format.fprintf formatter "await %a" pp_expression_t operand
       | BinaryOperator { BinaryOperator.left; operator; right; origin = _ } ->
           Format.fprintf
             formatter
@@ -1823,7 +1849,7 @@ end
 
 module Mapper = struct
   type 'a t = {
-    map_await: mapper:'a t -> location:Location.t -> Expression.t -> 'a;
+    map_await: mapper:'a t -> location:Location.t -> Await.t -> 'a;
     map_binary_operator: mapper:'a t -> location:Location.t -> BinaryOperator.t -> 'a;
     map_boolean_operator: mapper:'a t -> location:Location.t -> BooleanOperator.t -> 'a;
     map_call: mapper:'a t -> location:Location.t -> Call.t -> 'a;
@@ -1884,7 +1910,7 @@ module Mapper = struct
       { Node.value; location }
     =
     match value with
-    | Expression.Await expression -> map_await ~mapper ~location expression
+    | Expression.Await await -> map_await ~mapper ~location await
     | Expression.BinaryOperator binary_operator ->
         map_binary_operator ~mapper ~location binary_operator
     | Expression.BooleanOperator boolean_operator ->
@@ -2006,7 +2032,9 @@ module Mapper = struct
     default_map_substrings_with_location ~mapper ~map_location:Fn.id substrings
 
 
-  let default_map_await ~mapper awaited = map ~mapper awaited
+  let default_map_await ~mapper { Await.operand; origin } =
+    { Await.operand = map ~mapper operand; origin }
+
 
   let default_map_await_node ~mapper ~location awaited =
     Node.create ~location (Expression.Await (default_map_await ~mapper awaited))
@@ -2530,7 +2558,7 @@ end
 
 module Folder = struct
   type 'a t = {
-    fold_await: folder:'a t -> state:'a -> location:Location.t -> Expression.t -> 'a;
+    fold_await: folder:'a t -> state:'a -> location:Location.t -> Await.t -> 'a;
     fold_binary_operator: folder:'a t -> state:'a -> location:Location.t -> BinaryOperator.t -> 'a;
     fold_boolean_operator:
       folder:'a t -> state:'a -> location:Location.t -> BooleanOperator.t -> 'a;
@@ -2600,7 +2628,7 @@ module Folder = struct
       { Node.value; location }
     =
     match value with
-    | Expression.Await expression -> fold_await ~folder ~state ~location expression
+    | Expression.Await await -> fold_await ~folder ~state ~location await
     | Expression.BinaryOperator binary_operator ->
         fold_binary_operator ~folder ~state ~location binary_operator
     | Expression.BooleanOperator boolean_operator ->
@@ -2740,7 +2768,7 @@ module Folder = struct
       substrings
 
 
-  let default_fold_await ~folder ~state awaited = fold ~folder ~state awaited
+  let default_fold_await ~folder ~state { Await.operand; origin = _ } = fold ~folder ~state operand
 
   let default_fold_binary_operator
       ~folder
@@ -2997,6 +3025,7 @@ let origin { Node.value; _ } =
   | Expression.Subscript { Subscript.origin; _ } -> origin
   | Expression.WalrusOperator { WalrusOperator.origin; _ } -> origin
   | Expression.Slice { Slice.origin; _ } -> origin
+  | Expression.Await { Await.origin; _ } -> origin
   | _ -> None
 
 
@@ -3560,6 +3589,9 @@ let remove_origins expression =
       origin = None;
     }
   in
+  let map_await ~mapper { Await.operand; origin = _ } =
+    { Await.operand = Mapper.map ~mapper operand; origin = None }
+  in
   Mapper.map
     ~mapper:
       (Mapper.create_transformer
@@ -3572,5 +3604,6 @@ let remove_origins expression =
          ~map_subscript
          ~map_walrus_operator
          ~map_slice
+         ~map_await
          ())
     expression

source/ast/expression.mli

@@ -103,6 +103,16 @@ and Call : sig
   val location_insensitive_compare : t -> t -> int
 end
 
+and Await : sig
+  type t = {
+    operand: Expression.t;
+    origin: Origin.t option;
+  }
+  [@@deriving equal, compare, sexp, show, hash, to_yojson]
+
+  val location_insensitive_compare : t -> t -> int
+end
+
 and ComparisonOperator : sig
   type operator =
     | Equals
@@ -408,8 +418,10 @@ and Origin : sig
     | SubscriptGetItem (* `d[a]` is turned into `d.__getitem__(a)` *)
     | ForIter (* `for e in l:` is turned into `l.__iter__().__next__()` *)
     | ForNext (* `for e in l:` is turned into `l.__iter__().__next__()` *)
+    | ForAwait (* `for e in l:` might be turned into `await l.__iter__().__next__()` *)
     | GeneratorIter (* `(e for e in l)` is turned into `l.__iter__().__next__()` *)
     | GeneratorNext (* `(e for e in l)` is turned into `l.__iter__().__next__()` *)
+    | GeneratorAwait (* `(e for e in l)` might be turned into `await l.__iter__().__next__()` *)
     | With (* `with e1 as e2` is turned into `e2 = e1.__enter__()` *)
     | InContains (* `e in l` can be turned into `l.__contains__(e)` *)
     | InIter (* `e in l` can be turned into `l.__iter__().__next__().__eq__(e)` *)
@@ -507,7 +519,7 @@ end
 
 and Expression : sig
   type expression =
-    | Await of t
+    | Await of Await.t
     | BinaryOperator of BinaryOperator.t
     | BooleanOperator of BooleanOperator.t
     | Call of Call.t
@@ -548,7 +560,7 @@ module Mapper : sig
   val map_option : mapper:'a t -> Expression.t option -> 'a option
 
   val create
-    :  map_await:(mapper:'a t -> location:Location.t -> Expression.t -> 'a) ->
+    :  map_await:(mapper:'a t -> location:Location.t -> Await.t -> 'a) ->
     map_binary_operator:(mapper:'a t -> location:Location.t -> BinaryOperator.t -> 'a) ->
     map_boolean_operator:(mapper:'a t -> location:Location.t -> BooleanOperator.t -> 'a) ->
     map_call:(mapper:'a t -> location:Location.t -> Call.t -> 'a) ->
@@ -579,7 +591,7 @@ module Mapper : sig
     'a t
 
   val create_default
-    :  ?map_await:(mapper:Expression.t t -> location:Location.t -> Expression.t -> Expression.t) ->
+    :  ?map_await:(mapper:Expression.t t -> location:Location.t -> Await.t -> Expression.t) ->
     ?map_binary_operator:
       (mapper:Expression.t t -> location:Location.t -> BinaryOperator.t -> Expression.t) ->
     ?map_boolean_operator:
@@ -621,7 +633,7 @@ module Mapper : sig
     Expression.t t
 
   val create_transformer
-    :  ?map_await:(mapper:Expression.t t -> Expression.t -> Expression.t) ->
+    :  ?map_await:(mapper:Expression.t t -> Await.t -> Await.t) ->
     ?map_binary_operator:(mapper:Expression.t t -> BinaryOperator.t -> BinaryOperator.t) ->
     ?map_boolean_operator:(mapper:Expression.t t -> BooleanOperator.t -> BooleanOperator.t) ->
     ?map_call:(mapper:Expression.t t -> Call.t -> Call.t) ->
@@ -667,7 +679,7 @@ module Folder : sig
   val fold_option : folder:'a t -> state:'a -> Expression.t option -> 'a
 
   val create
-    :  ?fold_await:(folder:'a t -> state:'a -> location:Location.t -> Expression.t -> 'a) ->
+    :  ?fold_await:(folder:'a t -> state:'a -> location:Location.t -> Await.t -> 'a) ->
     ?fold_binary_operator:(folder:'a t -> state:'a -> location:Location.t -> BinaryOperator.t -> 'a) ->
     ?fold_boolean_operator:
       (folder:'a t -> state:'a -> location:Location.t -> BooleanOperator.t -> 'a) ->
@@ -706,7 +718,7 @@ module Folder : sig
     'a t
 
   val create_with_uniform_location_fold
-    :  ?fold_await:(folder:'a t -> state:'a -> Expression.t -> 'a) ->
+    :  ?fold_await:(folder:'a t -> state:'a -> Await.t -> 'a) ->
     ?fold_binary_operator:(folder:'a t -> state:'a -> BinaryOperator.t -> 'a) ->
     ?fold_boolean_operator:(folder:'a t -> state:'a -> BooleanOperator.t -> 'a) ->
     ?fold_call:(folder:'a t -> state:'a -> Call.t -> 'a) ->