Export the explicit annotation of class attributes to pysa

Summary:
We need the explicit annotation of class attributes to power the `type_annotation.original.matches/.equals` constraint in our model queries.
Let's export it from pyrefly.

Reviewed By: tianhan0

Differential Revision: D82984172

fbshipit-source-id: d0bac5af6a9c786d3d95db7f5287ac70c9316ecc

Author: arthaud

source/interprocedural/pyreflyApi.ml

@@ -843,6 +843,7 @@ module ModuleInfoFile = struct
     type t = {
       name: string;
       type_: JsonType.t;
+      explicit_annotation: string option;
       location: Location.t option;
     }
     [@@deriving equal, show]
@@ -852,11 +853,13 @@ module ModuleInfoFile = struct
       JsonUtil.get_member json "type"
       >>= JsonType.from_json
       >>= fun type_ ->
+      JsonUtil.get_optional_string_member json "explicit_annotation"
+      >>= fun explicit_annotation ->
       JsonUtil.get_optional_string_member json "location"
       >>= (function
             | Some location -> parse_location location >>| Option.some
             | None -> Ok None)
-      >>| fun location -> { name; type_; location }
+      >>| fun location -> { name; type_; explicit_annotation; location }
   end
 
   module ClassDefinition = struct
@@ -1335,6 +1338,7 @@ module CallableIdToQualifiedNameSharedMemory =
 module ClassField = struct
   type t = {
     type_: PysaType.t;
+    explicit_annotation: string option;
     location: Location.t option;
   }
   [@@deriving equal, compare, show]
@@ -2435,8 +2439,13 @@ module ReadWrite = struct
               };
             let fields =
               fields
-              |> List.map ~f:(fun { ModuleInfoFile.JsonClassField.name; type_; location } ->
-                     name, { ClassField.type_ = create_pysa_type type_; location })
+              |> List.map
+                   ~f:(fun
+                        { ModuleInfoFile.JsonClassField.name; type_; explicit_annotation; location }
+                      ->
+                     ( name,
+                       { ClassField.type_ = create_pysa_type type_; explicit_annotation; location }
+                     ))
               |> SerializableStringMap.of_alist_exn
             in
             ClassFieldsSharedMemory.add class_fields_shared_memory qualified_name fields;
@@ -3157,6 +3166,20 @@ module ReadOnly = struct
     |> fun { ClassField.type_; _ } -> type_
 
 
+  let get_class_attribute_explicit_annotation
+      { class_fields_shared_memory; _ }
+      ~class_name
+      ~attribute
+    =
+    ClassFieldsSharedMemory.get
+      class_fields_shared_memory
+      (FullyQualifiedName.from_reference_unchecked (Reference.create class_name))
+    |> assert_shared_memory_key_exists "missing class fields for class"
+    |> SerializableStringMap.find_opt attribute
+    |> assert_shared_memory_key_exists "missing class field"
+    |> fun { ClassField.explicit_annotation; _ } -> explicit_annotation
+
+
   let get_global_inferred_type { module_globals_shared_memory; _ } ~qualifier ~name =
     ModuleGlobalsSharedMemory.get
       module_globals_shared_memory

source/interprocedural/pyreflyApi.mli

@@ -135,6 +135,12 @@ module ReadOnly : sig
 
   val get_class_attribute_inferred_type : t -> class_name:string -> attribute:string -> PysaType.t
 
+  val get_class_attribute_explicit_annotation
+    :  t ->
+    class_name:string ->
+    attribute:string ->
+    string option
+
   val get_global_inferred_type : t -> qualifier:Ast.Reference.t -> name:string -> PysaType.t option
 
   module Type : sig
@@ -335,6 +341,7 @@ module ModuleInfoFile : sig
     type t = {
       name: string;
       type_: JsonType.t;
+      explicit_annotation: string option;
       location: Ast.Location.t option;
     }
     [@@deriving equal, show]

source/interprocedural_analyses/taint/modelParseResult.ml

@@ -877,71 +877,98 @@ end
 module TypeAnnotation : sig
   type t
 
-  val from_original_annotation
-    :  pyre_api:PyrePysaApi.ReadOnly.t ->
-    preserve_original:bool ->
-    Expression.t ->
+  module ExplicitAnnotation : sig
+    type t =
+      | Unsupported
+      | NotFound
+      | Found of string
+  end
+
+  val from_pyre1 : pyre_api:Analysis.PyrePysaEnvironment.ReadOnly.t -> Expression.t option -> t
+
+  val create
+    :  inferred_type:PyrePysaApi.PysaType.t option ->
+    explicit_annotation:ExplicitAnnotation.t ->
     t
 
-  val from_pysa_type : PyrePysaApi.PysaType.t -> t
+  val from_inferred_type : PyrePysaApi.PysaType.t option -> t
 
   val is_annotated : t -> bool
 
