parser support for PEP695 type aliases

Summary:
this diff adds a TypeAlias statement node for the PEP 695 type alias syntax.

some of the handling is stubbed out; locations marked with T196994965 need to be revisited and implemented later.

Reviewed By: samwgoldman

Differential Revision: D60260862

fbshipit-source-id: 980b87c8097722d99fc9416be50af41a0de4c27f

source/analysis/functionDefinition.ml

@@ -91,6 +91,7 @@ let collect_typecheck_units { Source.statements; _ } =
     | Nonlocal _
     | Pass
     | Raise _
+    | TypeAlias _
     | Return _ ->
         sofar
   in

source/analysis/globalLeakCheck.ml

@@ -659,6 +659,7 @@ module State (Context : Context) = struct
       | Global _
       | Import _
       | Nonlocal _
+      | TypeAlias _
       | Pass ->
           []
     in

source/analysis/preprocessing.ml

@@ -923,6 +923,7 @@ module Qualify (Context : QualifyContext) = struct
                 body;
                 orelse;
               } )
+      | TypeAlias _ (* TODO(yangdanny): handle TypeAlias? *)
       | Break
       | Continue
       | Import _
@@ -3317,6 +3318,7 @@ module AccessCollector = struct
     | Global _
     | Import _
     | Nonlocal _
+    | TypeAlias _ (* TODO(T196994965): handle TypeAlias *)
     | Pass ->
         collected
 

source/analysis/scope.ml

@@ -304,6 +304,7 @@ module Binding = struct
     | Delete _
     | Global _
     | Nonlocal _
+    | TypeAlias _ (* TODO(T196994965): handle TypeAlias *)
     | Pass ->
         sofar
 
@@ -372,7 +373,8 @@ let rec globals_of_statement sofar { Node.value; _ } =
   | Nonlocal _
   | Pass
   | Raise _
-  | Return _ ->
+  | Return _
+  | TypeAlias _ ->
       sofar
 
 
@@ -414,7 +416,8 @@ let rec nonlocals_of_statement sofar { Node.value; _ } =
   | Import _
   | Pass
   | Raise _
-  | Return _ ->
+  | Return _
+  | TypeAlias _ ->
       sofar
 
 

source/analysis/typeCheck.ml

@@ -5596,6 +5596,24 @@ module State (Context : Context) = struct
         (* lower augmented assignment to regular assignment *)
         let call = AugmentedAssign.lower ~location augmented_assignment in
         forward_assignment ~resolution ~location ~target ~annotation:None ~value:(Some call)
+    (* TODO(T196994965): handle type alias *)
+    | TypeAlias { TypeAlias.name; type_params; value } ->
+        (* TODO: remove after PEP 695 is supported *)
+        let type_params_errors =
+          match type_params with
+          | { Node.location; _ } :: _ ->
+              emit_error
+                ~errors:[]
+                ~location
+                ~kind:(Error.ParserFailure "PEP 695 type params are unsupported")
+          | _ -> []
+        in
+        forward_type_alias_definition
+          ~resolution
+          ~location
+          ~errors:type_params_errors
+          ~target:name
+          ~value
     | Assert { Assert.test; origin; message } ->
         let message_errors =
           Option.value

source/analysis/unawaitedAwaitableCheck.ml

@@ -1032,6 +1032,7 @@ module State (Context : Context) = struct
     | Global _
     | Import _
     | Nonlocal _
+    | TypeAlias _
     | Pass ->
         state
 

source/analysis/uninitializedLocalCheck.ml

@@ -210,6 +210,7 @@ let extract_reads_in_statement { Node.value; _ } =
     (* TODO(T107105911): Handle access for match statement. *)
     | Match _
     | Nonlocal _
+    | TypeAlias _
     | Pass
     | Try _ ->
         []

source/ast/statement.ml

@@ -31,6 +31,28 @@ let name_location
     }
 
 
