Add origins to typing.Union nodes created from the union shorthand a | b

Summary: We turn `a | b` into `typing.Union[a, b]` during preprocessing. Again, those nodes need an origin, since they are not in the original code.

Reviewed By: tianhan0

Differential Revision: D74739793

fbshipit-source-id: 31d8b15c7fe9bfa1f3c6657620b53b2718a892c0

Author: arthaud

source/analysis/preprocessing.ml

@@ -3889,7 +3889,7 @@ let replace_union_shorthand_in_annotation_expression =
     in
     let value =
       match value with
-      | Expression.BinaryOperator { operator = BinaryOperator.BitOr; left; right; origin = _ } ->
+      | Expression.BinaryOperator { operator = BinaryOperator.BitOr; left; right; origin } ->
           let indices =
             [left; right]
             (* Recursively transform them into `typing.Union[...]` form *)
@@ -3899,6 +3899,7 @@ let replace_union_shorthand_in_annotation_expression =
             |> List.rev
           in
           let index = { Node.value = Expression.Tuple indices; location } in
+          let origin = Some (Origin.create ?base:origin ~location Origin.UnionShorthand) in
           Expression.Subscript
             {
               base =
@@ -3910,11 +3911,11 @@ let replace_union_shorthand_in_annotation_expression =
                          {
                            base = { Node.location; value = Name (Name.Identifier "typing") };
                            attribute = "Union";
-                           origin = None;
+                           origin;
                          });
                 };
               index;
-              origin = None;
+              origin;
             }
       | Subscript { Subscript.base; index; origin } ->
           Subscript { base; index = transform_expression index; origin }

source/ast/expression.ml

@@ -1179,6 +1179,7 @@ and Origin : sig
     | InGetItem (* `e in l` can be turned into `l.__getitem__(0).__eq__(e)` *)
     | InGetItemEq (* `e in l` can be turned into `l.__getitem__(0).__eq__(e)` *)
     | Slice (* `1:2` is turned into `slice(1,2,None)` *)
+    | UnionShorthand (* `a | b` is turned into `typing.Union[a, b]` when in typing context *)
     | Negate (* `if cond:` is turned into `assert(cond)` and `assert(not cond)` *)
     | NegateIs (* `not(a is not b)` is turned into `a is b` *)
     | NegateIsNot (* `not(a is b)` is turned into `a is not b` *)
@@ -1287,6 +1288,7 @@ end = struct
     | InGetItem
     | InGetItemEq
     | Slice
+    | UnionShorthand
     | Negate
     | NegateIs
     | NegateIsNot

source/ast/expression.mli

@@ -428,6 +428,7 @@ and Origin : sig
     | InGetItem (* `e in l` can be turned into `l.__getitem__(0).__eq__(e)` *)
     | InGetItemEq (* `e in l` can be turned into `l.__getitem__(0).__eq__(e)` *)
     | Slice (* `1:2` is turned into `slice(1,2,None)` *)
+    | UnionShorthand (* `a | b` is turned into `typing.Union[a, b]` when in typing context *)
     | Negate (* `if cond:` is turned into `assert(cond)` and `assert(not cond)` *)
     | NegateIs (* `not(a is not b)` is turned into `a is b` *)
     | NegateIsNot (* `not(a is b)` is turned into `a is not b` *)