Allow all kind of expressions to appear as decorators

Reviewed By: kbansal

Differential Revision: D31865845

fbshipit-source-id: 57e8ccb13982ef2e8a3103367c830071a769bd20

Author: grievejia

source/analysis/attributeResolution.ml

@@ -1037,7 +1037,7 @@ module ClassSummary = struct
     name: Reference.t;
     qualifier: Reference.t;
     bases: bases;
-    decorators: Decorator.t list;
+    decorators: Expression.t list;
     class_attributes: ClassAttributes.t;
   }
   [@@deriving compare, eq, sexp, show, hash]
@@ -1101,7 +1101,6 @@ module ClassSummary = struct
 
 
   let has_decorator { decorators; _ } decorator =
-    let decorators = List.map decorators ~f:Decorator.to_expression in
     Expression.exists_in_list ~expression_list:decorators decorator
 
 

source/analysis/classSummary.ml

@@ -111,7 +111,7 @@ type t = {
   name: Reference.t;
   qualifier: Reference.t;
   bases: bases;
-  decorators: Decorator.t list;
+  decorators: Expression.t list;
   class_attributes: ClassAttributes.t;
 }
 [@@deriving compare, eq, sexp, show, hash]

source/analysis/classSummary.mli

@@ -212,8 +212,7 @@ module Visit = struct
                ~location:(Define.name_location ~body_location:statement.location define)
                name);
           List.iter parameters ~f:visit_parameter;
-          List.map decorators ~f:Ast.Statement.Decorator.to_expression
-          |> List.iter ~f:postcondition_visit;
+          List.iter decorators ~f:postcondition_visit;
           Option.iter ~f:postcondition_visit return_annotation
       | Import { Import.from; imports } ->
           let visit_import { Node.value = { Import.name; _ }; location = import_location } =

source/analysis/lookup.ml

@@ -744,7 +744,7 @@ let qualify
           List.map
             decorators
             ~f:
-              (qualify_decorator
+              (qualify_expression
                  ~qualify_strings:DoNotQualify
                  ~scope:{ scope with use_forward_references = true })
         in
@@ -778,7 +778,7 @@ let qualify
           }
         in
         let decorators =
-          List.map decorators ~f:(qualify_decorator ~qualify_strings:DoNotQualify ~scope)
+          List.map decorators ~f:(qualify_expression ~qualify_strings:DoNotQualify ~scope)
         in
         let body =
           let qualifier = qualify_if_needed ~qualifier name in
@@ -797,21 +797,25 @@ let qualify
                   let return_annotation =
                     return_annotation >>| qualify_expression ~scope ~qualify_strings:Qualify
                   in
-                  let qualify_decorator
-                      ({ Decorator.name = { Node.value = name; _ }; _ } as decorator)
-                    =
-                    match name |> Reference.as_list |> List.rev with
-                    | ["staticmethod"]
-                    | ["classmethod"]
-                    | ["property"]
-                    | "getter" :: _
-                    | "setter" :: _
-                    | "deleter" :: _ ->
+                  let qualify_decorator decorator =
+                    let is_reserved name =
+                      match Reference.as_list name |> List.rev with
+                      | ["staticmethod"]
+                      | ["classmethod"]
+                      | ["property"]
+                      | "getter" :: _
+                      | "setter" :: _
+                      | "deleter" :: _ ->
+                          true
+                      | _ -> false
+                    in
+                    match Decorator.from_expression decorator with
+                    | Some { Decorator.name = { Node.value = name; _ }; _ } when is_reserved name ->
                         decorator
                     | _ ->
                         (* TODO (T41755857): Decorator qualification logic should be slightly more
                            involved than this. *)
-                        qualify_decorator ~qualify_strings:DoNotQualify ~scope decorator
+                        qualify_expression ~qualify_strings:DoNotQualify ~scope decorator
                   in
                   let decorators = List.map decorators ~f:qualify_decorator in
                   let signature =
