[flow][apply-code-actions] Suggest Imports

Summary:
Exposes import suggestions for each unbound name error in the file. The JSON output maps each name to an array of suggestions for that name.

--pretty prettifies the output.

The tests need a `sed` to replace filenames because we generate new directories to run the tests.

Changelog: [internal]

Reviewed By: gkz

Differential Revision: D69545414

fbshipit-source-id: 1172968109c207a6f38d7fca79dd8503db33058f

src/commands/applyCodeActionCommand.ml

@@ -11,21 +11,6 @@ open CommandSpec
 (* This module implements the flow command `apply-code-action` which exposes LSP code-actions via the CLI *)
 let handle_error ?(code = Exit.Unknown_error) msg = Exit.(exit ~msg code)
 
-let handle_ok in_place patch source_path input =
-  let write_patch content =
-    let output_channel =
-      if in_place then
-        open_out source_path
-      else
-        stdout
-    in
-    output_string output_channel @@ Replacement_printer.print patch content;
-    close_out output_channel
-  in
-  match File_input.content_of_file_input input with
-  | Ok content -> write_patch content
-  | Error msg -> handle_error msg
-
 module SourceAddMissingImports = struct
   let spec =
     {
@@ -48,6 +33,21 @@ module SourceAddMissingImports = struct
         );
     }
 
+  let handle_ok in_place patch source_path input =
+    let write_patch content =
+      let output_channel =
+        if in_place then
+          open_out source_path
+        else
+          stdout
+      in
+      output_string output_channel @@ Replacement_printer.print patch content;
+      close_out output_channel
+    in
+    match File_input.content_of_file_input input with
+    | Ok content -> write_patch content
+    | Error msg -> handle_error msg
+
   let main base_flags connect_params _json _pretty root_arg path wait_for_recheck in_place file () =
     let source_path = expand_path file in
     let input = get_file_from_filename_or_stdin ~cmd:spec.name path (Some source_path) in
@@ -85,7 +85,16 @@ module SuggestImports = struct
         );
     }
 
-  let main base_flags connect_params _json _pretty root_arg path wait_for_recheck in_place file () =
+  let handle_ok pretty lsp_result =
+    let result_json =
+      Hh_json.JSON_Object
+        (SMap.bindings lsp_result
+        |> List.map (fun (name, result) -> (name, Lsp_fmt.print_codeActionResult ~key:"" result))
+        )
+    in
+    print_endline (Hh_json.json_to_string ~sort_keys:pretty ~pretty result_json)
+
+  let main base_flags connect_params _json pretty root_arg path wait_for_recheck _in_place file () =
     let source_path = expand_path file in
     let input = get_file_from_filename_or_stdin ~cmd:spec.name path (Some source_path) in
     let root = get_the_root ~base_flags ~input root_arg in
@@ -96,7 +105,7 @@ module SuggestImports = struct
     in
     let result = connect_and_make_request flowconfig_name connect_params root request in
     match result with
-    | ServerProt.Response.APPLY_CODE_ACTION (Ok patch) -> handle_ok in_place patch source_path input
+    | ServerProt.Response.SUGGEST_IMPORTS (Ok result) -> handle_ok pretty result
     | _ -> handle_error "Flow: invalid server response"
 
   let command = CommandSpec.command spec main

src/hack_forked/utils/lsp/lsp_fmt.mli

@@ -106,6 +106,8 @@ val parse_documentOnTypeFormatting : Hh_json.json option -> Lsp.DocumentOnTypeFo
 
 val print_documentOnTypeFormatting : Lsp.DocumentOnTypeFormatting.result -> Hh_json.json
 
+val print_codeActionResult : key:string -> Lsp.CodeAction.result -> Hh_json.json
+
 val parse_initialize : Hh_json.json option -> Lsp.Initialize.params
 
 val error_of_exn : exn -> Lsp.Error.t

src/server/command_handler/commandHandler.ml

@@ -544,6 +544,18 @@ let autofix_imports_cli ~options ~env ~profiling ~loc_of_aloc ~module_system_inf
     ~file_key
     ~file_content
 
+let suggest_imports_cli ~options ~env ~profiling ~loc_of_aloc ~module_system_info ~input =
+  let file_key = file_key_of_file_input ~options ~env input in
+  File_input.content_of_file_input input >>= fun file_content ->
+  Code_action_service.suggest_imports_cli
+    ~options
+    ~profiling
+    ~env
+    ~loc_of_aloc
+    ~module_system_info
+    ~file_key
+    ~file_content
+
 let autocomplete
     ~trigger_character
     ~reader
@@ -1519,7 +1531,18 @@ let handle_apply_code_action ~options ~reader ~profiling ~env action file_input
       in
       Lwt.return (ServerProt.Response.APPLY_CODE_ACTION result, None)
     | SuggestImports ->
