Add a variable callback which corresponds to get_type_alias

Summary: This diff contains several refactors 1- get_type_alias now returns an optional type instead of a RawAlias 2- we add get_variable which returns a Type.Variable.t option as a callback 3- we add plumbing for modify_variables The next step will be to actually use modify_variables as we parse type variables as variable declarations. Reviewed By: stroxler Differential Revision: D60144901 fbshipit-source-id: 3863421f77a5c2df4adf287eee1de92e7e52223c
facebook · Jul 24, 2024 · 9940b72 · 9940b72
1 parent 7438d40
commit 9940b72
Show file tree

Hide file tree

Showing 10 changed files with 187 additions and 231 deletions.
diff --git a/source/analysis/attributeResolution.ml b/source/analysis/attributeResolution.ml
@@ -58,10 +58,9 @@ module Queries = struct
     exists_matching_class_decorator: names:string list -> ClassSummary.t Ast.Node.t -> bool;
     class_exists: string -> bool;
     parse_annotation_without_validating_type_parameters:
-      ?modify_aliases:
-        (?replace_unbound_parameters_with_any:bool ->
-        TypeAliasEnvironment.RawAlias.t ->
-        TypeAliasEnvironment.RawAlias.t) ->
+      ?modify_aliases:(?replace_unbound_parameters_with_any:bool -> Type.t -> Type.t) ->
+      ?modify_variables:
+        (?replace_unbound_parameters_with_any:bool -> Type.Variable.t -> Type.Variable.t) ->
       ?allow_untracked:bool ->
       Ast.Expression.t ->
       Type.t;
@@ -2471,13 +2470,17 @@ class base ~queries:(Queries.{ controls; _ } as queries) =
         expression =
       let { Queries.parse_annotation_without_validating_type_parameters; _ } = queries in
       let modify_aliases ?replace_unbound_parameters_with_any = function
-        | TypeAliasEnvironment.RawAlias.TypeAlias alias ->
+        | alias ->
             self#check_invalid_type_parameters
               ?replace_unbound_parameters_with_any
               alias
               ~assumptions
             |> snd
-            |> fun alias -> TypeAliasEnvironment.RawAlias.TypeAlias alias
+            |> fun alias -> alias
+      in
+      let modify_variables ?replace_unbound_parameters_with_any =
+        let _replace = replace_unbound_parameters_with_any in
+        function
         | result -> result
       in
       let allow_untracked =
@@ -2490,6 +2493,7 @@ class base ~queries:(Queries.{ controls; _ } as queries) =
       let annotation =
         parse_annotation_without_validating_type_parameters
           ~modify_aliases
+          ~modify_variables
           ~allow_untracked
           expression
       in
@@ -3485,12 +3489,10 @@ class base ~queries:(Queries.{ controls; _ } as queries) =
                           annotation
                       | normal, had_descriptors, _ -> Type.union (normal @ had_descriptors)
                     in
