DSchau · DSchau · Oct 31, 2018 · Nov 1, 2018 · Nov 1, 2018 · Nov 1, 2018
diff --git a/packages/gatsby-remark-prismjs/highlight-line-range.js b/packages/gatsby-remark-prismjs/highlight-line-range.js
@@ -0,0 +1,99 @@
+"use strict"
+
+const COMMENT_START = /(#|\/\/|\{\/\*|\/\*+|<!--)/
+const COMMENT_END = /(-->|\*\/\}|\*\/)?/
+const DIRECTIVE = /highlight-(next-line|line|start|end)/
+const END_DIRECTIVE = /highlight-end/
+const plainTextWithLFTest = /<span class="token plain-text">[^<]*\n[^<]*<\/span>/g
+
+const stripComment = line =>
+  line.replace(
+    new RegExp(
+      `\\s*${COMMENT_START.source}\\s*${DIRECTIVE.source}\\s*${
+        COMMENT_END.source
+      }`
+    ),
+    ""
+  )
+
+const wrap = line =>
+  [
+    `<span class="gatsby-highlight-code-line">`,
+    `${stripComment(line)}\n`,
+    `</span>`,
+  ].join("")
+
+const getHighlights = (line, code, index) => {
+  const [, directive] = line.match(DIRECTIVE)
+
+  switch (directive) {
+    case "next-line":
+      return [
+        {
+          code: code[index + 1],
+          highlighted: true,
+        },
+        index + 1,
+      ]
+
+    case "start":
+      const endIndex = code.findIndex(line => END_DIRECTIVE.test(line))
+      const end = endIndex === -1 ? code.length : endIndex
+      const highlighted = code.slice(index + 1, end).map(line => ({
+        code: wrap(line),
+        highlighted: true,
+      }))
+      return [highlighted, end]
+
+    case "line":
+    default:
+      return [
+        {
+          code: wrap(line),
+          highlighted: true,
+        },
+        index,
+      ]
+  }
+}
+
+module.exports = function highlightLineRange(code, highlights = []) {
+  let highlighted = []
+  const split = code.split("\n")
+
+  if (highlights.length > 0) {
+    // HACK split plain-text spans with line separators inside into multiple plain-text spans
+    // separatered by line separator - this fixes line highlighting behaviour for jsx
+    code = code.replace(plainTextWithLFTest, match =>
+      match.replace(/\n/g, `</span>\n<span class="token plain-text">`)
+    )
+    return split.map((line, i) => {
+      if (highlights.includes(i + 1)) {
+        return {
+          highlighted: true,
+          code: line,
+        }
+      }
+
+      return {
+        code: line,
+      }
+    })
+  }
+
+  for (let i = 0; i < split.length; i++) {
+    const line = split[i]
+
+    if (DIRECTIVE.test(line)) {
+      const [highlights, index] = getHighlights(line, split, i)
+      highlighted = highlighted.concat(highlights)
+      i = index
+    } else {
+      highlighted.push({
+        code: line,
+      })
+    }
+  }
+
+  return highlighted
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-html.html b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-html.html
@@ -0,0 +1,9 @@
+<!doctype html>
+<html>
+  <head>
+    <title>Helloooooooooo</title>
+  </head>
+  <body>
+    <h1>hello world</h1> <!-- highlight-line -->
+  </body>
+</html>
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-jsx-comment.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-jsx-comment.js
@@ -0,0 +1,9 @@
+import React from "react"
+
+export default function Button() {
+  return (
+    <div>
+      <button>sup</button> {/* highlight-line */}
+    </div>
+  )
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-kitchen-sink.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-kitchen-sink.js
@@ -0,0 +1,28 @@
+import React, { Component } from "react" // highlight-line
+
+export default class Counter extends Component {
+  /* highlight-start */
+  state = {
+    count: 0,
+  }
+  // highlight-end
+
+  updateCount = () => {
+    this.setState(state => ({
+      // highlight-next-line
+      count: state.count + 1,
+    }))
+  }
+
+  render() {
+    const { count } = this.state // highlight-line
+    return (
+      <div>
+        <span>clicked {count}</span>
+        {/* highlight-start */}
+        <button onClick={this.updateCount}>Click me</button>
+        {/* highlight-end */}
+      </div>
+    )
+  }
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-line.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-line.js
@@ -0,0 +1,3 @@
+function test() {
+  return "hello world" // highlight-line
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-next-line.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-next-line.js
@@ -0,0 +1,4 @@
+function test() {
+  // highlight-next-line
+  return "hello world"
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-start-end.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-start-end.js
@@ -0,0 +1,6 @@
+function test() {
+  // highlight-start
+  var a = "b"
+  return a + "hello world"
+  // highlight-end
+}
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-start-without-end.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-start-without-end.js
@@ -0,0 +1,3 @@
+// highlight-start
+var a = "b"
+return a + "hello world"
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-yaml.yaml b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/highlight-yaml.yaml
@@ -0,0 +1,3 @@
+- title: uno
+- title: dos
+- title: catorce # highlight-line
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/index.js b/packages/gatsby-remark-prismjs/src/__tests__/__fixtures__/index.js
@@ -0,0 +1,14 @@
+const fs = require("fs-extra")
+const path = require("path")
+
+const base = __dirname
+
+module.exports = fs.readdirSync(base).reduce((lookup, file) => {
+  if (file !== "index.js") {
+    const name = file
+      .replace(/-(\w)/g, (_, char) => char.toUpperCase())
+      .replace(/\..+/, "")
+    lookup[name] = fs.readFileSync(path.join(base, file), "utf8")
+  }
+  return lookup
+}, {})
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/__snapshots__/highlight-line-range.js.snap b/packages/gatsby-remark-prismjs/src/__tests__/__snapshots__/highlight-line-range.js.snap
@@ -0,0 +1,20 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`highlighting a line range highlight-line highlights line 1`] = `"  return 'hello world'"`;
+
+exports[`highlighting a line range highlight-next-line highlights correct line 1`] = `"  return 'hello world'"`;
+
+exports[`highlighting a line range highlight-start / highlight-end highlights correct lines 1`] = `
+"  var a = 'b'
+  return a + 'hello world'"
+`;
+
+exports[`highlighting a line range highlight-start / highlight-end highlights without end directive 1`] = `
+"var a = 'b'
+return a + 'hello world'
+"
+`;
+
+exports[`highlighting a line range jsx comment highlights comment line 1`] = `"      <button>sup</button>"`;
+
+exports[`highlighting a line range yaml highlights yaml 1`] = `"- title: catorce"`;
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/highlight-line-range.js b/packages/gatsby-remark-prismjs/src/__tests__/highlight-line-range.js
@@ -0,0 +1,93 @@
+const highlightLineRange = require("../highlight-line-range")
+const fixtures = require("./__fixtures__")
+
+const output = highlighted => highlighted.map(({ code }) => code).join("\n")
+const getHighlighted = lines => lines.filter(line => line.highlighted)
+
+describe("highlighting a line range", () => {
+  describe("highlight-line", () => {
+    it("strips directive", () => {
+      const highlights = highlightLineRange(fixtures.highlightLine)
+      expect(output(highlights)).not.toContain("highlight-line")
+    })
+
+    it("highlights line", () => {
+      const highlights = highlightLineRange(fixtures.highlightLine)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+  })
+
+  describe("highlight-next-line", () => {
+    it("strips directive", () => {
+      const highlights = highlightLineRange(fixtures.highlightNextLine)
+
+      expect(output(highlights)).not.toContain("highlight-next-line")
+    })
+
+    it("highlights correct line", () => {
+      const highlights = highlightLineRange(fixtures.highlightNextLine)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+  })
+
+  describe("highlight-start / highlight-end", () => {
+    it("strips directives", () => {
+      const highlights = highlightLineRange(fixtures.highlightStartEnd)
+
+      const code = output(highlights)
+      ;["highlight-start", "highlight-end"].forEach(directive => {
+        expect(code).not.toContain(directive)
+      })
+    })
+
+    it("highlights correct lines", () => {
+      const highlights = highlightLineRange(fixtures.highlightStartEnd)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+
+    it("highlights without end directive", () => {
+      const highlights = highlightLineRange(fixtures.highlightStartWithoutEnd)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+  })
+
+  describe("jsx comment", () => {
+    it("removes directive", () => {
+      const highlights = highlightLineRange(fixtures.highlightJsxComment)
+
+      expect(output(highlights)).not.toContain("highlight-line")
+    })
+
+    it("highlights comment line", () => {
+      const highlights = highlightLineRange(fixtures.highlightJsxComment)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+  })
+
+  describe("yaml", () => {
+    it("strips directive", () => {
+      const highlights = highlightLineRange(fixtures.highlightYaml)
+
+      expect(highlights).not.toContain("highlight-line")
+    })
+
+    it("highlights yaml", () => {
+      const highlights = highlightLineRange(fixtures.highlightYaml)
+
+      expect(output(getHighlighted(highlights))).toMatchSnapshot()
+    })
+  })
+
+  describe("kitchen sink", () => {
+    it.skip("strips directives", () => {
+      const highlights = highlightLineRange(fixtures.highlightKitchenSink)
+
+      console.log(output(highlights))
+    })
+  })
+})
diff --git a/packages/gatsby-remark-prismjs/src/highlight-code.js b/packages/gatsby-remark-prismjs/src/highlight-code.js
@@ -2,8 +2,7 @@ const Prism = require(`prismjs`)
 const _ = require(`lodash`)
 
 const loadPrismLanguage = require(`./load-prism-language`)
-
-const plainTextWithLFTest = /<span class="token plain-text">[^<]*\n[^<]*<\/span>/g
+const highlightLineRange = require(`./highlight-line-range`)
 
 module.exports = (language, code, lineNumbersHighlight = []) => {
   // (Try to) load languages on demand.
@@ -22,35 +21,19 @@ module.exports = (language, code, lineNumbersHighlight = []) => {
 
   const grammar = Prism.languages[language]
 
-  let highlightedCode = Prism.highlight(code, grammar, language)
-  if (lineNumbersHighlight.length > 0) {
-    // HACK split plain-text spans with line separators inside into multiple plain-text spans
-    // separatered by line separator - this fixes line highlighting behaviour for jsx
-    highlightedCode = highlightedCode.replace(plainTextWithLFTest, match =>
-      match.replace(/\n/g, `</span>\n<span class="token plain-text">`)
-    )
+  const highlighted = Prism.highlight(code, grammar, language)
+  const codeSplits = highlightLineRange(highlighted, lineNumbersHighlight)
 
-    const codeSplits = highlightedCode.split(`\n`).map((split, i) => {
-      if (_.includes(lineNumbersHighlight, i + 1)) {
-        return {
-          highlighted: true,
-          code: `<span class="gatsby-highlight-code-line">${split}\n</span>`,
-        }
-      } else {
-        return { code: split }
-      }
-    })
+  let finalCode = ``
 
-    highlightedCode = ``
-    const lastIdx = codeSplits.length - 1
-    // Don't add back the new line character after highlighted lines
-    // as they need to be display: block and full-width.
-    codeSplits.forEach((split, idx) => {
-      split.highlighted
-        ? (highlightedCode += split.code)
-        : (highlightedCode += `${split.code}${idx == lastIdx ? `` : `\n`}`)
-    })
-  }
+  const lastIdx = codeSplits.length - 1
+  // Don't add back the new line character after highlighted lines
+  // as they need to be display: block and full-width.
+  codeSplits.forEach((split, idx) => {
+    split.highlighted
+      ? (finalCode += split.code)
+      : (finalCode += `${split.code}${idx == lastIdx ? `` : `\n`}`)
+  })
 
-  return highlightedCode
+  return finalCode
 }