From 7d779408dbe39028feb44f938debaf89b1af463a Mon Sep 17 00:00:00 2001
From: Siavash Askari Nasr <ciavash@protonmail.com>
Date: Mon, 16 May 2022 18:05:36 +0430
Subject: [PATCH] [html formatter] Add option to let users provide custom CSS
 styles

Example: WithCustomCSS(map[chroma.TokenType]string{chroma.Line: `display: inline;`})
---
 formatters/html/html.go      | 14 +++++++++++---
 formatters/html/html_test.go | 12 ++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/formatters/html/html.go b/formatters/html/html.go
index e2fb6f995..52577da2f 100644
--- a/formatters/html/html.go
+++ b/formatters/html/html.go
@@ -25,6 +25,13 @@ func WithClasses(b bool) Option { return func(f *Formatter) { f.Classes = b } }
 // WithAllClasses disables an optimisation that omits redundant CSS classes.
 func WithAllClasses(b bool) Option { return func(f *Formatter) { f.allClasses = b } }
 
+// WithCustomCSS sets user's custom CSS styles.
+func WithCustomCSS(css map[chroma.TokenType]string) Option {
+	return func(f *Formatter) {
+		f.customCSS = css
+	}
+}
+
 // TabWidth sets the number of characters for a tab. Defaults to 8.
 func TabWidth(width int) Option { return func(f *Formatter) { f.tabWidth = width } }
 
@@ -160,6 +167,7 @@ type Formatter struct {
 	prefix              string
 	Classes             bool // Exported field to detect when classes are being used
 	allClasses          bool
+	customCSS           map[chroma.TokenType]string
 	preWrapper          PreWrapper
 	tabWidth            int
 	wrapLongLines       bool
@@ -435,15 +443,15 @@ func (f *Formatter) styleToCSS(style *chroma.Style) map[chroma.TokenType]string
 		if t != chroma.Background {
 			entry = entry.Sub(bg)
 		}
-		if !f.allClasses && entry.IsZero() {
+		if !f.allClasses && entry.IsZero() && f.customCSS[t] == `` {
 			continue
 		}
-		classes[t] = StyleEntryToCSS(entry)
+		classes[t] = f.customCSS[t] + StyleEntryToCSS(entry)
 	}
 	classes[chroma.Background] += f.tabWidthStyle()
 	classes[chroma.PreWrapper] += classes[chroma.Background] + `;`
 	// Make PreWrapper a grid to show highlight style with full width.
-	if len(f.highlightRanges) > 0 {
+	if len(f.highlightRanges) > 0 && f.customCSS[chroma.PreWrapper] == `` {
 		classes[chroma.PreWrapper] += `display: grid;`
 	}
 	// Make PreWrapper wrap long lines.
diff --git a/formatters/html/html_test.go b/formatters/html/html_test.go
index f2539f1f6..0f133474f 100644
--- a/formatters/html/html_test.go
+++ b/formatters/html/html_test.go
@@ -108,6 +108,18 @@ func TestTableLineNumberNewlines(t *testing.T) {
 </span>`)
 }
 
+func TestWithCustomCSS(t *testing.T) {
+	f := New(WithClasses(false), WithCustomCSS(map[chroma.TokenType]string{chroma.Line: `display: inline;`}))
+	it, err := lexers.Get("bash").Tokenise(nil, "echo FOO")
+	assert.NoError(t, err)
+
+	var buf bytes.Buffer
+	err = f.Format(&buf, styles.Fallback, it)
+	assert.NoError(t, err)
+
+	assert.Regexp(t, `<span style="display:flex;display:inline;"><span><span style=".*">echo</span> FOO</span></span>`, buf.String())
+}
+
 func TestWrapLongLines(t *testing.T) {
 	f := New(WithClasses(false), WrapLongLines(true))
 	it, err := lexers.Get("go").Tokenise(nil, "package main\nfunc main()\n{\nprintln(\"hello world\")\n}\n")