Support citation style links

Author: Necoro

html2text.go

@@ -2,6 +2,7 @@ package html2text
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"regexp"
 	"strings"
@@ -18,6 +19,7 @@ type Options struct {
 	PrettyTables        bool                 // Turns on pretty ASCII rendering for table elements.
 	PrettyTablesOptions *PrettyTablesOptions // Configures pretty ASCII rendering for table elements.
 	OmitLinks           bool                 // Turns on omitting links
+	CitationStyleLinks  bool                 // Uses citation style links like [1]
 }
 
 // PrettyTablesOptions overrides tablewriter behaviors
@@ -72,11 +74,16 @@ func FromHTMLNode(doc *html.Node, o ...Options) (string, error) {
 	ctx := textifyTraverseContext{
 		buf:     bytes.Buffer{},
 		options: options,
+		citationMap: map[string]int{},
 	}
 	if err := ctx.traverse(doc); err != nil {
 		return "", err
 	}
 
+	if ctx.options.CitationStyleLinks && ctx.citationCount > 0 {
+		ctx.emitCitations()
+	}
+
 	text := strings.TrimSpace(newlineRe.ReplaceAllString(
 		strings.Replace(ctx.buf.String(), "\n ", "\n", -1), "\n\n"),
 	)
@@ -124,6 +131,8 @@ type textifyTraverseContext struct {
 	blockquoteLevel int
 	lineLength      int
 	isPre           bool
+	citationCount   int
+	citationMap     map[string]int
 }
 
 // tableTraverseContext holds table ASCII-form related context.
@@ -255,7 +264,11 @@ func (ctx *textifyTraverseContext) handleElement(node *html.Node) error {
 			attrVal = ctx.normalizeHrefLink(attrVal)
 			// Don't print link href if it matches link element content or if the link is empty.
 			if !ctx.options.OmitLinks && attrVal != "" && linkText != attrVal {
-				hrefLink = "( " + attrVal + " )"
+				if ctx.options.CitationStyleLinks {
+					hrefLink = ctx.addCitation(attrVal)
+				} else {
+					hrefLink = "( " + attrVal + " )"
+				}
 			}
 		}
 
@@ -503,6 +516,40 @@ func (ctx *textifyTraverseContext) normalizeHrefLink(link string) string {
 	return link
 }
 
+func formatCitation(idx int) string {
+	return fmt.Sprintf("[%d] ", idx)
+}
+
+func (ctx *textifyTraverseContext) addCitation(url string) string {
+	idx, ok := ctx.citationMap[url]
+
+	if !ok {
+		ctx.citationCount += 1
+		idx = ctx.citationCount
+		ctx.citationMap[url] = idx
+	}
+
+	return formatCitation(idx)
+}
+
+func (ctx *textifyTraverseContext) emitCitations() {
+	// this method writes to the buffer directly instead of using `emit`, b/c we do not want to split long links
+	ctx.buf.WriteString("\n\n")
+
+	// citations are ordered by link --> bring them into the correct order first
+	links := make([]string, ctx.citationCount)
+
+	for k, v := range ctx.citationMap {
+		links[v-1] = k // arrays are 0-based, our citations are 1-based
+	}
+
+	for i, link := range links {
+		ctx.buf.WriteString(formatCitation(i + 1))
+		ctx.buf.WriteString(link)
+		ctx.buf.WriteByte('\n')
+	}
+}
+
 // renderEachChild visits each direct child of a node and collects the sequence of
 // textuual representaitons separated by a single newline.
 func (ctx *textifyTraverseContext) renderEachChild(node *html.Node) (string, error) {

html2text_test.go

@@ -521,6 +521,58 @@ func TestOmitLinks(t *testing.T) {
 	}
 }
 
+func TestCitationStyleLinks(t *testing.T) {
+	testCases := []struct {
+		input  string
+		output string
+	}{
+		{
+			`<a></a>`,
+			``,
+		},
+		{
+			`<a href=""></a>`,
+			``,
+		},
+		{
+			`<a href="http://example.com/"></a>`,
+			"[1] \n\n[1] http://example.com/",
+		},
+		{
+			`<a href="">Link</a>`,
+			"Link",
+		},
+		{
+			`<a href="http://example1.com/">Link1</a><a href="http://example2.com/">Link2</a>`,
+			"Link1 [1] Link2 [2] \n\n[1] http://example1.com/\n[2] http://example2.com/",
+		},
+		{
+			`<a href="http://example1.com/">Link1</a><a href="http://example1.com/">Link1 again</a>`,
+			"Link1 [1] Link1 again [1] \n\n[1] http://example1.com/",
+		},
+		{
+			`<a href="http://example.com/"><span class="a">Link</span></a>`,
+			"Link [1] \n\n[1] http://example.com/",
+		},
+		{
+			"<a href='http://example.com/'>\n\t<span class='a'>Link</span>\n\t</a>",
+			"Link [1] \n\n[1] http://example.com/",
+		},
+		{
+			`<a href="http://example.com/"><img src="http://example.ru/hello.jpg" alt="Example"></a>`,
+			"Example [1] \n\n[1] http://example.com/",
+		},
+	}
+
+	for _, testCase := range testCases {
+		if msg, err := wantString(testCase.input, testCase.output, Options{CitationStyleLinks: true}); err != nil {
+			t.Error(err)
+		} else if len(msg) > 0 {
+			t.Log(msg)
+		}
+	}
+}
+
 func TestImageAltTags(t *testing.T) {
 	testCases := []struct {
 		input  string