feat: Add code block indentation token support

ansi/blockstack.go

@@ -59,10 +59,10 @@
 
 // Width returns the available rendering width.
 func (s BlockStack) Width(ctx RenderContext) uint {
-	if s.Indent()*s.Margin() > uint(ctx.options.WordWrap) {
+	if s.Indent()+s.Margin()*2 > uint(ctx.options.WordWrap) {
 		return 0
 	}
-	return uint(ctx.options.WordWrap) - s.Indent()*s.Margin()
+	return uint(ctx.options.WordWrap) - s.Indent() - s.Margin()*2
 }
 
 // Parent returns the current BlockElement's parent.

ansi/codeblock.go

@@ -1,13 +1,13 @@
 package ansi
 
 import (
+	"bytes"
 	"io"
 	"sync"
 
 	"github.com/alecthomas/chroma/v2"
 	"github.com/alecthomas/chroma/v2/quick"
 	"github.com/alecthomas/chroma/v2/styles"
-	"github.com/muesli/reflow/indent"
 	"github.com/muesli/termenv"
 )
 
@@ -62,82 +62,90 @@ func chromaStyle(style StylePrimitive) string {
 
 func (e *CodeBlockElement) Render(w io.Writer, ctx RenderContext) error {
 	bs := ctx.blockStack
-
-	var indentation uint
-	var margin uint
 	rules := ctx.options.Styles.CodeBlock
-	if rules.Indent != nil {
-		indentation = *rules.Indent
-	}
-	if rules.Margin != nil {
-		margin = *rules.Margin
+
+	be := BlockElement{
+		Block: &bytes.Buffer{},
+		Style: rules.StyleBlock,
 	}
-	theme := rules.Theme
+	bs.Push(be)
+	return nil
+}
+
+func (e *CodeBlockElement) Finish(w io.Writer, ctx RenderContext) error {
+	bs := ctx.blockStack
+	rules := bs.Current().Style
+
+	cb := ctx.options.Styles.CodeBlock
+	theme := cb.Theme
+	chromaRules := cb.Chroma
 
-	if rules.Chroma != nil && ctx.options.ColorProfile != termenv.Ascii {
+	if chromaRules != nil && ctx.options.ColorProfile != termenv.Ascii {
 		theme = chromaStyleTheme
 		mutex.Lock()
 		// Don't register the style if it's already registered.
 		_, ok := styles.Registry[theme]
 		if !ok {
 			styles.Register(chroma.MustNewStyle(theme,
 				chroma.StyleEntries{
-					chroma.Text:                chromaStyle(rules.Chroma.Text),
-					chroma.Error:               chromaStyle(rules.Chroma.Error),
-					chroma.Comment:             chromaStyle(rules.Chroma.Comment),
-					chroma.CommentPreproc:      chromaStyle(rules.Chroma.CommentPreproc),
-					chroma.Keyword:             chromaStyle(rules.Chroma.Keyword),
-					chroma.KeywordReserved:     chromaStyle(rules.Chroma.KeywordReserved),
-					chroma.KeywordNamespace:    chromaStyle(rules.Chroma.KeywordNamespace),
-					chroma.KeywordType:         chromaStyle(rules.Chroma.KeywordType),
-					chroma.Operator:            chromaStyle(rules.Chroma.Operator),
-					chroma.Punctuation:         chromaStyle(rules.Chroma.Punctuation),
-					chroma.Name:                chromaStyle(rules.Chroma.Name),
-					chroma.NameBuiltin:         chromaStyle(rules.Chroma.NameBuiltin),
-					chroma.NameTag:             chromaStyle(rules.Chroma.NameTag),
-					chroma.NameAttribute:       chromaStyle(rules.Chroma.NameAttribute),
-					chroma.NameClass:           chromaStyle(rules.Chroma.NameClass),
-					chroma.NameConstant:        chromaStyle(rules.Chroma.NameConstant),
-					chroma.NameDecorator:       chromaStyle(rules.Chroma.NameDecorator),
-					chroma.NameException:       chromaStyle(rules.Chroma.NameException),
-					chroma.NameFunction:        chromaStyle(rules.Chroma.NameFunction),
-					chroma.NameOther:           chromaStyle(rules.Chroma.NameOther),
-					chroma.Literal:             chromaStyle(rules.Chroma.Literal),
-					chroma.LiteralNumber:       chromaStyle(rules.Chroma.LiteralNumber),
-					chroma.LiteralDate:         chromaStyle(rules.Chroma.LiteralDate),
-					chroma.LiteralString:       chromaStyle(rules.Chroma.LiteralString),
-					chroma.LiteralStringEscape: chromaStyle(rules.Chroma.LiteralStringEscape),
-					chroma.GenericDeleted:      chromaStyle(rules.Chroma.GenericDeleted),
-					chroma.GenericEmph:         chromaStyle(rules.Chroma.GenericEmph),
-					chroma.GenericInserted:     chromaStyle(rules.Chroma.GenericInserted),
-					chroma.GenericStrong:       chromaStyle(rules.Chroma.GenericStrong),
-					chroma.GenericSubheading:   chromaStyle(rules.Chroma.GenericSubheading),
-					chroma.Background:          chromaStyle(rules.Chroma.Background),
+					chroma.Text:                chromaStyle(chromaRules.Text),
+					chroma.Error:               chromaStyle(chromaRules.Error),
+					chroma.Comment:             chromaStyle(chromaRules.Comment),
+					chroma.CommentPreproc:      chromaStyle(chromaRules.CommentPreproc),
+					chroma.Keyword:             chromaStyle(chromaRules.Keyword),
+					chroma.KeywordReserved:     chromaStyle(chromaRules.KeywordReserved),
+					chroma.KeywordNamespace:    chromaStyle(chromaRules.KeywordNamespace),
+					chroma.KeywordType:         chromaStyle(chromaRules.KeywordType),
+					chroma.Operator:            chromaStyle(chromaRules.Operator),
+					chroma.Punctuation:         chromaStyle(chromaRules.Punctuation),
+					chroma.Name:                chromaStyle(chromaRules.Name),
+					chroma.NameBuiltin:         chromaStyle(chromaRules.NameBuiltin),
+					chroma.NameTag:             chromaStyle(chromaRules.NameTag),
+					chroma.NameAttribute:       chromaStyle(chromaRules.NameAttribute),
+					chroma.NameClass:           chromaStyle(chromaRules.NameClass),
+					chroma.NameConstant:        chromaStyle(chromaRules.NameConstant),
+					chroma.NameDecorator:       chromaStyle(chromaRules.NameDecorator),
+					chroma.NameException:       chromaStyle(chromaRules.NameException),
+					chroma.NameFunction:        chromaStyle(chromaRules.NameFunction),
+					chroma.NameOther:           chromaStyle(chromaRules.NameOther),
+					chroma.Literal:             chromaStyle(chromaRules.Literal),
+					chroma.LiteralNumber:       chromaStyle(chromaRules.LiteralNumber),
+					chroma.LiteralDate:         chromaStyle(chromaRules.LiteralDate),
+					chroma.LiteralString:       chromaStyle(chromaRules.LiteralString),
+					chroma.LiteralStringEscape: chromaStyle(chromaRules.LiteralStringEscape),
+					chroma.GenericDeleted:      chromaStyle(chromaRules.GenericDeleted),
+					chroma.GenericEmph:         chromaStyle(chromaRules.GenericEmph),
+					chroma.GenericInserted:     chromaStyle(chromaRules.GenericInserted),
+					chroma.GenericStrong:       chromaStyle(chromaRules.GenericStrong),
+					chroma.GenericSubheading:   chromaStyle(chromaRules.GenericSubheading),
+					chroma.Background:          chromaStyle(chromaRules.Background),
 				}))
 		}
 		mutex.Unlock()
 	}
 
-	iw := indent.NewWriterPipe(w, indentation+margin, func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, " ")
-	})
-
+	mw := NewMarginWriter(ctx, w, bs.Current().Style)
+	renderText(mw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockPrefix)
 	if len(theme) > 0 {
-		renderText(iw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockPrefix)
-
-		err := quick.Highlight(iw, e.Code, e.Language, "terminal256", theme)
+		err := quick.Highlight(mw, e.Code, e.Language, "terminal256", theme)
 		if err != nil {
 			return err
 		}
-		renderText(iw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockSuffix)
-		return nil
-	}
+	} else {
+		// fallback rendering
+		el := &BaseElement{
+			Token: e.Code,
+			Style: rules.StylePrimitive,
+		}
 
-	// fallback rendering
-	el := &BaseElement{
-		Token: e.Code,
-		Style: rules.StylePrimitive,
+		err := el.Render(mw, ctx)
+		if err != nil {
+			return err
+		}
 	}
+	renderText(mw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockSuffix)
 
-	return el.Render(iw, ctx)
+	bs.Current().Block.Reset()
+	bs.Pop()
+	return nil
 }

ansi/elements.go

@@ -282,12 +282,14 @@ func (tr *ANSIRenderer) NewElement(node ast.Node, source []byte) Element {
 			line := n.Lines().At(i)
 			s += string(line.Value(source))
 		}
+		e := &CodeBlockElement{
+			Code:     s,
+			Language: string(n.Language(source)),
+		}
 		return Element{
 			Entering: "\n",
-			Renderer: &CodeBlockElement{
-				Code:     s,
-				Language: string(n.Language(source)),
-			},
+			Renderer: e,
+			Finisher: e,
 		}
 
 	case ast.KindCodeBlock:
@@ -298,11 +300,13 @@ func (tr *ANSIRenderer) NewElement(node ast.Node, source []byte) Element {
 			line := n.Lines().At(i)
 			s += string(line.Value(source))
 		}
+		e := &CodeBlockElement{
+			Code: s,
+		}
 		return Element{
 			Entering: "\n",
-			Renderer: &CodeBlockElement{
-				Code: s,
-			},
+			Renderer: e,
+			Finisher: e,
 		}
 
 	case ast.KindCodeSpan:

ansi/heading.go

@@ -4,7 +4,6 @@
 	"bytes"
 	"io"
 
-	"github.com/muesli/reflow/indent"
 	"github.com/muesli/reflow/wordwrap"
 )
 
@@ -60,28 +59,16 @@
 func (e *HeadingElement) Finish(w io.Writer, ctx RenderContext) error {
 	bs := ctx.blockStack
 	rules := bs.Current().Style
+	mw := NewMarginWriter(ctx, w, rules)
 
-	var indentation uint
-	var margin uint
-	if rules.Indent != nil {
-		indentation = *rules.Indent
-	}
-	if rules.Margin != nil {
-		margin = *rules.Margin
-	}
-
-	iw := indent.NewWriterPipe(w, indentation+margin, func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, bs.Parent().Style.StylePrimitive, " ")
-	})
-
-	flow := wordwrap.NewWriter(int(bs.Width(ctx) - indentation - margin*2))
+	flow := wordwrap.NewWriter(int(bs.Width(ctx)))
 	_, err := flow.Write(bs.Current().Block.Bytes())
 	if err != nil {
 		return err
 	}
 	flow.Close()
 