@@ -1252,11 +1256,6 @@ let qualify
       | Constant _ -> value
     in
     { expression with Node.value }
-  and qualify_decorator ~qualify_strings ~scope { Decorator.name; arguments } =
-    {
-      Decorator.name = Node.map name ~f:(qualify_reference ~scope);
-      arguments = arguments >>| List.map ~f:(qualify_argument ~qualify_strings ~scope);
-    }
   and qualify_argument { Call.Argument.name; value } ~qualify_strings ~scope =
     let name =
       let rename identifier =
@@ -3093,15 +3092,11 @@ module AccessCollector = struct
           List.fold base_arguments ~init:collected ~f:(fun sofar { Call.Argument.value; _ } ->
               from_expression sofar value)
         in
-        List.map decorators ~f:Decorator.to_expression
-        |> List.fold ~init:collected ~f:from_expression
+        List.fold ~init:collected ~f:from_expression decorators
     | Define
         { Define.signature = { Define.Signature.decorators; parameters; return_annotation; _ }; _ }
       ->
-        let collected =
-          List.map decorators ~f:Decorator.to_expression
-          |> List.fold ~init:collected ~f:from_expression
-        in
+        let collected = List.fold ~init:collected ~f:from_expression decorators in
         let collected =
           List.fold
             parameters
@@ -3828,11 +3823,15 @@ let inline_six_metaclass ({ Source.statements; _ } as source) =
       let transform_class
           ~class_statement:({ Class.base_arguments; decorators; _ } as class_statement)
         =
-        let is_six_add_metaclass_decorator { Decorator.name; _ } =
-          Identifier.equal (Node.value name |> Reference.show) "six.add_metaclass"
+        let is_six_add_metaclass_decorator expression =
+          match Decorator.from_expression expression with
+          | Some ({ Decorator.name = { Node.value = name; _ }; _ } as decorator)
+            when Reference.equal name (Reference.create_from_list ["six"; "add_metaclass"]) ->
+              Either.First decorator
+          | _ -> Either.Second expression
         in
         let six_add_metaclass_decorators, rest =
