Collect standard formatters in expect tests

So that the output of #-directives is captured

testsuite/tests/typing-modules/printing.ml

@@ -22,4 +22,5 @@ module M = struct module N = struct let x = 1 end end;;
 #show_module M;;
 [%%expect{|
 module M : sig module N : sig val x : int end end
+module M : sig module N : sig ... end end
 |}];;

testsuite/tools/expect_test.ml

@@ -158,6 +158,33 @@ module Compiler_messages = struct
       f
 end
 
+let collect_formatters buf pps ~f =
+  List.iter (fun pp -> Format.pp_print_flush pp ()) pps;
+  let save = List.map (fun pp -> Format.pp_get_formatter_out_functions pp ()) pps in
+  let restore () =
+    List.iter2
+      (fun pp out_functions ->
+         Format.pp_print_flush pp ();
+         Format.pp_set_formatter_out_functions pp out_functions)
+      pps save
+  in
+  let out_string str ofs len = Buffer.add_substring buf str ofs len
+  and out_flush = ignore
+  and out_newline () = Buffer.add_char buf '\n'
+  and out_spaces n = for i = 1 to n do Buffer.add_char buf ' ' done in
+  let out_functions = { Format.out_string; out_flush; out_newline; out_spaces } in
+  List.iter
+    (fun pp -> Format.pp_set_formatter_out_functions pp out_functions)
+    pps;
+  match f () with
+  | x             -> restore (); x
+  | exception exn -> restore (); raise exn
+
+(* Invariant: ppf = Format.formatter_of_buffer buf *)
+let capture_everything buf ppf ~f =
+  collect_formatters buf [Format.std_formatter; Format.err_formatter] ~f:(fun () ->
+    Compiler_messages.capture ppf ~f)
+
 let exec_phrase ppf phrase =
   if !Clflags.dump_parsetree then Printast. top_phrase ppf phrase;
   if !Clflags.dump_source    then Pprintast.top_phrase ppf phrase;
@@ -242,7 +269,7 @@ let eval_expect_file _fname ~file_contents =
     Misc.delete_eol_spaces s
   in
   let corrected_expectations =
-    Compiler_messages.capture ppf ~f:(fun () ->
+    capture_everything buf ppf ~f:(fun () ->
       List.fold_left chunks ~init:[] ~f:(fun acc chunk ->
         let output = exec_phrases chunk.phrases in
         match eval_expectation chunk.expectation ~output with
@@ -254,7 +281,7 @@ let eval_expect_file _fname ~file_contents =
     match trailing_code with
     | None -> ""
     | Some phrases ->
-      Compiler_messages.capture ppf ~f:(fun () -> exec_phrases phrases)
+      capture_everything buf ppf ~f:(fun () -> exec_phrases phrases)
   in
   { corrected_expectations; trailing_output }