-      Lwt.return (ServerProt.Response.APPLY_CODE_ACTION (Error "Not yet implemented"), None)
+      let result =
+        try_with (fun () ->
+            suggest_imports_cli
+              ~options
+              ~profiling
+              ~env
+              ~loc_of_aloc:(Parsing_heaps.Reader.loc_of_aloc ~reader)
+              ~module_system_info:(mk_module_system_info ~options ~reader)
+              ~input:file_input
+        )
+      in
+      Lwt.return (ServerProt.Response.SUGGEST_IMPORTS result, None)
   )
 
 let handle_autocomplete

src/server/protocol/serverProt.ml

@@ -329,6 +329,8 @@ module Response = struct
 
   type rage_response = (string * string) list
 
+  type suggest_imports_response = (Lsp.CodeAction.command_or_action list SMap.t, string) result
+
   type graph_response = (graph_response_subgraph, string) result
 
   and graph_response_subgraph = (string * string list) list
@@ -369,6 +371,7 @@ module Response = struct
         lazy_stats: lazy_stats;
       }
     | SAVE_STATE of (string, string) result
+    | SUGGEST_IMPORTS of suggest_imports_response
 
   let to_string = function
     | APPLY_CODE_ACTION _ -> "apply-code-action response"
@@ -390,4 +393,5 @@ module Response = struct
     | RAGE _ -> "rage response"
     | STATUS _ -> "status response"
     | SAVE_STATE _ -> "save_state response"
+    | SUGGEST_IMPORTS _ -> "suggest imports response"
 end

src/services/code_action/code_action_service.ml

@@ -1611,6 +1611,41 @@ let with_type_checked_file ~options ~profiling ~env ~file_key ~file_content ~f =
   | Ok (Parse_artifacts { ast; _ }, Typecheck_artifacts { cx; _ }) -> f ~cx ~ast
   | _ -> Error "Failed to parse or check file"
 
+let suggest_imports_cli
+    ~options ~profiling ~env ~loc_of_aloc ~module_system_info ~file_key ~file_content =
+  let uri = File_key.to_string file_key |> Lsp_helpers.path_to_lsp_uri ~default_path:"" in
+  let get_edits ~cx ~ast =
+    let errors = Context.errors cx in
+    let { ServerEnv.exports; _ } = env in
+    let (imports, _) =
+      Flow_error.ErrorSet.fold
+        (fun error (imports, unbound_names) ->
+          match
+            Flow_error.msg_of_error error |> Error_message.map_loc_of_error_message loc_of_aloc
+          with
+          | Error_message.EBuiltinNameLookupFailed { loc = error_loc; name }
+            when Options.autoimports options ->
+            let ranked_imports =
+              suggest_imports
+                ~layout_options:(Code_action_utils.layout_options options)
+                ~module_system_info
+                ~ast
+                ~diagnostics:[]
+                ~imports_ranked_usage:true
+                ~exports
+                ~name
+                uri
+                error_loc
+            in
+            (SMap.add name ranked_imports imports, SSet.add name unbound_names)
+          | _ -> (imports, unbound_names))
+        errors
+        (SMap.empty, SSet.empty)
+    in
+    Ok imports
+  in
+  with_type_checked_file ~options ~profiling ~env ~file_key ~file_content ~f:get_edits
+
 let autofix_imports_cli
     ~options ~profiling ~env ~loc_of_aloc ~module_system_info ~file_key ~file_content =
   let src_dir = File_key.to_string file_key |> Filename.dirname |> Base.Option.return in

src/services/code_action/code_action_service.mli

@@ -54,25 +54,35 @@ val code_actions_at_loc :
   loc:Loc.t ->
   (Lsp.CodeAction.command_or_action list, string) result
 
-val autofix_imports_lsp :
+val autofix_imports_cli :
   options:Options.t ->
+  profiling:Profiling_js.running ->
   env:ServerEnv.env ->
   loc_of_aloc:(ALoc.t -> Loc.t) ->
   module_system_info:Lsp_module_system_info.t ->
-  cx:Context.t ->
-  ast:(Loc.t, Loc.t) Flow_ast.Program.t ->
-  uri:Lsp.DocumentUri.t ->
-  Lsp.TextEdit.t list
+  file_key:File_key.t ->
+  file_content:string ->
+  (Replacement_printer.patch, string) result
 
-val autofix_imports_cli :
+val suggest_imports_cli :
   options:Options.t ->
   profiling:Profiling_js.running ->
   env:ServerEnv.env ->
   loc_of_aloc:(ALoc.t -> Loc.t) ->
   module_system_info:Lsp_module_system_info.t ->
   file_key:File_key.t ->
   file_content:string ->
