Source file main.ml

open Expect_test_common.Std
open Base
open Ppxlib
open Ast_builder.Default

let lifter ~loc = object
  inherit [expression] Lifter.lift
  inherit Ppxlib_metaquot_lifters.expression_lifters loc

  method filename file_name =
    eapply ~loc (evar ~loc "Expect_test_common.Std.File.Name.of_string")
      [ estring ~loc (File.Name.to_string file_name) ]
end

let lift_location ~loc of_loc =
  (lifter ~loc)#location of_loc
;;

let lift_expectation ~loc expect =
  let exp = (lifter ~loc)#raw expect in
  (* the constraint is added in order to avoid warnings if the code
     is compiled with -principal. *)
  [%expr ([%e exp] : string Expect_test_common.Std.Expectation.t)]
;;

let eoption ~loc x =
  match x with
  | None -> pexp_construct ~loc (Located.mk ~loc (lident "None")) None
  | Some e ->
    pexp_construct ~loc (Located.mk ~loc (lident "Some")) (Some e)
;;

let estring_option ~loc x =
  eoption ~loc (Option.map x ~f:(estring ~loc))
;;

(* Grab a list of all the output expressions *)
let collect_expectations = object
  inherit [(Location.t * Expectation.Raw.t) list] Ast_traverse.fold as super

  method! expression expr acc =
    match Expect_extension.match_expectation expr with
    | None ->
      super#expression expr acc
    | Some ext ->
      assert_no_attributes expr.pexp_attributes;
      (expr.pexp_loc, ext) :: acc
end

let replace_expects = object
  inherit Ast_traverse.map as super

  method! expression ({ pexp_attributes; pexp_loc; _ } as expr) =
    match Expect_extension.match_expectation expr with
    | None -> super#expression expr
    | Some ext ->
      let f_var =
        match ext.body with
        | Exact _ | Pretty _ | Unreachable -> "Expect_test_collector.save_output"
        | Output -> "Expect_test_collector.save_and_return_output"
      in
      let loc = { pexp_loc with loc_end = pexp_loc.loc_start } in
      let expr =
        [%expr
          [%e evar ~loc f_var]
            [%e lift_location ~loc ext.extid_location]
        ]
      in
      { expr with pexp_attributes }
end

let file_digest =
  let cache = Hashtbl.create (module String) ~size:32 in
  fun fname ->
    Hashtbl.find_or_add cache fname ~default:(fun () ->
      Caml.Digest.file fname |> Caml.Digest.to_hex)

let rewrite_test_body ~descr ~tags ~uncaught_exn pstr_loc body =
  let loc = pstr_loc in
  let expectations =
    List.map (collect_expectations#expression body [])
      ~f:(fun (loc, expect_extension) -> lift_expectation ~loc expect_extension)
    |> elist ~loc
  in

  let uncaught_exn =
    Option.map uncaught_exn ~f:(fun (loc, expectation) ->
      lift_expectation ~loc expectation)
    |> eoption ~loc
  in

  let body = replace_expects#expression body in

  let absolute_filename =
    Ppx_here_expander.expand_filename pstr_loc.loc_start.pos_fname
  in

  let hash = file_digest loc.loc_start.pos_fname in
  [%expr
    let module Expect_test_collector = Expect_test_collector.Make(Expect_test_config) in
    Expect_test_collector.run
      ~file_digest:        (Expect_test_common.Std.File.Digest.of_string [%e estring ~loc hash])
      ~location:           [%e lift_location ~loc (Ppx_expect_payload.transl_loc pstr_loc)]
      ~absolute_filename:  [%e estring ~loc absolute_filename]
      ~description:        [%e estring_option ~loc descr]
      ~tags:               [%e elist ~loc (List.map tags ~f:(estring ~loc))]
      ~expectations:       [%e expectations]
      ~uncaught_exn_expectation: [%e uncaught_exn]
      ~inline_test_config: (module Inline_test_config)
      (fun () -> [%e body])
  ]

module P = struct
  open Ast_pattern

  let uncaught_exn =
    Attribute.declare_with_name_loc
      "@expect.uncaught_exn"
      Attribute.Context.value_binding
      (map1' (Ppx_expect_payload.pattern ()) ~f:(fun loc x -> (loc, x)))
      (fun ~name_loc (loc, x) ->
         (loc,
          Ppx_expect_payload.make x ~kind:Normal ~extension_id_loc:name_loc))

  let opt_name () =
    map (pstring __) ~f:(fun f x -> f (Some x)) |||
    map ppat_any     ~f:(fun f   -> f None)

  let pattern () =
    pstr ((
      pstr_value nonrecursive (
        (Attribute.pattern uncaught_exn
           (value_binding
              ~pat:(
                map
                  (Attribute.pattern Ppx_inline_test.tags (opt_name ()))
                  ~f:(fun f attributes name_opt ->
                    f ~name:name_opt
                      ~tags:(match attributes with
                        | None -> []
                        | Some x -> x)))
              ~expr:__)
         ^:: nil)
      ) ^:: nil))
end

(* Set to [true] when we see a [%expect_test] extension *)
module Has_tests =
  Driver.Create_file_property
    (struct let name = "ppx_expect.has_tests" end)
    (Bool)

let expect_test =
  Extension.declare_inline "expect_test" Structure_item
    (P.pattern ())
    (fun ~loc ~path:_ uncaught_exn ~name ~tags code ->
       Has_tests.set true;
       Ppx_inline_test.validate_extension_point_exn
         ~name_of_ppx_rewriter:"ppx_expect" ~loc ~tags;
       rewrite_test_body ~descr:name ~tags ~uncaught_exn loc code
       |> Ppx_inline_test.maybe_drop loc)
;;

let () =
  Driver.register_transformation "expect_test"
    ~rules:[ Context_free.Rule.extension expect_test ]
    ~enclose_impl:(fun whole_loc ->
      match whole_loc, Ppx_inline_test_libname.get () with
      | None, _ | _, None -> ([], [])
      | Some loc, Some _ ->
        let maybe_drop = Ppx_inline_test.maybe_drop in
        let absolute_filename =
          Ppx_here_expander.expand_filename loc.loc_start.pos_fname
        in
        let header =
          let loc = { loc with loc_end = loc.loc_start } in
          maybe_drop loc [%expr Expect_test_collector.Current_file.set
                                  ~absolute_filename:[%e estring ~loc absolute_filename]]
        and footer =
          let loc = { loc with loc_start = loc.loc_end } in
          maybe_drop loc [%expr Expect_test_collector.Current_file.unset ()]
        in
        (header, footer)
    )
;;