-	_, err = iw.Write(flow.Bytes())
+	_, err = mw.Write(flow.Bytes())
 	if err != nil {
 		return err
 	}

ansi/margin.go

@@ -5,11 +5,19 @@ import (
 
 	"github.com/muesli/reflow/indent"
 	"github.com/muesli/reflow/padding"
+	"github.com/muesli/termenv"
 )
 
 // MarginWriter is a Writer that applies indentation and padding around
 // whatever you write to it.
 type MarginWriter struct {
+	indentation, margin  uint
+	indentPos, marginPos uint
+	indentToken          string
+
+	profile            termenv.Profile
+	rules, parentRules StylePrimitive
+
 	w  io.Writer
 	pw *padding.Writer
 	iw *indent.Writer
@@ -18,35 +26,55 @@ type MarginWriter struct {
 // NewMarginWriter returns a new MarginWriter.
 func NewMarginWriter(ctx RenderContext, w io.Writer, rules StyleBlock) *MarginWriter {
 	bs := ctx.blockStack
+	mw := &MarginWriter{
+		w:           w,
+		profile:     ctx.options.ColorProfile,
+		rules:       rules.StylePrimitive,
+		parentRules: bs.Parent().Style.StylePrimitive,
+	}
 
-	var indentation uint
-	var margin uint
 	if rules.Indent != nil {
-		indentation = *rules.Indent
+		mw.indentation = *rules.Indent
+		mw.indentToken = " "
+		if rules.IndentToken != nil {
+			mw.indentToken = *rules.IndentToken
+		}
 	}
 	if rules.Margin != nil {
-		margin = *rules.Margin
+		mw.margin = *rules.Margin
 	}
 
-	pw := padding.NewWriterPipe(w, bs.Width(ctx), func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, rules.StylePrimitive, " ")
-	})
-
-	ic := " "
-	if rules.IndentToken != nil {
-		ic = *rules.IndentToken
-	}
-	iw := indent.NewWriterPipe(pw, indentation+margin, func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, bs.Parent().Style.StylePrimitive, ic)
+	mw.pw = padding.NewWriterPipe(mw.w, bs.Width(ctx), func(wr io.Writer) {
+		renderText(mw.w, mw.profile, mw.rules, " ")
 	})
 
-	return &MarginWriter{
-		w:  w,
-		pw: pw,
-		iw: iw,
-	}
+	mw.iw = indent.NewWriterPipe(mw.pw, mw.indentation+mw.margin, mw.indentFunc)
+	return mw
 }
 
 func (w *MarginWriter) Write(b []byte) (int, error) {
 	return w.iw.Write(b)
 }
+
+// indentFunc is called when writing each the margin and indentation tokens.
+// The margin is written first, using an empty space character as the token.
+// The indentation is written next, using the token specified in the rules.
+func (w *MarginWriter) indentFunc(iw io.Writer) {
+	ic := " "
+	switch {
+	case w.margin == 0 && w.indentation == 0:
+		return
+	case w.margin >= 1 && w.indentation == 0:
+		break
+	case w.margin >= 1 && w.marginPos < w.margin:
+		w.marginPos++
+	case w.indentation >= 1 && w.indentPos < w.indentation:
+		w.indentPos++
+		ic = w.indentToken
+		if w.indentPos == w.indentation {
+			w.marginPos = 0
+			w.indentPos = 0
+		}
+	}
+	renderText(w.w, w.profile, w.parentRules, ic)
+}