From 8f709e20335155f9d6821c41188e06e0166e908e Mon Sep 17 00:00:00 2001
From: Mariusz Gronczewski <xani666@gmail.com>
Date: Sun, 5 Apr 2020 22:15:09 +0200
Subject: [PATCH] expose chroma's CSS stylesheet

---
 README.md        |  9 ++++++---
 renderer.go      | 20 ++++++++++++++++++++
 renderer_test.go | 18 ++++++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 3053747..56bab5b 100644
--- a/README.md
+++ b/README.md
@@ -55,13 +55,14 @@ By default when no language information is written in the code block, this
 renderer will try to auto-detect the used language. This option disables
 this behavior and will fallback to a sane default when no language
 information is available.
+- `EmbedCSS()` 
+This option will embed CSS needed for chroma's `html.WithClasses()` at the beginning of blackfriday document. 
+CSS can also be extracted separately by calling `Renderer`'s.`ChromaCSS(w)` method, which will return styleshet for currently set style 
 - `Extend(bf.Renderer)`  
 This option allows to define the base blackfriday that will be extended.
 - `ChromaOptions(...html.Option)`  
 This option allows you to pass Chroma's html options in the renderer. Such
 options can be found [here](https://github.com/alecthomas/chroma#the-html-formatter).
-There is currently an issue with the `html.WithClasses()` option as it expects
-the CSS classes to be written separately. I'll come up with a fix later.
 
 ### Option examples
 
@@ -92,6 +93,8 @@ r := bfchroma.NewRenderer(bfchroma.Style("dracula"))
 r = bfchroma.NewRenderer(bfchroma.ChromaStyle(styles.Dracula))
 ```
 
+
+
 ## Examples
 
 ```go
@@ -194,7 +197,7 @@ r := bfchroma.NewRenderer(
 var css template.CSS
 
 b := new(bytes.Buffer)
-if err := r.Formatter.WriteCSS(b, r.Style); err != nil {
+if err := r.ChromaCSS(b); err != nil {
 	logrus.WithError(err).Warning("Couldn't write CSS")
 }
 css = template.CSS(b.String())
diff --git a/renderer.go b/renderer.go
index 7cc5d72..e9e997a 100644
--- a/renderer.go
+++ b/renderer.go
@@ -40,6 +40,13 @@ func WithoutAutodetect() Option {
 	}
 }
 
+// EmbedCSS will embed CSS needed for html.WithClasses() in beginning of the document
+func EmbedCSS() Option {
+	return func(r *Renderer) {
+		r.embedCSS = true
+	}
+}
+
 // ChromaOptions allows to pass Chroma html.Option such as Standalone()
 // WithClasses(), ClassPrefix(prefix)...
 func ChromaOptions(options ...html.Option) Option {
@@ -101,11 +108,19 @@ type Renderer struct {
 	ChromaOptions []html.Option
 	Style         *chroma.Style
 	Formatter     *html.Formatter
+	embedCSS bool
 }
 
 // RenderNode satisfies the Renderer interface
 func (r *Renderer) RenderNode(w io.Writer, node *bf.Node, entering bool) bf.WalkStatus {
 	switch node.Type {
+	case bf.Document:
+		if entering && r.embedCSS {
+			w.Write([]byte("<style>"))
+			r.Formatter.WriteCSS(w, r.Style)
+			w.Write([]byte("</style>"))
+		}
+		return r.Base.RenderNode(w, node, entering)
 	case bf.CodeBlock:
 		if err := r.RenderWithChroma(w, node.Literal, node.CodeBlockData); err != nil {
 			return r.Base.RenderNode(w, node, entering)
@@ -125,3 +140,8 @@ func (r *Renderer) RenderHeader(w io.Writer, ast *bf.Node) {
 func (r *Renderer) RenderFooter(w io.Writer, ast *bf.Node) {
 	r.Base.RenderFooter(w, ast)
 }
+
+// ChromaCSS returns CSS used with chroma's html.WithClasses() option
+func (r *Renderer) ChromaCSS(w io.Writer) error {
+	return r.Formatter.WriteCSS(w, r.Style)
+}
\ No newline at end of file
diff --git a/renderer_test.go b/renderer_test.go
index 765eae0..0912c35 100644
--- a/renderer_test.go
+++ b/renderer_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/alecthomas/chroma/formatters/html"
 	"github.com/alecthomas/chroma/styles"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	bf "github.com/russross/blackfriday/v2"
 )
 
@@ -147,6 +148,23 @@ func TestRender(t *testing.T) {
 	assert.Contains(t, string(h), "<pre")
 }
 
+func TestRender_EmbedCSS(t *testing.T) {
+	r := NewRenderer(EmbedCSS())
+	h := bf.Run([]byte(""), bf.WithRenderer(r))
+	assert.Contains(t, string(h),"<style>")
+	assert.Contains(t, string(h),".chroma")
+	assert.Contains(t, string(h),"</style>")
+}
+
+func TestRenderer_ChromaCSS(t *testing.T) {
+	r := NewRenderer()
+	var w bytes.Buffer
+	err := r.ChromaCSS(&w)
+	require.NoError(t,err)
+	assert.Contains(t, w.String(),".chroma")
+
+}
+
 func ExampleNewRenderer() {
 	// Complex example on how to initialize the renderer
 	md := "```go\npackage main\n\nfunc main() {\n}\n```"