-  val as_type : t -> PyrePysaApi.PysaType.t
+  val inferred_type : t -> PyrePysaApi.PysaType.t option
 
-  val as_original_annotation : t -> Ast.Expression.t option
+  val explicit_annotation : t -> ExplicitAnnotation.t
 
   (* Show the original annotation, as written by the user. *)
-  val show_original_annotation : t -> string
+  val show_explicit_annotation : t -> string option
 
   (* Show the fully qualified type annotation from the type checker. *)
   val show_fully_qualified_annotation : t -> string
 end = struct
+  module ExplicitAnnotation = struct
+    type t =
+      | Unsupported
+      | NotFound
+      | Found of string
+  end
+
   type t = {
-    expression: Expression.t option;
-        (* The original annotation, as an expression. None if not supported. *)
-    type_: PyrePysaApi.PysaType.t Lazy.t;
+    explicit_annotation: ExplicitAnnotation.t Lazy.t;
+    inferred_type: PyrePysaApi.PysaType.t option Lazy.t;
   }
 
-  let from_original_annotation ~pyre_api ~preserve_original expression =
+  let from_pyre1 ~pyre_api = function
+    | Some expression ->
+        {
+          explicit_annotation = lazy (ExplicitAnnotation.Found (Expression.show expression));
+          inferred_type =
+            lazy
+              (Analysis.PyrePysaEnvironment.ReadOnly.parse_annotation pyre_api expression
+              |> PyrePysaApi.PysaType.from_pyre1_type
+              |> Option.some);
+        }
+    | None -> { explicit_annotation = lazy ExplicitAnnotation.NotFound; inferred_type = lazy None }
+
+
+  let create ~inferred_type ~explicit_annotation =
+    { explicit_annotation = lazy explicit_annotation; inferred_type = lazy inferred_type }
+
+
+  let from_inferred_type inferred_type =
     {
-      expression = Option.some_if preserve_original expression;
-      type_ =
-        lazy
-          (PyrePysaApi.ReadOnly.parse_annotation pyre_api expression
-          |> PyrePysaApi.PysaType.from_pyre1_type);
+      explicit_annotation = lazy ExplicitAnnotation.Unsupported;
+      inferred_type = lazy inferred_type;
     }
 
 
-  let from_pysa_type type_ = { expression = None; type_ = lazy type_ }
+  let is_annotated { explicit_annotation; _ } =
+    match Lazy.force explicit_annotation with
+    | ExplicitAnnotation.Unsupported -> failwith "is_annotated is not supported in this context"
+    | ExplicitAnnotation.Found annotation ->
+        (String.is_prefix ~prefix:"typing.Annotated[" annotation
+        || String.is_prefix ~prefix:"Annotated[" annotation)
+        && String.is_suffix ~suffix:"]" annotation
+    | ExplicitAnnotation.NotFound -> false
 
-  let is_annotated = function
-    | { expression = None; _ } -> failwith "is_annotated is not supported in this context"
-    | { expression = Some { Node.value = expression; _ }; _ } -> (
-        match expression with
-        | Expression.Expression.Subscript
-            {
-              base =
-                { Node.value = Name (Expression.Name.Attribute { attribute = "Annotated"; _ }); _ };
-              _;
-            } ->
-            true
-        | _ -> false)
 
+  let inferred_type { inferred_type; _ } = Lazy.force inferred_type
 
-  let as_type { type_; _ } = Lazy.force type_
-
-  let as_original_annotation { expression; _ } = expression
+  let explicit_annotation { explicit_annotation; _ } = Lazy.force explicit_annotation
 
   (* Show the original annotation, as written by the user. *)
-  let show_original_annotation = function
-    | { expression = Some expression; _ } -> Expression.show expression
-    | _ -> failwith "show_original_annotation is not supported in this context"
+  let show_explicit_annotation { explicit_annotation; _ } =
+    match Lazy.force explicit_annotation with
+    | ExplicitAnnotation.Unsupported ->
+        failwith "show_explicit_annotation is not supported in this context"
+    | ExplicitAnnotation.Found annotation -> Some annotation
+    | ExplicitAnnotation.NotFound -> None
 
 
   (* Show the parsed annotation from pyre *)
-  let show_fully_qualified_annotation { type_; _ } =
-    PyrePysaApi.PysaType.show_fully_qualified (Lazy.force type_)
+  let show_fully_qualified_annotation { inferred_type; _ } =
+    match Lazy.force inferred_type with
+    | Some inferred_type -> PyrePysaApi.PysaType.show_fully_qualified inferred_type
+    | None -> "typing.Any"
 end
 
 module Modelable = struct
