integrate code block extraction into doctests script

Author: zth

analysis/src/Protocol.ml

@@ -137,6 +137,11 @@ let optWrapInQuotes s =
   | None -> None
   | Some s -> Some (wrapInQuotes s)
 
+let stringifyResult = function
+  | Ok r -> stringifyObject [("TAG", Some (wrapInQuotes "Ok")); ("_0", Some r)]
+  | Error e ->
+    stringifyObject [("TAG", Some (wrapInQuotes "Error")); ("_0", Some e)]
+
 let stringifyCompletionItem c =
   stringifyObject
     [

runtime/RescriptTools_ExtractCodeBlocks.res

@@ -7,4 +7,4 @@ type codeBlock = {
 /**
 `decodeFromJson(json)` parse JSON generated from `rescript-tools extract-codeblocks` command
 */
-external decodeFromJson: Stdlib_JSON.t => array<codeBlock> = "%identity"
+external decodeFromJson: Stdlib_JSON.t => result<array<codeBlock>, string> = "%identity"

tests/docstring_tests/DocTest.res

@@ -1,14 +1,5 @@
 open Node
 
-module Docgen = RescriptTools.Docgen
-
-type example = {
-  id: string,
-  kind: string,
-  name: string,
-  docstrings: array<string>,
-}
-
 // Only major version
 let nodeVersion =
   Process.version
@@ -49,189 +40,21 @@ let getOutput = buffer =>
 let extractDocFromFile = async file => {
   let toolsBin = Path.join([Process.cwd(), "cli", "rescript-tools.js"])
 
-  let {stdout} = await SpawnAsync.run(~command=toolsBin, ~args=["doc", file])
+  let {stdout} = await SpawnAsync.run(
+    ~command=toolsBin,
+    ~args=["extract-codeblocks", file, "--transform-assert-equal"],
+  )
 
   try {
     stdout
     ->getOutput
     ->JSON.parseOrThrow
-    ->Docgen.decodeFromJson
+    ->RescriptTools.ExtractCodeBlocks.decodeFromJson
   } catch {
-  | JsExn(_) => JsError.panic(`Failed to generate docstrings from ${file}`)
-  | _ => assert(false)
-  }
-}
-
-let getExamples = ({items}: Docgen.doc) => {
-  let rec loop = (items: list<Docgen.item>, acc: list<example>) => {
-    switch items {
-    | list{Value({docstrings, id, name}), ...rest} =>
-      loop(rest, list{{id, name, docstrings, kind: "value"}, ...acc})
-    | list{Type({docstrings, id, name}), ...rest} =>
-      loop(rest, list{{id, name, docstrings, kind: "type"}, ...acc})
-    | list{Module({id, name, docstrings, items}), ...rest} =>
-      loop(
-        list{...rest, ...List.fromArray(items)},
-        list{{id, name, docstrings, kind: "module"}, ...acc},
-      )
-    | list{ModuleType({id, name, docstrings, items}), ...rest} =>
-      loop(
-        list{...rest, ...List.fromArray(items)},
-        list{{id, name, docstrings, kind: "moduleType"}, ...acc},
-      )
-    | list{ModuleAlias({id, name, docstrings, items}), ...rest} =>
-      loop(
-        list{...rest, ...List.fromArray(items)},
-        list{{id, name, docstrings, kind: "moduleAlias"}, ...acc},
-      )
-    | list{} => acc
-    }
+  | JsExn(e) =>
+    Console.error(e)
+    JsError.panic(`Failed to extract code blocks from ${file}`)
   }
-
-  items
-  ->List.fromArray
-  ->loop(list{})
-  ->List.toArray
-  ->Array.filter(({docstrings}) => Array.length(docstrings) > 0)
-}
-
-let getCodeBlocks = example => {
-  let rec loopEndCodeBlock = (lines, acc) => {
-    switch lines {
-    | list{hd, ...rest} =>
-      if (
-        hd
-        ->String.trim
-        ->String.endsWith("```")
-      ) {
-        acc
-      } else {
-        loopEndCodeBlock(rest, list{hd, ...acc})
-      }
-    | list{} => panic(`Failed to find end of code block for ${example.kind}: ${example.id}`)
-    }
-  }
-
-  // Transform lines that contain == patterns to use assertEqual
-  let transformEqualityAssertions = code => {
-    let lines = code->String.split("\n")
-    lines
-    ->Array.mapWithIndex((line, idx) => {
-      let trimmedLine = line->String.trim
-
-      // Check if the line contains == and is not inside a comment
-      if (
-        trimmedLine->String.includes("==") &&
-        !(trimmedLine->String.startsWith("//")) &&
-        !(trimmedLine->String.startsWith("/*")) &&
-        // Not an expression line
-        !(trimmedLine->String.startsWith("if")) &&
-        !(trimmedLine->String.startsWith("|")) &&
-        !(trimmedLine->String.startsWith("let")) &&
-        !(trimmedLine->String.startsWith("~")) &&
-        !(trimmedLine->String.startsWith("module")) &&
-        !(trimmedLine->String.startsWith("->")) &&
-        !(trimmedLine->String.endsWith(","))
-      ) {
-        // Split string at == (not ===) and transform to assertEqual
-        let parts = {
-          let rec searchFrom = (currentLine: string, startIndex: int): option<(string, string)> => {
-            if startIndex >= currentLine->String.length {
-              // Base case: reached end of string without finding a suitable "=="
-              None
-            } else {
-              let lineSuffix = currentLine->String.sliceToEnd(~start=startIndex)
-              let idxEqEqInSuffix = lineSuffix->String.indexOfOpt("==")
-              let idxEqEqEqInSuffix = lineSuffix->String.indexOfOpt("===")
-
-              switch (idxEqEqInSuffix, idxEqEqEqInSuffix) {
-              | (None, _) =>
-                // No "==" found in the rest of the string.
-                None
-              | (Some(iEqEq), None) =>
-                // Found "==" but no "===" in the suffix.
-                // This "==" must be standalone.
-                // Calculate its absolute index in the original `currentLine`.
-                let actualIdx = startIndex + iEqEq
-                let left = currentLine->String.slice(~start=0, ~end=actualIdx)
-                let right = currentLine->String.sliceToEnd(~start=actualIdx + 2)
-                Some((left, right))
-              | (Some(iEqEq), Some(iEqEqEq)) =>
-                // Found both "==" and "===" in the suffix.
-                if iEqEq < iEqEqEq {
-                  // The "==" occurs before "===". This "==" is standalone.
-                  // Example: "a == b === c". In suffix "a == b === c", iEqEq is for "a ==".
-                  let actualIdx = startIndex + iEqEq
-                  let left = currentLine->String.slice(~start=0, ~end=actualIdx)
-                  let right = currentLine->String.sliceToEnd(~start=actualIdx + 2)
-                  Some((left, right))
-                } else {
-                  // iEqEq >= iEqEqEq
-                  // This means the first "==" encountered (at iEqEq relative to suffix)
-                  // is part of or comes after the first "===" encountered (at iEqEqEq relative to suffix).
-                  // Example: "a === b". In suffix "a === b", iEqEqEq is for "a ===", iEqEq is also for "a ==".
-                  // We must skip over the "===" found at iEqEqEq and search again.
-                  // The next search should start after this "===".
-                  // The "===" starts at (startIndex + iEqEqEq) in the original line. It has length 3.
-                  searchFrom(currentLine, startIndex + iEqEqEq + 3)
-                }
-              }
-            }
-          }
-          searchFrom(line, 0)
-        }
-        let parts = switch parts {
-        | Some((left, right)) if right->String.trim->String.length === 0 =>
-          Some((
-            left,
-            lines
-            ->Array.get(idx + 1)
-            ->Option.getOrThrow(
-              ~message="Expected to have an expected expression on the next line",
-            ),
-          ))
-        | _ => parts
-        }
-        switch parts {
-        | Some((left, right)) if !(right->String.includes(")") || right->String.includes("//")) =>
-          `(${left->String.trim})->assertEqual(${right->String.trim})`
-        | _ => line
-        }
-      } else {
-        line
-      }
-    })
-    ->Array.join("\n")
-  }
-
-  let rec loop = (lines: list<string>, acc: list<string>) => {
-    switch lines {
-    | list{hd, ...rest} =>
-      switch hd
-      ->String.trim
-      ->String.startsWith("```res") {
-      | true =>
-        let code = loopEndCodeBlock(rest, list{})
-        let codeString =
-          code
-          ->List.reverse
-          ->List.toArray
-          ->Array.join("\n")
-          ->transformEqualityAssertions
-        loop(rest, list{codeString, ...acc})
-      | false => loop(rest, acc)
-      }
-    | list{} => acc
-    }
-  }
-
-  example.docstrings
-  ->Array.reduce([], (acc, docstring) => acc->Array.concat(docstring->String.split("\n")))
-  ->List.fromArray
-  ->loop(list{})
-  ->List.toArray
-  ->Belt.Array.reverse
-  ->Array.join("\n\n")
 }
 
 let batchSize = OS.cpus()->Array.length
@@ -258,7 +81,10 @@ let extractExamples = async () => {
   let examples = []
   await docFiles->ArrayUtils.forEachAsyncInBatches(~batchSize, async f => {
     let doc = await extractDocFromFile(Path.join(["runtime", f]))
-    examples->Array.pushMany(doc->getExamples)
+    switch doc {
+    | Ok(doc) => examples->Array.pushMany(doc)
+    | Error(e) => Console.error(e)
+    }
   })
 
   examples->Array.sort((a, b) => String.compare(a.id, b.id))
@@ -303,7 +129,7 @@ let main = async () => {
         )
         None
       } else {
-        let code = getCodeBlocks(example)
+        let code = example.code
 
         if code->String.length === 0 {
           None