Improve content scraper

3 files changed, 22 insertions, 0 deletions

diff --git a/reader/rewrite/rewrite_functions.go b/reader/rewrite/rewrite_functions.go
index f1a5b83..6717deb 100644
--- a/reader/rewrite/rewrite_functions.go
+++ b/reader/rewrite/rewrite_functions.go
@@ -5,6 +5,7 @@
 package rewrite
 
 import (
+	"fmt"
 	"regexp"
 	"strings"
 
@@ -38,3 +39,10 @@ func addYoutubeVideo(entryURL, entryContent string) string {
 	}
 	return entryContent
 }
+
+func addPDFLink(entryURL, entryContent string) string {
+	if strings.HasSuffix(entryURL, ".pdf") {
+		return fmt.Sprintf(`<a href="%s">PDF</a><br>%s`, entryURL, entryContent)
+	}
+	return entryContent
+}
diff --git a/reader/rewrite/rewriter.go b/reader/rewrite/rewriter.go
index d4d5096..7862f4e 100644
--- a/reader/rewrite/rewriter.go
+++ b/reader/rewrite/rewriter.go
@@ -18,12 +18,16 @@ func Rewriter(entryURL, entryContent, customRewriteRules string) string {
 	}
 
 	rules := strings.Split(rulesList, ",")
+	rules = append(rules, "add_pdf_download_link")
+
 	for _, rule := range rules {
 		switch strings.TrimSpace(rule) {
 		case "add_image_title":
 			entryContent = addImageTitle(entryURL, entryContent)
 		case "add_youtube_video":
 			entryContent = addYoutubeVideo(entryURL, entryContent)
+		case "add_pdf_download_link":
+			entryContent = addPDFLink(entryURL, entryContent)
 		}
 	}
 
diff --git a/reader/rewrite/rewriter_test.go b/reader/rewrite/rewriter_test.go
index 5f13895..995508b 100644
--- a/reader/rewrite/rewriter_test.go
+++ b/reader/rewrite/rewriter_test.go
@@ -58,3 +58,13 @@ func TestRewriteWithXkcdAndNoImage(t *testing.T) {
 		t.Errorf(`Not expected output: got "%s" instead of "%s"`, output, expected)
 	}
 }
+
+func TestRewriteWithPDFLink(t *testing.T) {
+	description := "test"
+	output := Rewriter("https://example.org/document.pdf", description, ``)
+	expected := `<a href="https://example.org/document.pdf">PDF</a><br>test`
+
+	if expected != output {
+		t.Errorf(`Not expected output: got "%s" instead of "%s"`, output, expected)
+	}
+}