-          List.partition_tf decorators ~f:is_six_add_metaclass_decorator
+          List.partition_map decorators ~f:is_six_add_metaclass_decorator
         in
         match six_add_metaclass_decorators with
         | [

source/analysis/preprocessing.ml

@@ -152,9 +152,7 @@ module Binding = struct
         let sofar = of_expression sofar value in
         of_unannotated_target ~kind:(Kind.AssignTarget None) sofar target
     | Statement.Class { Class.name; base_arguments; decorators; _ } ->
-        let sofar =
-          List.map decorators ~f:Decorator.to_expression |> List.fold ~init:sofar ~f:of_expression
-        in
+        let sofar = List.fold ~init:sofar ~f:of_expression decorators in
         let sofar =
           List.fold
             base_arguments
@@ -163,9 +161,7 @@ module Binding = struct
         in
         { kind = Kind.ClassName; name = Reference.show name; location } :: sofar
     | Statement.Define { Define.signature = { name; decorators; parameters; _ } as signature; _ } ->
-        let sofar =
-          List.map decorators ~f:Decorator.to_expression |> List.fold ~init:sofar ~f:of_expression
-        in
+        let sofar = List.fold ~init:sofar ~f:of_expression decorators in
         let sofar =
           List.fold
             parameters

source/analysis/scope.ml

@@ -15,7 +15,6 @@ open Test
 
 let test_apply_decorators context =
   let resolution = ScratchProject.setup ~context [] |> ScratchProject.build_global_resolution in
-  let decorator ?arguments name = { Decorator.name = + !&name; arguments } in
   let create_define ~decorators ~parameters ~return_annotation =
     {
       Define.Signature.name = !&"define";
@@ -62,15 +61,15 @@ let test_apply_decorators context =
     Type.string;
   assert_apply_contextlib_decorators
     (create_define
-       ~decorators:[decorator "contextlib.contextmanager"]
+       ~decorators:[!"contextlib.contextmanager"]
        ~parameters:[]
        ~return_annotation:
          (Some
             (+Expression.Constant (Constant.String (StringLiteral.create "typing.Iterator[str]")))))
     (Type.parametric "contextlib._GeneratorContextManager" [Single Type.string]);
   assert_apply_contextlib_decorators
     (create_define
-       ~decorators:[decorator "contextlib.contextmanager"]
+       ~decorators:[!"contextlib.contextmanager"]
        ~parameters:[]
        ~return_annotation:
          (Some
@@ -81,7 +80,7 @@ let test_apply_decorators context =
   let create_parameter ~name = Parameter.create ~location:Location.any ~name () in
   (* Custom decorators. *)
   create_define
-    ~decorators:[decorator "$strip_first_parameter"]
+    ~decorators:[!"$strip_first_parameter"]
     ~parameters:[create_parameter ~name:"self"; create_parameter ~name:"other"]
     ~return_annotation:None
   |> (fun define ->

source/analysis/test/resolvedCallableTest.ml

@@ -1400,7 +1400,22 @@ let test_updates context =
                     };
                 ];
               decorators =
-                [{ name = node ~start:(3, 1) ~stop:(3, 9) !&"typing.overload"; arguments = None }];
+                [
+                  node
+                    ~start:(3, 1)
+                    ~stop:(3, 9)
+                    (Expression.Name
+                       (Name.Attribute
+                          {
+                            Name.Attribute.base =
+                              node
+                                ~start:(3, 1)
+                                ~stop:(3, 9)
+                                (Expression.Name (Name.Identifier "typing"));
+                            attribute = "overload";
+                            special = false;
+                          }));
+                ];
               return_annotation =
                 Some (node ~start:(4, 19) ~stop:(4, 22) (Expression.Name (Name.Identifier "int")));
               async = false;
@@ -1445,7 +1460,22 @@ let test_updates context =
                     };
                 ];
               decorators =
-                [{ name = node ~start:(7, 1) ~stop:(7, 9) !&"typing.overload"; arguments = None }];
+                [
+                  node
+                    ~start:(7, 1)
+                    ~stop:(7, 9)
+                    (Expression.Name
+                       (Name.Attribute
+                          {
+                            Name.Attribute.base =
+                              node
+                                ~start:(7, 1)
+                                ~stop:(7, 9)
+                                (Expression.Name (Name.Identifier "typing"));
+                            attribute = "overload";
+                            special = false;
+                          }));
+                ];
               return_annotation =
                 Some (node ~start:(8, 19) ~stop:(8, 22) (Expression.Name (Name.Identifier "str")));
               async = false;
@@ -1552,7 +1582,12 @@ let test_updates context =
                            { Parameter.name = "$parameter$self"; value = None; annotation = None };
                        ];
                      decorators =
-                       [{ name = node ~start:(3, 3) ~stop:(3, 11) !&"property"; arguments = None }];
+                       [
+                         node
+                           ~start:(3, 3)
+                           ~stop:(3, 11)
+                           (Expression.Name (Name.Identifier "property"));
+                       ];
                      return_annotation =
                        Some
                          (node
@@ -1605,10 +1640,20 @@ let test_updates context =
                             ];
                           decorators =
                             [
-                              {
-                                name = node ~start:(5, 3) ~stop:(5, 13) !&"foo.setter";
-                                arguments = None;
-                              };
+                              node
+                                ~start:(5, 3)
+                                ~stop:(5, 13)
+                                (Expression.Name
+                                   (Name.Attribute
+                                      {
+                                        Name.Attribute.base =
+                                          node
+                                            ~start:(5, 3)
+                                            ~stop:(5, 6)
+                                            (Expression.Name (Name.Identifier "foo"));
+                                        attribute = "setter";
+                                        special = false;
+                                      }));
                             ];
                           return_annotation =
                             Some