From 24627da0d2e321d93bdeca4fd2124a84b055beb7 Mon Sep 17 00:00:00 2001
From: Gabriel Garrido <gabriel@garrido.io>
Date: Sun, 19 May 2024 21:29:36 +0200
Subject: [PATCH] Support visibility filtering using flag

---
 README.md        |  3 ++-
 client/client.go | 17 ++++++++++++-----
 client/post.go   |  8 ++++++--
 main.go          | 10 +++++-----
 4 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index a473140..a6e903f 100644
--- a/README.md
+++ b/README.md
@@ -45,7 +45,8 @@ Usage of mastodon-markdown-archive:
         Thread replies for a post in a single file
   -user string
         URL of Mastodon account whose toots will be fetched
-
+  -visibility string
+        Filter out posts whose visibility does not match the passed visibility value
 ```
 
 ## Example
diff --git a/client/client.go b/client/client.go
index 94bc4e0..9b1fd9b 100644
--- a/client/client.go
+++ b/client/client.go
@@ -6,6 +6,11 @@ import (
 	"strings"
 )
 
+type ClientOptions struct {
+	Visibility string
+	Threaded   bool
+}
+
 type Client struct {
 	handle  string
 	baseURL string
@@ -19,7 +24,8 @@ type Client struct {
 	// Map of Post.Id:*Post.
 	postIdMap map[string]*Post
 	// List of Post.Id. Tracks the posts which will be written as individual files.
-	output []string
+	output  []string
+	options ClientOptions
 }
 
 type PostsFilter struct {
@@ -31,7 +37,7 @@ type PostsFilter struct {
 	MaxId          string
 }
 
-func New(userURL string, filters PostsFilter, threaded bool) (Client, error) {
+func New(userURL string, filters PostsFilter, opts ClientOptions) (Client, error) {
 	var client Client
 	parsedURL, err := url.Parse(userURL)
 
@@ -63,7 +69,7 @@ func New(userURL string, filters PostsFilter, threaded bool) (Client, error) {
 	for i := range posts {
 		post := posts[i]
 		postIdMap[post.Id] = &post
-		if !threaded {
+		if !opts.Threaded && !post.ShouldSkip(opts.Visibility) {
 			output = append(output, post.Id)
 		}
 	}
@@ -77,9 +83,10 @@ func New(userURL string, filters PostsFilter, threaded bool) (Client, error) {
 		replies:   replies,
 		orphans:   orphans,
 		output:    output,
+		options:   opts,
 	}
 
-	if threaded {
+	if opts.Threaded {
 		for _, post := range posts {
 			client.threadPost(post.Id)
 		}
@@ -154,7 +161,7 @@ func (c *Client) flushReplies(post *Post, descendants *[]*Post) {
 func (c *Client) threadPost(postId string) {
 	post := c.postIdMap[postId]
 
-	if post.InReplyToId == "" {
+	if post.InReplyToId == "" && !post.ShouldSkip(c.options.Visibility) {
 		c.flushReplies(post, &post.descendants)
 		c.output = append(c.output, post.Id)
 		return
diff --git a/client/post.go b/client/post.go
index 998efab..c3fac0a 100644
--- a/client/post.go
+++ b/client/post.go
@@ -57,8 +57,12 @@ type Post struct {
 	descendants        []*Post
 }
 
-func (p Post) ShouldSkip() bool {
-	return p.Visibility != "public"
+func (p Post) ShouldSkip(visibility string) bool {
+	if visibility == "" {
+		return false
+	}
+
+	return p.Visibility != visibility
 }
 
 func (p Post) Descendants() []*Post {
diff --git a/main.go b/main.go
index 533f1c2..01b6a89 100644
--- a/main.go
+++ b/main.go
@@ -27,6 +27,7 @@ func main() {
 	filenameTemplate := flag.String("filename", "", "Template for post filename")
 	porcelain := flag.Bool("porcelain", false, "Prints the amount of fetched posts to stdout in a parsable manner")
 	downloadMedia := flag.String("download-media", "", "Download media in a post. Omit or pass an empty string to not download media. Pass 'bundle' to download the media inline in a single directory with its original post. Pass a path to a directory to download all media there.")
+	visibility := flag.String("visibility", "", "Filter out posts whose visibility does not match the passed visibility value")
 
 	flag.Parse()
 
@@ -37,7 +38,10 @@ func main() {
 		SinceId:        *sinceId,
 		MaxId:          *maxId,
 		MinId:          *minId,
-	}, *threaded)
+	}, client.ClientOptions{
+		Threaded:   *threaded,
+		Visibility: *visibility,
+	})
 
 	if err != nil {
 		log.Panicln(err)
@@ -58,10 +62,6 @@ func main() {
 	}
 
 	for _, post := range posts {
-		if post.ShouldSkip() {
-			continue
-		}
-
 		if err := fileWriter.Write(post); err != nil {
 			log.Panicln("error writing post to file: %w", err)
 		}