+module TypeAlias = struct
+  type t = {
+    name: Expression.t;
+    type_params: Expression.TypeParam.t list;
+    value: Expression.t;
+  }
+  [@@deriving equal, compare, sexp, show, hash, to_yojson]
+
+  let location_insensitive_compare left right =
+    match Expression.location_insensitive_compare left.name right.name with
+    | x when not (Int.equal x 0) -> x
+    | _ -> (
+        match
+          List.compare
+            Expression.TypeParam.location_insensitive_compare
+            left.type_params
+            right.type_params
+        with
+        | x when not (Int.equal x 0) -> x
+        | _ -> Expression.location_insensitive_compare left.value right.value)
+end
+
 module Assign = struct
   type t = {
     target: Expression.t;
@@ -1574,6 +1596,7 @@ and Statement : sig
     | Raise of Raise.t
     | Return of Return.t
     | Try of Try.t
+    | TypeAlias of TypeAlias.t
     | With of With.t
     | While of While.t
   [@@deriving equal, compare, sexp, hash, to_yojson]
@@ -1606,6 +1629,7 @@ end = struct
     | Raise of Raise.t
     | Return of Return.t
     | Try of Try.t
+    | TypeAlias of TypeAlias.t
     | With of With.t
     | While of While.t
   [@@deriving equal, compare, sexp, show, hash, to_yojson]
@@ -1634,6 +1658,7 @@ end = struct
     | Raise left, Raise right -> Raise.location_insensitive_compare left right
     | Return left, Return right -> Return.location_insensitive_compare left right
     | Try left, Try right -> Try.location_insensitive_compare left right
+    | TypeAlias left, TypeAlias right -> TypeAlias.location_insensitive_compare left right
     | With left, With right -> With.location_insensitive_compare left right
     | While left, While right -> While.location_insensitive_compare left right
     | Assign _, _ -> -1
@@ -1655,6 +1680,7 @@ end = struct
     | Raise _, _ -> -1
     | Return _, _ -> -1
     | Try _, _ -> -1
+    | TypeAlias _, _ -> -1
     | With _, _ -> -1
     | While _, _ -> -1
 
@@ -2018,6 +2044,16 @@ module PrettyPrinter = struct
           orelse
           pp_finally_block
           finally
+    | TypeAlias { TypeAlias.name; type_params; value } ->
+        Format.fprintf
+          formatter
+          "type %a%a = %a"
+          Expression.pp
+          name
+          Expression.pp_type_param_list
+          type_params
+          Expression.pp
+          value
     | With { With.items; body; async } ->
         let pp_item formatter (expression, expression_option) =
           Format.fprintf
@@ -2168,6 +2204,8 @@ let is_generator statements =
         || List.exists handlers ~f:(fun { Try.Handler.body; _ } -> is_statements_generator body)
         || is_statements_generator orelse
         || is_statements_generator finally
+    | TypeAlias { TypeAlias.name; value; _ } ->
+        is_expression_generator name || is_expression_generator value
     | With { With.items; body; _ } ->
         List.exists items ~f:(fun (expression, _) -> is_expression_generator expression)
         || is_statements_generator body

source/ast/statement.mli

@@ -33,6 +33,17 @@ module AugmentedAssign : sig
   val lower : location:Location.t -> t -> Expression.t
 end
 
+module TypeAlias : sig
+  type t = {
+    name: Expression.t;
+    type_params: Expression.TypeParam.t list;
+    value: Expression.t;
+  }
+  [@@deriving equal, compare, sexp, show, hash, to_yojson]
+
+  val location_insensitive_compare : t -> t -> int
+end
+
 module Import : sig
   type import = {
     name: Reference.t;
@@ -453,6 +464,7 @@ and Statement : sig
     | Raise of Raise.t
     | Return of Return.t
     | Try of Try.t
+    | TypeAlias of TypeAlias.t
     | With of With.t
     | While of While.t
   [@@deriving equal, compare, sexp, hash, to_yojson]

source/ast/transform.ml

@@ -69,6 +69,15 @@ module Make (Transformer : Transformer) = struct
     let transform_argument { Call.Argument.name; value } ~transform_expression =
       { Call.Argument.name; value = transform_expression value }
     in
+    let transform_type_param ({ Node.value; _ } as node) ~transform_expression =
+      let new_value =
+        match value with
+        | TypeParam.TypeVar { TypeParam.name; bound = Some bound } ->
+            TypeParam.TypeVar { TypeParam.name; bound = Some (transform_expression bound) }
+        | _ -> value
+      in
+      { node with Node.value = new_value }
+    in
     let transform_parameter
         ({ Node.value = { Parameter.name; value; annotation }; _ } as node)
         ~transform_expression
@@ -281,7 +290,8 @@ module Make (Transformer : Transformer) = struct
                 body = transform_list body ~f:transform_statement |> List.concat;
                 decorators = transform_list decorators ~f:transform_expression;
                 top_level_unbound_names;
-                type_params;
+                type_params =
+                  transform_list type_params ~f:(transform_type_param ~transform_expression);
               }
         | Continue -> value
         | Define { Define.signature; captures; unbound_names; body } ->
@@ -308,7 +318,8 @@ module Make (Transformer : Transformer) = struct
                 parent;
                 nesting_define;
                 generator;
-                type_params;
+                type_params =
+                  transform_list type_params ~f:(transform_type_param ~transform_expression);
               }
             in
             let transform_capture { Define.Capture.name; kind } =
@@ -424,6 +435,14 @@ module Make (Transformer : Transformer) = struct
             let orelse = transform_list orelse ~f:transform_statement |> List.concat in
             let finally = transform_list finally ~f:transform_statement |> List.concat in
             Try { Try.body; handlers; orelse; finally; handles_exception_group }
+        | TypeAlias { TypeAlias.name; type_params; value } ->
+            TypeAlias
+              {
+                TypeAlias.name = transform_expression name;
+                type_params =
+                  transform_list type_params ~f:(transform_type_param ~transform_expression);
+                value = transform_expression value;
+              }
         | With { With.items; body; async } ->
             let transform_item (item, alias) =
               transform_expression item, alias >>| transform_expression
@@ -488,6 +507,7 @@ module MakeStatementTransformer (Transformer : StatementTransformer) = struct
         | Pass
         | Raise _
         | Return _