-
                     let elements_and_get_results =
                       List.map elements ~f:(fun element ->
                           element, get_descriptor_method element ~kind:`DunderGet)
                     in
-
                     let get_type = List.unzip elements_and_get_results |> snd |> collect in
                     let set_type =
                       if accessed_through_class then

diff --git a/source/analysis/classHierarchyEnvironment.ml b/source/analysis/classHierarchyEnvironment.ml
@@ -30,14 +30,13 @@ module IncomingDataComputation = struct
       get_class_summary: string -> ClassSummary.t Ast.Node.t option;
       is_from_empty_stub: Ast.Reference.t -> bool;
       get_type_alias:
-        ?replace_unbound_parameters_with_any:bool ->
-        Type.Primitive.t ->
-        TypeAliasEnvironment.RawAlias.t option;
+        ?replace_unbound_parameters_with_any:bool -> Type.Primitive.t -> Type.t option;
+      get_variable:
+        ?replace_unbound_parameters_with_any:bool -> Type.Primitive.t -> Type.Variable.t option;
       parse_annotation_without_validating_type_parameters:
-        ?modify_aliases:
-          (?replace_unbound_parameters_with_any:bool ->
-          TypeAliasEnvironment.RawAlias.t ->
-          TypeAliasEnvironment.RawAlias.t) ->
+        ?modify_aliases:(?replace_unbound_parameters_with_any:bool -> Type.t -> Type.t) ->
+        ?modify_variables:
+          (?replace_unbound_parameters_with_any:bool -> Type.Variable.t -> Type.Variable.t) ->
         ?allow_untracked:bool ->
         Ast.Expression.t ->
         Type.t;
@@ -228,11 +227,7 @@ module IncomingDataComputation = struct
           extract_supertype (Type.expression base)
           >>= fun (name, parameters) -> Some { ClassHierarchy.Target.target = name; parameters }
         in
-        let resolved_aliases ?replace_unbound_parameters_with_any:_ name =
-          match get_type_alias name with
-          | Some (TypeAliasEnvironment.RawAlias.TypeAlias t) -> Some t
-          | _ -> None
-        in
+        let resolved_aliases ?replace_unbound_parameters_with_any:_ name = get_type_alias name in
         let has_placeholder_stub_parent =
           compute_extends_placeholder_stub_class
             class_summary
@@ -310,6 +305,7 @@ module Edges = Environment.EnvironmentTable.WithCache (struct
             EmptyStubEnvironment.ReadOnly.is_from_empty_stub empty_stub_environment ?dependency;
           get_type_alias =
             TypeAliasEnvironment.ReadOnly.get_type_alias alias_environment ?dependency;
+          get_variable = TypeAliasEnvironment.ReadOnly.get_variable alias_environment ?dependency;
           parse_annotation_without_validating_type_parameters =
             TypeAliasEnvironment.ReadOnly.parse_annotation_without_validating_type_parameters
               alias_environment
@@ -318,18 +314,7 @@ module Edges = Environment.EnvironmentTable.WithCache (struct
     in
 
     let variable_aliases name =
-      match queries.get_type_alias ?replace_unbound_parameters_with_any:(Some true) name with
-      | Some (TypeAliasEnvironment.RawAlias.VariableAlias variable) ->
-          let type_variables =
-            Type.Variable.of_declaration
-              ~create_type:
-                (Type.create
-                   ~aliases:Type.resolved_empty_aliases
-                   ~variables:Type.resolved_empty_variables)
-              variable
-          in
-          Some type_variables
-      | _ -> None
+      queries.get_variable ?replace_unbound_parameters_with_any:(Some true) name
     in
 
     IncomingDataComputation.get_parents queries key ~variable_aliases

diff --git a/source/analysis/globalResolution.ml b/source/analysis/globalResolution.ml
@@ -169,6 +169,10 @@ let get_type_alias ({ dependency; _ } as resolution) =
   TypeAliasEnvironment.ReadOnly.get_type_alias ?dependency (alias_environment resolution)
 
 
+let get_variable ({ dependency; _ } as resolution) =
+  TypeAliasEnvironment.ReadOnly.get_variable ?dependency (alias_environment resolution)
+
+
 let parse_annotation_without_validating_type_parameters ({ dependency; _ } as resolution) =
   TypeAliasEnvironment.ReadOnly.parse_annotation_without_validating_type_parameters
     ?dependency
@@ -351,7 +355,7 @@ let immediate_parents resolution = ClassHierarchy.immediate_parents (class_hiera
 let base_is_from_placeholder_stub variable_aliases resolution =
   let resolved_aliases ?replace_unbound_parameters_with_any:_ name =
     match (get_type_alias resolution) name with
-    | Some (TypeAliasEnvironment.RawAlias.TypeAlias t) -> Some t
+    | Some t -> Some t
     | _ -> None
   in
   AnnotatedBases.base_is_from_placeholder_stub

diff --git a/source/analysis/globalResolution.mli b/source/analysis/globalResolution.mli
@@ -59,14 +59,19 @@ val get_type_alias
   :  t ->
   ?replace_unbound_parameters_with_any:bool ->
   Type.Primitive.t ->
-  TypeAliasEnvironment.RawAlias.t option
+  Type.t option
+
+val get_variable
+  :  t ->
+  ?replace_unbound_parameters_with_any:bool ->
+  Type.Primitive.t ->
+  Type.Variable.t option
 
 val parse_annotation_without_validating_type_parameters
   :  t ->
-  ?modify_aliases:
-    (?replace_unbound_parameters_with_any:bool ->
-    TypeAliasEnvironment.RawAlias.t ->
-    TypeAliasEnvironment.RawAlias.t) ->
+  ?modify_aliases:(?replace_unbound_parameters_with_any:bool -> Type.t -> Type.t) ->
+  ?modify_variables:
+    (?replace_unbound_parameters_with_any:bool -> Type.Variable.t -> Type.Variable.t) ->
   ?allow_untracked:bool ->
   Expression.t ->
   Type.t

diff --git a/source/analysis/resolution.ml b/source/analysis/resolution.ml
@@ -333,23 +333,14 @@ let is_consistent_with ({ global_resolution; _ } as resolution) =
 
 let global_resolution { global_resolution; _ } = global_resolution
 
-let get_type_alias resolution =
+let get_variable resolution =
   let global_resolution = global_resolution resolution in
-  GlobalResolution.get_type_alias global_resolution
+  GlobalResolution.get_variable global_resolution
 
 
 let variables resolution name =
-  match get_type_alias ?replace_unbound_parameters_with_any:(Some true) resolution name with
-  | Some (TypeAliasEnvironment.RawAlias.VariableAlias variable) ->
-      let type_variables =
-        Type.Variable.of_declaration
-          ~create_type:
-            (Type.create
-               ~aliases:Type.resolved_empty_aliases
-               ~variables:Type.resolved_empty_variables)
-          variable
-      in
-      Some type_variables
+  match get_variable ?replace_unbound_parameters_with_any:(Some true) resolution name with
+  | Some variable -> Some variable
   | _ -> None
 
 

diff --git a/source/analysis/test/constraintsSetTest.ml b/source/analysis/test/constraintsSetTest.ml
@@ -16,22 +16,6 @@ open Assumptions
 
 let ( ! ) concretes = List.map concretes ~f:(fun single -> Type.Parameter.Single single)
 
-let variable_aliases aliases name =
-  match aliases ?replace_unbound_parameters_with_any:(Some true) name with
-  | Some (TypeAliasEnvironment.RawAlias.VariableAlias variable) ->
-      let type_variables =
-        Type.Variable.of_declaration
-          ~create_type:
-            (Type.create
-               ~aliases:Type.resolved_empty_aliases
-               ~variables:Type.resolved_empty_variables)
-          variable
-      in
-
-      Some type_variables
-  | _ -> None
-
-
 let make_attributes ~class_name =
   let parse_attribute (name, annotation) =
     AnnotatedAttribute.create
@@ -170,58 +154,53 @@ let make_assert_functions context =
   let resolution = GlobalResolution.create environment in
   let default_postprocess annotation = Type.Variable.mark_all_variables_as_bound annotation in
   let prep annotation =
+    let s =
+      [
+        "Base";
+        "Child";
+        "Unrelated";
+        "T_Unconstrained";
+        "T_Bound_Base";
+        "T_Bound_Child";
+        "T_Bound_Union_Base_Unrelated";
+        "T_Bound_Union";
+        "T_Bound_ReadOnly";
+        "T_Bound_object";
+        "T_Base_Unrelated";
+        "T_Child_Unrelated";
+        "T_Base_Unrelated_int";
+        "V";
+        "P";
+        "P2";
+        "T";
+        "T1";
+        "T2";
+        "T3";
+        "T4";
+        "G_invariant";
+        "G_covariant";
+        "T_Baseovariant";
+        "ClassWithOverloadedConstructor";
+        "Constructable";
+        "UserDefinedVariadic";
+        "UserDefinedVariadicSimpleChild";
+        "UserDefinedVariadicMapChild";
+        "Ts";
+        "Ts2";
+        "Tensor";
+      ]
+      |> Type.Primitive.Set.of_list
+    in
     let aliases ?replace_unbound_parameters_with_any:_ a =
-      let s =
-        [
-          "Base";
-          "Child";
-          "Unrelated";
-          "T_Unconstrained";
-          "T_Bound_Base";
-          "T_Bound_Child";
-          "T_Bound_Union_Base_Unrelated";
-          "T_Bound_Union";
-          "T_Bound_ReadOnly";
-          "T_Bound_object";
-          "T_Base_Unrelated";
-          "T_Child_Unrelated";
-          "T_Base_Unrelated_int";
-          "V";
-          "P";
-          "P2";
-          "T";
-          "T1";
-          "T2";
-          "T3";
-          "T4";
-          "G_invariant";
-          "G_covariant";
-          "T_Baseovariant";
-          "ClassWithOverloadedConstructor";
-          "Constructable";
-          "UserDefinedVariadic";
-          "UserDefinedVariadicSimpleChild";
-          "UserDefinedVariadicMapChild";
-          "Ts";
-          "Ts2";
-          "Tensor";
-        ]
-        |> Type.Primitive.Set.of_list
-      in
       if Set.mem s a then
-        Some (TypeAliasEnvironment.RawAlias.TypeAlias (Type.Primitive ("test." ^ a)))
+        Some (Type.Primitive ("test." ^ a))
       else
         GlobalResolution.get_type_alias resolution a
     in
-    let resolved_aliases ?replace_unbound_parameters_with_any:_ name =
-      match aliases name with
-      | Some (TypeAliasEnvironment.RawAlias.TypeAlias t) -> Some t
-      | _ -> None
+    let variables ?replace_unbound_parameters_with_any:_ a =
+      GlobalResolution.get_variable resolution a
     in
-
-    annotation
-    |> Type.create ~variables:(variable_aliases aliases) ~aliases:resolved_aliases
-    |> Type.expression
+    annotation |> Type.create ~variables ~aliases |> Type.expression
   in
   let parse_annotation annotation ~do_prep =
     annotation
@@ -310,13 +289,9 @@ let make_assert_functions context =
                 | _ -> failwith "expected tuple"
               in
               let global_resolution = GlobalResolution.create environment in
-              match GlobalResolution.get_type_alias global_resolution primitive with
-              | Some (TypeAliasEnvironment.RawAlias.VariableAlias variable_declaration) -> (
-                  match
-                    Type.Variable.of_declaration
-                      ~create_type:(GlobalResolution.parse_annotation global_resolution)
-                      variable_declaration
-                  with
+              match GlobalResolution.get_variable global_resolution primitive with
+              | Some variable_declaration -> (
+                  match variable_declaration with
                   | ParamSpecVariable variable ->
                       Type.Variable.ParamSpecPair (variable, parse_parameters value)
                   | TypeVarTupleVariable variable -> (

diff --git a/source/analysis/test/environmentTest.ml b/source/analysis/test/environmentTest.ml
@@ -377,8 +377,8 @@ let test_register_aliases context =
     let environment = register_all sources in
     let global_resolution = GlobalResolution.create environment in
     let assert_alias (alias, target) =
-      match GlobalResolution.get_type_alias global_resolution alias with
-      | Some alias -> assert_equal ~printer:TypeAliasEnvironment.RawAlias.show target alias
+      match GlobalResolution.get_variable global_resolution alias with
+      | Some alias -> assert_equal ~printer:Type.Variable.show target alias
       | None -> failwith "Alias is missing"
     in
     List.iter aliases ~f:assert_alias
@@ -387,11 +387,7 @@ let test_register_aliases context =
     ["test.py", {|
           Tparams = pyre_extensions.ParameterSpecification('Tparams')
       |}]
-    [
-      ( "test.Tparams",
-        TypeAliasEnvironment.RawAlias.VariableAlias
-          (Type.Variable.Declaration.DParamSpec { name = "test.Tparams" }) );
-    ];
+    ["test.Tparams", ParamSpecVariable (Type.Variable.ParamSpec.create "test.Tparams")];
   ()