Recognise the metadata field on Odoc code blocks

lib/mli_parser.ml

@@ -1,7 +1,14 @@
 open! Compat
 
 module Code_block = struct
-  type t = { location : Odoc_parser.Loc.span; contents : string }
+  type t = {
+    location : Odoc_parser.Loc.span;
+    metadata :
+      (string Odoc_parser.Loc.with_location
+      * string Odoc_parser.Loc.with_location option)
+      option;
+    contents : string;
+  }
 end
 
 let drop_last lst =
@@ -51,8 +58,8 @@ let extract_code_blocks ~(location : Lexing.position) ~docstring =
     List.map
       (fun block ->
         match Odoc_parser.Loc.value block with
-        | `Code_block (_metadata, { Odoc_parser.Loc.value = contents; _ }) ->
-            [ { Code_block.location = block.location; contents } ]
+        | `Code_block (metadata, { Odoc_parser.Loc.value = contents; _ }) ->
+            [ { Code_block.location = block.location; metadata; contents } ]
         | `List (_, _, lists) -> List.map acc lists |> List.concat
         | _ -> [])
       blocks
@@ -118,6 +125,50 @@ let docstring_code_blocks str =
     (docstrings (Lexing.from_string str))
   |> List.concat
 
+let make_block ~loc code_block =
+  let handle_header = function
+    | Some (lang_tag, labels) -> (
+        let header = Block.Header.of_string (Odoc_parser.Loc.value lang_tag) in
+        match labels with
+        | None -> Ok (header, [])
+        | Some labels -> (
+            match Label.of_string (Odoc_parser.Loc.value labels) with
+            | Ok labels -> Ok (header, labels)
+            | Error msgs ->
+                Error (List.hd msgs) (* TODO: Report precise location *)))
+    | None ->
+        (* If not specified, blocks are run as ocaml blocks *)
+        Ok (Some OCaml, [])
+  in
+  match handle_header code_block.Code_block.metadata with
+  | Error _ as e -> e
+  | Ok (header, labels) ->
+      let contents = Compat.String.split_on_char '\n' code_block.contents in
+      Block.mk ~loc ~section:None ~labels ~header ~contents ~legacy_labels:false
+        ~errors:[]
+
+let code_block_markup code_block =
+  let open Document in
+  let opening =
+    match code_block.Code_block.metadata with
+    | Some (lang_tag, labels) ->
+        let labels =
+          match labels with
+          | Some s -> [ Text " "; Text (Odoc_parser.Loc.value s) ]
+          | None -> []
+        in
+        [ Text "{@"; Text (Odoc_parser.Loc.value lang_tag) ]
+        @ labels @ [ Text "[" ]
+    | None -> [ Text "{[" ]
+  in
+  let hpad =
+    let has_several_lines = String.contains code_block.contents '\n' in
+    let column = code_block.location.start.column in
+    if not has_several_lines then ""
+    else Astring.String.v ~len:column (fun _ -> ' ')
+  in
+  (opening, [ Text (hpad ^ "]}") ])
+
 let parse_mli file_contents =
   (* Find the locations of the code blocks within [file_contents], then slice it up into
      [Text] and [Block] parts by using the starts and ends of those blocks as
@@ -132,22 +183,14 @@ let parse_mli file_contents =
           Document.Text
             (slice lines ~start:!cursor ~end_:code_block.location.start)
         in
-        let column = code_block.location.start.column in
-        let contents = Compat.String.split_on_char '\n' code_block.contents in
         let block =
-          match
-            Block.mk ~loc ~section:None ~labels:[] ~header:(Some OCaml)
-              ~contents ~legacy_labels:false ~errors:[]
-          with
+          match make_block ~loc code_block with
           | Ok block -> Document.Block block
           | Error _ -> failwith "Error creating block"
         in
-        let hpad =
-          if List.length contents = 1 then ""
-          else Astring.String.v ~len:column (fun _ -> ' ')
-        in
+        let opening, closing = code_block_markup code_block in
         cursor := code_block.location.end_;
-        [ pre_text; Text "{["; block; Text (hpad ^ "]}") ])
+        [ pre_text ] @ opening @ [ block ] @ closing)
       code_blocks
     |> List.concat
   in

test/bin/mdx-test/expect/simple-mli/test-case.mli

@@ -19,14 +19,30 @@
       print_endline (String.concat " " [the; last; phrase])
     ]}
 
-    {[
+    With the optional header:
+
+    {@ocaml[
       # List.map (fun x -> x * x) [(1 + 9); 2; 3]
       - : int list = [100; 4; 9]
       # List.map (fun x -> x * x) [1; 2; 3]
       - : int list = [1; 4; 9]
     ]}
+
+    A shell block:
+
+    {@sh set-FOO=Hello,set-BAR=Bash[$ echo $FOO $BAR
+Hello Bash
+
+    ]}
+
+    A block that doesn't run:
+
+    {@text[
+      # 1
+      = 2 ?
+    ]}
 *)
 val foo : string
 
-(** {[1 + 1 = 3]} *)
+(** {@ocaml[1 + 1 = 3]} *)
 val bar : string

test/lib/test_mli_parser.ml

@@ -11,7 +11,7 @@ let mli =
 
     {[List.map (fun x -> x * x) [1; 2; 3]]}
 
-    {[
+    {@ocaml [
       # List.map (fun x -> x * x) [(1 + 9); 2; 3]
       - : int list = [100; 4; 9]
       # List.map (fun x -> x * x) [1; 2; 3]
@@ -20,7 +20,7 @@ let mli =
 *)
 val foo : string
 
-(** {[1 + 1 = 3]} *)
+(** {@ocaml [1 + 1 = 3]} *)
 val bar : string
 |}
 
@@ -53,15 +53,16 @@ let test_parse_mli =
  Block {loc: File "_none_", line 9; section: None; labels: [];
         header: Some ocaml;
         contents: ["List.map (fun x -> x * x) [1; 2; 3]"]; value: OCaml};
- Text "]}"; Text "\n\n    "; Text "{[";
+ Text "]}"; Text "\n\n    "; Text "{@"; Text "ocaml"; Text "[";
  Block {loc: File "_none_", lines 11-16; section: None; labels: [];
         header: Some ocaml;
         contents: ["# List.map (fun x -> x * x) [(1 + 9); 2; 3]";
                    "- : int list = [100; 4; 9]";
                    "# List.map (fun x -> x * x) [1; 2; 3]";
                    "- : int list = [1; 4; 9]"];
         value: Toplevel};
- Text "    ]}"; Text "\n*)\nval foo : string\n\n(** "; Text "{[";
+ Text "    ]}"; Text "\n*)\nval foo : string\n\n(** "; Text "{@";
+ Text "ocaml"; Text "[";
  Block {loc: File "_none_", line 20; section: None; labels: [];
         header: Some ocaml; contents: ["1 + 1 = 3"]; value: OCaml};
  Text "]}";|x})