-  (Replacement_printer.patch, string) result
+  (Lsp.CodeAction.command_or_action list SMap.t, string) result
+
+val autofix_imports_lsp :
+  options:Options.t ->
+  env:ServerEnv.env ->
+  loc_of_aloc:(ALoc.t -> Loc.t) ->
+  module_system_info:Lsp_module_system_info.t ->
+  cx:Context.t ->
+  ast:(Loc.t, Loc.t) Flow_ast.Program.t ->
+  uri:Lsp.DocumentUri.t ->
+  Lsp.TextEdit.t list
 
 val autofix_exports :
   options:Options.t ->

tests/apply_code_action_command/ExportFoo1.js

@@ -0,0 +1,3 @@
+// @flow
+
+export const foo = 3;

tests/apply_code_action_command/ExportFoo2.js

@@ -0,0 +1,3 @@
+// @flow
+
+export const foo = 3;

tests/apply_code_action_command/apply_code_action_command.exp

@@ -5,6 +5,7 @@ import OtherModule from "./OtherModule";
 
 OtherModule;
 > Confirm no errors
+Found 0 errors
 > apply-code-action 'source.addMissingImports' tmp/multi.js
 // @flow
 
@@ -13,4 +14,71 @@ import OtherModule from "./OtherModule";
 OtherModule;
 OtherModule;
 > Confirm no errors
-No errors!
+Found 0 errors
+> apply-code-action suggestImports tmp/suggest_imports.js
+{
+  "OtherModule":[
+    {
+      "command":{
+        "arguments":["textDocument/codeAction","import","Import default from ./OtherModule"],
+        "command":"log:",
+        "title":""
+      },
+      "diagnostics":[],
+      "edit":{
+        "changes":{
+          "file:tmp/suggest_imports.js": [
+            {
+              "newText":"import OtherModule from \"./OtherModule\";\n\n",
+              "range":{"end":{"character":0,"line":1},"start":{"character":0,"line":1}}
+            }
+          ]
+        }
+      },
+      "kind":"quickfix",
+      "title":"Import default from ./OtherModule"
+    }
+  ],
+  "foo":[
+    {
+      "command":{
+        "arguments":["textDocument/codeAction","import","Import from ./ExportFoo1"],
+        "command":"log:",
+        "title":""
+      },
+      "diagnostics":[],
+      "edit":{
+        "changes":{
+          "file:tmp/suggest_imports.js": [
+            {
+              "newText":"import { foo } from \"./ExportFoo1\";\n\n",
+              "range":{"end":{"character":0,"line":1},"start":{"character":0,"line":1}}
+            }
+          ]
+        }
+      },
+      "kind":"quickfix",
+      "title":"Import from ./ExportFoo1"
+    },
+    {
+      "command":{
+        "arguments":["textDocument/codeAction","import","Import from ./ExportFoo2"],
+        "command":"log:",
+        "title":""
+      },
+      "diagnostics":[],
+      "edit":{
+        "changes":{
+          "file:tmp/suggest_imports.js": [
+            {
+              "newText":"import { foo } from \"./ExportFoo2\";\n\n",
+              "range":{"end":{"character":0,"line":1},"start":{"character":0,"line":1}}
+            }
+          ]
+        }
+      },
+      "kind":"quickfix",
+      "title":"Import from ./ExportFoo2"
+    }
+  ]
+}

tests/apply_code_action_command/suggest_imports.js

@@ -0,0 +1,3 @@
+// @flow
+OtherModule;
+foo;

tests/apply_code_action_command/test.sh

@@ -14,11 +14,11 @@ echo '> apply-code-action '\''source.addMissingImports'\'' tmp/a.js'
 assert_ok "$FLOW" apply-code-action 'source.addMissingImports' tmp/a.js
 assert_ok "$FLOW" apply-code-action 'source.addMissingImports' --in-place tmp/a.js
 echo '> Confirm no errors'
-assert_ok "$FLOW" force-recheck tmp/a.js
+assert_ok "$FLOW" focus-check tmp/a.js
 echo '> apply-code-action '\''source.addMissingImports'\'' tmp/multi.js'
 assert_ok "$FLOW" apply-code-action 'source.addMissingImports' tmp/multi.js
 assert_ok "$FLOW" apply-code-action 'source.addMissingImports' --in-place tmp/multi.js
 echo '> Confirm no errors'
-assert_ok "$FLOW" force-recheck tmp/multi.js
-
-assert_ok "$FLOW" status tmp
+assert_ok "$FLOW" focus-check tmp/multi.js
+echo '> apply-code-action suggestImports tmp/suggest_imports.js'
+"$FLOW" apply-code-action suggestImports --pretty tmp/suggest_imports.js | sed -e 's/file:.*:/file:tmp\/suggest_imports.js": /'