renderer: add support for extending node renderers

renderer.go

@@ -4,6 +4,7 @@ package markdown
 import (
 	"bytes"
 	"io"
+	"sync"
 
 	"github.com/yuin/goldmark/ast"
 	"github.com/yuin/goldmark/renderer"
@@ -13,8 +14,10 @@ import (
 // NewRenderer returns a new markdown Renderer that is configured by default values.
 func NewRenderer(options ...Option) *Renderer {
 	r := &Renderer{
-		config: NewConfig(),
-		rc:     renderContext{},
+		config:               NewConfig(),
+		rc:                   renderContext{},
+		maxKind:              20, // a random number slightly larger than the number of default ast kinds
+		nodeRendererFuncsTmp: map[ast.NodeKind]renderer.NodeRendererFunc{},
 	}
 	for _, opt := range options {
 		opt.SetMarkdownOption(r.config)
@@ -24,10 +27,16 @@ func NewRenderer(options ...Option) *Renderer {
 
 // Renderer is an implementation of renderer.Renderer that renders nodes as Markdown
 type Renderer struct {
-	config *Config
-	rc     renderContext
+	config               *Config
+	rc                   renderContext
+	nodeRendererFuncsTmp map[ast.NodeKind]renderer.NodeRendererFunc
+	maxKind              int
+	nodeRendererFuncs    []nodeRenderer
+	initSync             sync.Once
 }
 
+var _ renderer.Renderer = &Renderer{}
+
 // AddOptions implements renderer.Renderer.AddOptions
 func (r *Renderer) AddOptions(opts ...renderer.Option) {
 	config := renderer.NewConfig()
@@ -37,64 +46,74 @@ func (r *Renderer) AddOptions(opts ...renderer.Option) {
 	for name, value := range config.Options {
 		r.config.SetOption(name, value)
 	}
-	// TODO handle any config.NodeRenderers set by opts
+
+	// handle any config.NodeRenderers set by opts
+	config.NodeRenderers.Sort()
+	l := len(config.NodeRenderers)
+	for i := l - 1; i >= 0; i-- {
+		v := config.NodeRenderers[i]
+		nr, _ := v.Value.(renderer.NodeRenderer)
+		nr.RegisterFuncs(r)
+	}
+}
+
+func (r *Renderer) Register(kind ast.NodeKind, fun renderer.NodeRendererFunc) {
+	r.nodeRendererFuncsTmp[kind] = fun
+	if int(kind) > r.maxKind {
+		r.maxKind = int(kind)
+	}
 }
 
 // Render implements renderer.Renderer.Render
 func (r *Renderer) Render(w io.Writer, source []byte, n ast.Node) error {
 	r.rc = newRenderContext(w, source, r.config)
-	/* TODO
-	reg.Register(ast.KindString, r.renderString)
-	*/
+	r.initSync.Do(func() {
+		r.nodeRendererFuncs = make([]nodeRenderer, r.maxKind+1)
+		// add default functions
+		// blocks
+		r.nodeRendererFuncs[ast.KindDocument] = r.renderBlockSeparator
+		r.nodeRendererFuncs[ast.KindHeading] = r.chainRenderers(r.renderBlockSeparator, r.renderHeading)
+		r.nodeRendererFuncs[ast.KindBlockquote] = r.chainRenderers(r.renderBlockSeparator, r.renderBlockquote)
+		r.nodeRendererFuncs[ast.KindCodeBlock] = r.chainRenderers(r.renderBlockSeparator, r.renderCodeBlock)
+		r.nodeRendererFuncs[ast.KindFencedCodeBlock] = r.chainRenderers(r.renderBlockSeparator, r.renderFencedCodeBlock)
+		r.nodeRendererFuncs[ast.KindHTMLBlock] = r.chainRenderers(r.renderBlockSeparator, r.renderHTMLBlock)
+		r.nodeRendererFuncs[ast.KindList] = r.chainRenderers(r.renderBlockSeparator, r.renderList)
+		r.nodeRendererFuncs[ast.KindListItem] = r.chainRenderers(r.renderBlockSeparator, r.renderListItem)
+		r.nodeRendererFuncs[ast.KindParagraph] = r.renderBlockSeparator
+		r.nodeRendererFuncs[ast.KindTextBlock] = r.renderBlockSeparator
+		r.nodeRendererFuncs[ast.KindThematicBreak] = r.chainRenderers(r.renderBlockSeparator, r.renderThematicBreak)
+
+		// inlines
+		r.nodeRendererFuncs[ast.KindAutoLink] = r.renderAutoLink
+		r.nodeRendererFuncs[ast.KindCodeSpan] = r.renderCodeSpan
+		r.nodeRendererFuncs[ast.KindEmphasis] = r.renderEmphasis
+		r.nodeRendererFuncs[ast.KindImage] = r.renderImage
+		r.nodeRendererFuncs[ast.KindLink] = r.renderLink
+		r.nodeRendererFuncs[ast.KindRawHTML] = r.renderRawHTML
+		r.nodeRendererFuncs[ast.KindText] = r.renderText
+		// TODO: add KindString
+		// r.nodeRendererFuncs[ast.KindString] = r.renderString
+
+		for kind, fun := range r.nodeRendererFuncsTmp {
+			r.nodeRendererFuncs[kind] = r.transform(fun)
+		}
+		r.nodeRendererFuncsTmp = nil
+	})
 	return ast.Walk(n, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
-		return r.getRenderer(n)(n, entering), r.rc.writer.Err()
+		return r.nodeRendererFuncs[n.Kind()](n, entering), r.rc.writer.Err()
 	})
 }
 
-// nodeRenderer is a markdown node renderer func.
-type nodeRenderer func(ast.Node, bool) ast.WalkStatus
-
-func (r *Renderer) getRenderer(node ast.Node) nodeRenderer {
-	renderers := []nodeRenderer{}
-	switch node.Type() {
-	case ast.TypeBlock:
-		renderers = append(renderers, r.renderBlockSeparator)
-	}
-	switch node.Kind() {
-	case ast.KindAutoLink:
-		renderers = append(renderers, r.renderAutoLink)
-	case ast.KindHeading:
-		renderers = append(renderers, r.renderHeading)
-	case ast.KindBlockquote:
-		renderers = append(renderers, r.renderBlockquote)
-	case ast.KindCodeBlock:
-		renderers = append(renderers, r.renderCodeBlock)
-	case ast.KindCodeSpan:
-		renderers = append(renderers, r.renderCodeSpan)
-	case ast.KindEmphasis:
-		renderers = append(renderers, r.renderEmphasis)
-	case ast.KindThematicBreak:
-		renderers = append(renderers, r.renderThematicBreak)
-	case ast.KindFencedCodeBlock:
-		renderers = append(renderers, r.renderFencedCodeBlock)
-	case ast.KindHTMLBlock:
-		renderers = append(renderers, r.renderHTMLBlock)
-	case ast.KindImage:
-		renderers = append(renderers, r.renderImage)
-	case ast.KindList:
-		renderers = append(renderers, r.renderList)
-	case ast.KindListItem:
-		renderers = append(renderers, r.renderListItem)
-	case ast.KindRawHTML:
-		renderers = append(renderers, r.renderRawHTML)
-	case ast.KindText:
-		renderers = append(renderers, r.renderText)
-	case ast.KindLink:
-		renderers = append(renderers, r.renderLink)
+func (r *Renderer) transform(fn renderer.NodeRendererFunc) nodeRenderer {
+	return func(n ast.Node, entering bool) ast.WalkStatus {
+		status, _ := fn(r.rc.writer, r.rc.source, n, entering)
+		return status
 	}
-	return r.chainRenderers(renderers...)
 }
 
+// nodeRenderer is a markdown node renderer func.
+type nodeRenderer func(ast.Node, bool) ast.WalkStatus
+
 func (r *Renderer) chainRenderers(renderers ...nodeRenderer) nodeRenderer {
 	return func(node ast.Node, entering bool) ast.WalkStatus {
 		var walkStatus ast.WalkStatus

writer.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"io"
 	"unicode"
+
+	"github.com/yuin/goldmark/util"
 )
 
 // Line delimiter
@@ -33,6 +35,8 @@ type markdownWriter struct {
 	err error
 }
 
+var _ util.BufWriter = &markdownWriter{}
+
 // newMarkdownWriter returns a new markdownWriter
 func newMarkdownWriter(w io.Writer, config *Config) *markdownWriter {
 	result := &markdownWriter{
@@ -55,7 +59,7 @@ func (m *markdownWriter) Reset(w io.Writer) {
 
 // WriteLine writes the given bytes as a finished line, regardless of trailing newline.
 func (m *markdownWriter) WriteLine(line []byte) (n int) {
-	n = m.Write(line)
+	n, _ = m.Write(line)
 	m.FlushLine()
 
 	return n
@@ -98,9 +102,9 @@ func (p *markdownWriter) PopPrefix() {
 
 // Write writes the given data to an internal buffer, then writes any complete lines to the
 // underlying writer.
-func (m *markdownWriter) Write(data []byte) (n int) {
+func (m *markdownWriter) Write(data []byte) (n int, err error) {
 	if m.err != nil {
-		return 0
+		return 0, m.err
 	}
 	// Writing to a bytes.Buffer always returns a nil error
 	n, _ = m.buf.Write(data)
@@ -123,15 +127,45 @@ func (m *markdownWriter) Write(data []byte) (n int) {
 		_, err := m.output.Write(prefixedLine.Bytes())
 		if err != nil {
 			m.err = err
-			return 0
+			return 0, m.err
 		}
 		m.line += 1
 		prefixedLine.Reset()
 	}
-	return n
+	return n, nil
 }
 
 // Err returns the last write error, or nil.
 func (m *markdownWriter) Err() error {
 	return m.err
 }
+
+// returns how many bytes are unused in the buffer.
+func (m *markdownWriter) Available() int {
+	return m.buf.Cap() - m.buf.Len()
+}
+
+func (m *markdownWriter) Buffered() int {
+	return m.buf.Len()
+}
+
+func (m *markdownWriter) Flush() error {
+	_, err := m.output.Write(m.buf.Bytes())
+	if err != nil {
+		return err
+	}
+	m.buf.Reset()
+	return nil
+}
+
+func (m *markdownWriter) WriteByte(c byte) error {
+	return m.buf.WriteByte(c)
+}
+
+func (m *markdownWriter) WriteRune(r rune) (size int, err error) {
+	return m.buf.WriteRune(r)
+}
+
+func (m *markdownWriter) WriteString(s string) (n int, err error) {
+	return m.buf.WriteString(s)
+}

writer_test.go

@@ -160,15 +160,19 @@ func TestWriteError(t *testing.T) {
 	writer := newMarkdownWriter(&ew, NewConfig())
 	data := []byte("foo\n")
 
-	assert.Equal(len(data), writer.Write(data), "Writes should succeed before error")
+	var n int
+	n, _ = writer.Write(data)
+	assert.Equal(len(data), n, "Writes should succeed before error")
 	assert.Equal(len(data), writer.WriteLine(data), "Writes should succeed before error")
 	ew.err = err
-	assert.Equal(0, writer.Write(data), "Once error is set, writes become no-op")
+	n, _ = writer.Write(data)
+	assert.Equal(0, n, "Once error is set, writes become no-op")
 	assert.Equal(0, writer.WriteLine(data), "Once error is set, writes become no-op")
 	assert.Equal(err, writer.Err(), "Err() should match error returned by errorWriter")
 
 	ew.err = nil
 	writer.Reset(&ew)
-	assert.Equal(len(data), writer.Write(data), "Writes should succeed after Reset")
+	n, _ = writer.Write(data)
+	assert.Equal(len(data), n, "Writes should succeed after Reset")
 	assert.Equal(len(data), writer.WriteLine(data), "Writes should succeed after Reset")
 }