From 59f2ba9fae27dee37758b101d704e2b6cc3d4ec7 Mon Sep 17 00:00:00 2001
From: Manuel Puyol <manuelpuyol@github.com>
Date: Wed, 17 Feb 2021 09:36:20 -0600
Subject: [PATCH] run before and after tokenize hooks

---
 src/components/Highlight.js    | 25 ++++++++++++++++++++++++-
 src/types.js                   |  8 +++++++-
 src/vendor/prism/prism-core.js |  1 +
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/components/Highlight.js b/src/components/Highlight.js
index 788c38b..8196461 100755
--- a/src/components/Highlight.js
+++ b/src/components/Highlight.js
@@ -12,8 +12,10 @@ import type {
   TokenInputProps,
   TokenOutputProps,
   RenderProps,
+  PrismGrammar,
   PrismLib,
   PrismTheme,
+  PrismToken,
 } from "../types";
 
 type Props = {
@@ -121,6 +123,25 @@ class Highlight extends Component<Props, *> {
     return output;
   };
 
+  tokenize = (
+    Prism: PrismLib,
+    code: string,
+    grammar: PrismGrammar,
+    language: Language
+  ): Array<PrismToken> => {
+    const env = {
+      code,
+      grammar,
+      language,
+    };
+
+    Prism.hooks.run("before-tokenize", env);
+    env.tokens = Prism.tokenize(env.code, env.grammar, env.language);
+    Prism.hooks.run("after-tokenize", env);
+
+    return env.tokens;
+  };
+
   render() {
     const { Prism, language, code, children } = this.props;
 
@@ -128,7 +149,9 @@ class Highlight extends Component<Props, *> {
 
     const grammar = Prism.languages[language];
     const mixedTokens =
-      grammar !== undefined ? Prism.tokenize(code, grammar, language) : [code];
+      grammar !== undefined
+        ? this.tokenize(Prism, code, grammar, language)
+        : [code];
     const tokens = normalizeTokens(mixedTokens);
 
     return children({
diff --git a/src/types.js b/src/types.js
index 21c792b..f1a4d13 100755
--- a/src/types.js
+++ b/src/types.js
@@ -5,7 +5,7 @@ import includedLangs from "./vendor/prism/includeLangs";
 
 export type Language = $Keys<typeof includedLangs>;
 
-type PrismGrammar = {
+export type PrismGrammar = {
   [key: string]: mixed,
 };
 
@@ -37,6 +37,12 @@ export type PrismLib = {
     grammar: PrismGrammar,
     language: Language
   ) => string,
+  hooks: {
+    run: (
+      name: string,
+      env: { code: string, grammar: PrismGrammar, language: Language }
+    ) => void,
+  },
 };
 
 export type StyleObj = {
diff --git a/src/vendor/prism/prism-core.js b/src/vendor/prism/prism-core.js
index 61189a4..029da46 100755
--- a/src/vendor/prism/prism-core.js
+++ b/src/vendor/prism/prism-core.js
@@ -339,6 +339,7 @@ var Prism = (function () {
 
     hooks: {
       add: function () {},
+      run: function (name, env) {},
     },
 
     tokenize: function (text, grammar, language) {