+        | TypeAlias _
         | Nonlocal _ ->
             value
         | Class ({ Class.body; _ } as value) ->

source/ast/visit.ml

@@ -62,6 +62,12 @@ module MakeNodeVisitor (Visitor : NodeVisitor) = struct
     Option.iter ~f:visit_expression annotation
 
 
+  let visit_type_param { Node.value; _ } ~visit_expression =
+    match value with
+    | TypeParam.TypeVar { TypeParam.bound; _ } -> Option.iter ~f:visit_expression bound
+    | _ -> ()
+
+
   let rec visit_expression ~state ?visitor_override expression =
     let visitor = Option.value visitor_override ~default:Visitor.node in
     let visit_expression = visit_expression ~state ?visitor_override in
@@ -183,23 +189,26 @@ module MakeNodeVisitor (Visitor : NodeVisitor) = struct
       | AugmentedAssign { AugmentedAssign.target; value; _ } ->
           visit_expression target;
           visit_expression value
-      | Class ({ Class.name; base_arguments; body; decorators; _ } as class_) ->
+      | Class ({ Class.name; base_arguments; body; decorators; type_params; _ } as class_) ->
           visit_node
             ~state
             ~visitor
             (Reference
                (Node.create ~location:(Class.name_location ~body_location:location class_) name));
+          List.iter type_params ~f:(visit_type_param ~visit_expression);
           List.iter base_arguments ~f:(visit_argument ~visit_expression);
           List.iter body ~f:visit_statement;
           List.iter ~f:visit_expression decorators
       | Define ({ Define.signature; captures; body; unbound_names = _ } as define) ->
-          let iter_signature { Define.Signature.name; parameters; decorators; return_annotation; _ }
+          let iter_signature
+              { Define.Signature.name; parameters; decorators; return_annotation; type_params; _ }
             =
             visit_node
               ~state
               ~visitor
               (Reference
                  (Node.create ~location:(Define.name_location ~body_location:location define) name));
+            List.iter type_params ~f:(visit_type_param ~visit_expression);
             List.iter parameters ~f:(visit_parameter ~state ~visitor ~visit_expression);
             List.iter ~f:visit_expression decorators;
             Option.iter ~f:visit_expression return_annotation
@@ -277,6 +286,10 @@ module MakeNodeVisitor (Visitor : NodeVisitor) = struct
           visit_expression test;
           List.iter body ~f:visit_statement;
           List.iter orelse ~f:visit_statement
+      | TypeAlias { TypeAlias.name; type_params; value; _ } ->
+          visit_expression name;
+          List.iter type_params ~f:(visit_type_param ~visit_expression);
+          visit_expression value
       | Import _
       | Nonlocal _
       | Global _
@@ -352,6 +365,7 @@ module MakeStatementVisitor (Visitor : StatementVisitor) = struct
         | Pass
         | Raise _
         | Return _
+        | TypeAlias _
         | Nonlocal _ ->
             ()
         | Class { Class.body; _ }

source/cpython_parser/pyreCPythonParser.ml

@@ -889,9 +889,8 @@ let statement =
     in
     List.map targets ~f:create_assign_for_target
   in
-  let type_alias ~location ~name ~type_params ~value ~context =
-    ignore type_params (* TODO(yangdanny) support 695 syntax for type aliases *);
-    assign ~location ~targets:[name] ~value ~type_comment:None ~context
+  let type_alias ~location ~name ~type_params ~value ~context:_ =
+    [Statement.TypeAlias { TypeAlias.name; type_params; value } |> Ast.Node.create ~location]
   in
   let aug_assign ~location ~target ~op ~value ~context:_ =
     [

source/cpython_parser/test/statementTest.ml

@@ -3030,6 +3030,43 @@ let test_match =
   PyreCPythonParser.with_context do_test
 
 
+let test_type_alias =
+  let do_test context =
+    let assert_parsed = assert_parsed ~context in
+    test_list
+      [
+        labeled_test_case __FUNCTION__ __LINE__
+        @@ assert_parsed
+             "type A = B"
+             ~expected:
+               [+Statement.TypeAlias { TypeAlias.name = !"A"; type_params = []; value = !"B" }];
+        labeled_test_case __FUNCTION__ __LINE__
+        @@ assert_parsed
+             "type A[T, T2: str, *Ts, **P] = B"
+             ~expected:
+               [
+                 +Statement.TypeAlias
+                    {
+                      TypeAlias.name = !"A";
+                      type_params =
+                        [
+                          +TypeParam.TypeVar { name = "T"; bound = None };
+                          +TypeParam.TypeVar
+                             {
+                               name = "T2";
+                               bound = Some (+Expression.Name (Name.Identifier "str"));
+                             };
+                          +TypeParam.TypeVarTuple "Ts";
+                          +TypeParam.ParamSpec "P";
+                        ];
+                      value = !"B";
+                    };
+               ];
+      ]
+  in
+  PyreCPythonParser.with_context do_test
+
+
 let () =
   "parse_statements"
   >::: [
@@ -3045,5 +3082,6 @@ let () =
          test_define;
          test_class;
          test_match;
+         test_type_alias;
        ]
   |> Test.run