@@ -957,11 +984,11 @@ module Modelable = struct
       }
     | Attribute of {
         target_name: Reference.t;
-        type_annotation: TypeAnnotation.t option Lazy.t;
+        type_annotation: TypeAnnotation.t Lazy.t;
       }
     | Global of {
         target_name: Reference.t;
-        type_annotation: TypeAnnotation.t option Lazy.t;
+        type_annotation: TypeAnnotation.t Lazy.t;
       }
 
   let create_callable ~pyre_api ~callables_to_definitions_map target =
@@ -1024,14 +1051,25 @@ module Modelable = struct
                ~include_generated_attributes:false
                ~class_name
                ~attribute
-             >>| TypeAnnotation.from_original_annotation ~pyre_api ~preserve_original:true
+             |> TypeAnnotation.from_pyre1 ~pyre_api:pyre1_api
          | PyrePysaApi.ReadOnly.Pyrefly pyrefly_api ->
-             Interprocedural.PyreflyApi.ReadOnly.get_class_attribute_inferred_type
-               pyrefly_api
-               ~class_name
-               ~attribute
-             |> TypeAnnotation.from_pysa_type
-             |> Option.some)
+             let inferred_type =
+               Interprocedural.PyreflyApi.ReadOnly.get_class_attribute_inferred_type
+                 pyrefly_api
+                 ~class_name
+                 ~attribute
+               |> Option.some
+             in
+             let explicit_annotation =
+               Interprocedural.PyreflyApi.ReadOnly.get_class_attribute_explicit_annotation
+                 pyrefly_api
+                 ~class_name
+                 ~attribute
+               |> function
+               | Some annotation -> TypeAnnotation.ExplicitAnnotation.Found annotation
+               | None -> TypeAnnotation.ExplicitAnnotation.NotFound
+             in
+             TypeAnnotation.create ~inferred_type ~explicit_annotation)
     in
     Attribute { target_name; type_annotation }
 
@@ -1043,15 +1081,19 @@ module Modelable = struct
         (match pyre_api with
         | PyrePysaApi.ReadOnly.Pyre1 pyre1_api ->
             Analysis.PyrePysaEnvironment.ReadOnly.get_global_annotation pyre1_api target_name
-            >>| TypeAnnotation.from_original_annotation ~pyre_api ~preserve_original:true
+            |> TypeAnnotation.from_pyre1 ~pyre_api:pyre1_api
         | PyrePysaApi.ReadOnly.Pyrefly pyrefly_api ->
             let qualifier = Option.value_exn (Reference.prefix target_name) in
             let name = Reference.last target_name in
-            Interprocedural.PyreflyApi.ReadOnly.get_global_inferred_type
-              pyrefly_api
-              ~qualifier
-              ~name
-            >>| TypeAnnotation.from_pysa_type)
+            let inferred_type =
+              Interprocedural.PyreflyApi.ReadOnly.get_global_inferred_type
+                pyrefly_api
+                ~qualifier
+                ~name
+            in
+            TypeAnnotation.create
+              ~inferred_type
+              ~explicit_annotation:TypeAnnotation.ExplicitAnnotation.Unsupported)
     in
     Global { target_name; type_annotation }
 

source/interprocedural_analyses/taint/modelParseResult.mli

@@ -445,22 +445,30 @@ end
 module TypeAnnotation : sig
   type t
 
-  val from_original_annotation
-    :  pyre_api:PyrePysaApi.ReadOnly.t ->
-    preserve_original:bool ->
-    Ast.Expression.t ->
+  module ExplicitAnnotation : sig
+    type t =
+      | Unsupported
+      | NotFound
+      | Found of string
+  end
+
+  val from_pyre1 : pyre_api:Analysis.PyrePysaEnvironment.ReadOnly.t -> Ast.Expression.t option -> t
+
+  val create
+    :  inferred_type:PyrePysaApi.PysaType.t option ->
+    explicit_annotation:ExplicitAnnotation.t ->
     t
 
-  val from_pysa_type : PyrePysaApi.PysaType.t -> t
+  val from_inferred_type : PyrePysaApi.PysaType.t option -> t
 
   val is_annotated : t -> bool
 
-  val as_type : t -> PyrePysaApi.PysaType.t
+  val inferred_type : t -> PyrePysaApi.PysaType.t option
 
-  val as_original_annotation : t -> Ast.Expression.t option
+  val explicit_annotation : t -> ExplicitAnnotation.t
 
   (* Show the original annotation, as written by the user. *)
-  val show_original_annotation : t -> string
+  val show_explicit_annotation : t -> string option
 
   (* Show the fully qualified type annotation from the type checker. *)
   val show_fully_qualified_annotation : t -> string
@@ -483,7 +491,7 @@ module Modelable : sig
 
   val target_name : t -> Ast.Reference.t
 
-  val type_annotation : t -> TypeAnnotation.t option
+  val type_annotation : t -> TypeAnnotation.t
 
   val undecorated_signatures : t -> PyrePysaApi.ModelQueries.FunctionSignature.t list