jq, but for markdown

mq-poc.R

#!/usr/bin/env Rapp
#| name: mq-poc
#| title: Markdown query proof of concept
#| description: Query Markdown headings and section bodies using CommonMark source positions.

#| description: Print CommonMark XML with source positions.
xml <- FALSE

#| description: Print the parsed Markdown node table.
nodes <- FALSE

#| description: Trim blank lines around section body output.
trim <- TRUE

#| description: Query to run, e.g. `h1 | text`, `paragraph | text`, `h1 | first | body`, or `h1 | nth(2) | source`.
query <- NULL

#| description: Markdown file path.
path <- NULL

stopifnot(requireNamespace("commonmark", quietly = TRUE))
stopifnot(requireNamespace("xml2", quietly = TRUE))

read_markdown <- function(path) {
  stopifnot(length(path) == 1L)
  stopifnot(file.exists(path))

  lines <- readLines(path, warn = FALSE, encoding = "UTF-8")
  paste(enc2utf8(lines), collapse = "\n")
}

newline_byte_positions <- function(text) {
  matches <- gregexpr("\n", text, fixed = TRUE, useBytes = TRUE)[[1L]]
  if (identical(matches, -1L)) {
    integer()
  } else {
    as.integer(matches)
  }
}

sourcepos_to_byte_range <- function(sourcepos, line_start_byte) {
  pos <- strsplit(sourcepos, "[-:]", perl = TRUE)[[1L]]
  stopifnot(length(pos) == 4L)

  pos <- as.integer(pos)
  start <- line_start_byte[[pos[[1L]]]] + pos[[2L]] - 1L
  end <- line_start_byte[[pos[[3L]]]] + pos[[4L]] - 1L

  c(start = start, end = end)
}

source_slice <- function(text, start, end) {
  if (start > end) {
    return("")
  }

  bytes <- charToRaw(text)
  rawToChar(bytes[start:end])
}

parse_markdown <- function(text) {
  xml_text <- commonmark::markdown_xml(
    text,
    sourcepos = TRUE,
    extensions = TRUE,
    footnotes = TRUE,
    normalize = TRUE
  )

  doc <- xml2::read_xml(charToRaw(xml_text), encoding = "UTF-8")
  xml2::xml_ns_strip(doc)

  elements <- xml2::xml_find_all(doc, "//*[@sourcepos]")
  line_start_byte <- c(1L, newline_byte_positions(text) + 1L)
  ranges <- lapply(
    xml2::xml_attr(elements, "sourcepos"),
    sourcepos_to_byte_range,
    line_start_byte = line_start_byte
  )
  ranges <- do.call(rbind, ranges)

  data.frame(
    node = seq_along(elements),
    type = xml2::xml_name(elements),
    level = suppressWarnings(as.integer(xml2::xml_attr(elements, "level"))),
    text = xml2::xml_text(elements),
    start = ranges[, "start"],
    end = ranges[, "end"],
    stringsAsFactors = FALSE
  )
}

new_selection <- function(table, rows, source, trim) {
  structure(
    list(table = table, rows = rows, source = source, trim = trim),
    class = "mq_selection"
  )
}

select_nodes <- function(context, type, level = NULL) {
  keep <- context$table$type == type
  if (!is.null(level)) {
    keep <- keep & context$table$level == level
  }

  new_selection(
    context$table,
    context$table[keep, , drop = FALSE],
    context$source,
    context$trim
  )
}

select_index <- function(selection, index) {
  stopifnot(inherits(selection, "mq_selection"))
  stopifnot(length(index) == 1L, index >= 1L, index <= nrow(selection$rows))

  new_selection(
    selection$table,
    selection$rows[index, , drop = FALSE],
    selection$source,
    selection$trim
  )
}

selection_text <- function(selection) {
  stopifnot(inherits(selection, "mq_selection"))
  selection$rows$text
}

selection_source <- function(selection) {
  stopifnot(inherits(selection, "mq_selection"))

  vapply(
    seq_len(nrow(selection$rows)),
    function(i) {
      source_slice(
        selection$source,
        selection$rows$start[[i]],
        selection$rows$end[[i]]
      )
    },
    character(1L)
  )
}

selection_body <- function(selection) {
  stopifnot(inherits(selection, "mq_selection"))
  stopifnot(all(selection$rows$type == "heading"))

  table <- selection$table
  source_len <- length(charToRaw(selection$source))

  vapply(
    seq_len(nrow(selection$rows)),
    function(i) {
      row <- selection$rows[i, , drop = FALSE]
      next_heading <- table[
        table$node > row$node[[1L]] &
          table$type == "heading" &
          table$level <= row$level[[1L]],
        ,
        drop = FALSE
      ]

      start <- row$end[[1L]] + 1L
      end <- if (nrow(next_heading)) {
        next_heading$start[[1L]] - 1L
      } else {
        source_len
      }

      body <- source_slice(selection$source, start, end)
      if (selection$trim) {
        body <- sub("^\\n+", "", body)
        body <- sub("\\n+$", "", body)
      }
      body
    },
    character(1L)
  )
}

commonmark_node_types <- function() {
  c(
    "document",
    "block_quote",
    "list",
    "item",
    "paragraph",
    "heading",
    "thematic_break",
    "code_block",
    "html_block",
    "custom_block",
    "text",
    "softbreak",
    "linebreak",
    "code",
    "html_inline",
    "custom_inline",
    "emph",
    "strong",
    "link",
    "image",
    "table",
    "table_header",
    "table_row",
    "table_cell"
  )
}

node_binding_name <- function(type) {
  switch(
    type,
    "text" = "text_node",
    type
  )
}

new_query_env <- function(table, source, trim) {
  context <- list(table = table, source = source, trim = trim)
  env <- new.env(parent = emptyenv())

  env$`|` <- function(lhs, rhs) {
    stopifnot(is.function(rhs))
    rhs(lhs)
  }

  env$first <- function(selection) select_index(selection, 1L)
  env$last <- function(selection) select_index(selection, nrow(selection$rows))
  env$nth <- function(index) {
    stopifnot(length(index) == 1L)
    force(index)
    function(selection) select_index(selection, index)
  }

  env$text <- selection_text
  env$body <- selection_body
  env$source <- selection_source
  env$node <- function(type) select_nodes(context, type)

  for (level in 1:6) {
    env[[paste0("h", level)]] <- select_nodes(context, "heading", level)
  }

  for (type in commonmark_node_types()) {
    name <- node_binding_name(type)
    if (!exists(name, env, inherits = FALSE)) {
      env[[name]] <- select_nodes(context, type)
    }
  }
  env$blockquote <- env$block_quote

  env
}

eval_query <- function(query, table, source, trim) {
  expr <- parse(text = query, keep.source = FALSE)
  stopifnot(length(expr) == 1L)

  eval(expr[[1L]], new_query_env(table, source, trim))
}

text <- read_markdown(path)

if (xml) {
  commonmark::markdown_xml(
    text,
    sourcepos = TRUE,
    extensions = TRUE,
    footnotes = TRUE,
    normalize = TRUE
  ) |>
    cat()
  quit(save = "no")
}

table <- parse_markdown(text)

if (nodes) {
  print(table, row.names = FALSE)
  quit(save = "no")
}

result <- eval_query(query, table, text, trim)

if (inherits(result, "mq_selection")) {
  result <- selection_source(result)
}

writeLines(result)