From 758fb26dad0692dd7fd165d2b0e74e6c7e2b18b8 Mon Sep 17 00:00:00 2001
From: Tyrrrz <1935960+Tyrrrz@users.noreply.github.com>
Date: Mon, 21 Aug 2023 17:44:54 +0300
Subject: [PATCH] Add `--include-threads` and `--include-archived-threads` to
 `exportall`

Related to #1119
---
 .../Commands/ExportAllCommand.cs              | 40 +++++++++++++++++++
 .../Commands/ExportGuildCommand.cs            | 36 ++++++++++-------
 .../Discord/Data/ChannelKind.cs               |  3 ++
 3 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/DiscordChatExporter.Cli/Commands/ExportAllCommand.cs b/DiscordChatExporter.Cli/Commands/ExportAllCommand.cs
index e32353e..18191c5 100644
--- a/DiscordChatExporter.Cli/Commands/ExportAllCommand.cs
+++ b/DiscordChatExporter.Cli/Commands/ExportAllCommand.cs
@@ -34,6 +34,18 @@ public class ExportAllCommand : ExportCommandBase
     )]
     public bool IncludeVoiceChannels { get; init; } = true;
 
+    [CommandOption(
+        "include-threads",
+        Description = "Include threads."
+    )]
+    public bool IncludeThreads { get; init; } = false;
+
+    [CommandOption(
+        "include-archived-threads",
+        Description = "Include archived threads."
+    )]
+    public bool IncludeArchivedThreads { get; init; } = false;
+
     [CommandOption(
         "data-package",
         Description =
@@ -46,6 +58,14 @@ public class ExportAllCommand : ExportCommandBase
     {
         await base.ExecuteAsync(console);
 
+        // Cannot include archived threads without including active threads as well
+        if (IncludeArchivedThreads && !IncludeThreads)
+        {
+            throw new CommandException(
+                "Option --include-archived-threads can only be used when --include-threads is also specified."
+            );
+        }
+
         var cancellationToken = console.RegisterCancellationHandler();
         var channels = new List<Channel>();
 
@@ -56,10 +76,26 @@ public class ExportAllCommand : ExportCommandBase
 
             await foreach (var guild in Discord.GetUserGuildsAsync(cancellationToken))
             {
+                // Regular channels
                 await foreach (var channel in Discord.GetGuildChannelsAsync(guild.Id, cancellationToken))
                 {
+                    if (channel.Kind == ChannelKind.GuildCategory)
+                        continue;
+
+                    if (!IncludeVoiceChannels && channel.Kind.IsVoice())
+                        continue;
+
                     channels.Add(channel);
                 }
+
+                // Threads
+                if (IncludeThreads)
+                {
+                    await foreach (var thread in Discord.GetGuildThreadsAsync(guild.Id, IncludeArchivedThreads, cancellationToken))
+                    {
+                        channels.Add(thread);
+                    }
+                }
             }
         }
         // Pull from the data package
@@ -105,6 +141,10 @@ public class ExportAllCommand : ExportCommandBase
             channels.RemoveAll(c => c.Kind.IsGuild());
         if (!IncludeVoiceChannels)
             channels.RemoveAll(c => c.Kind.IsVoice());
+        if (!IncludeThreads)
+            channels.RemoveAll(c => c.Kind.IsThread());
+        if (!IncludeArchivedThreads)
+            channels.RemoveAll(c => c.Kind.IsThread() && c.IsArchived);
 
         await ExportAsync(console, channels);
     }
diff --git a/DiscordChatExporter.Cli/Commands/ExportGuildCommand.cs b/DiscordChatExporter.Cli/Commands/ExportGuildCommand.cs
index a85de3f..7279dac 100644
--- a/DiscordChatExporter.Cli/Commands/ExportGuildCommand.cs
+++ b/DiscordChatExporter.Cli/Commands/ExportGuildCommand.cs
@@ -1,5 +1,4 @@
-using System;
-using System.Linq;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using CliFx.Attributes;
 using CliFx.Exceptions;
@@ -7,7 +6,6 @@ using CliFx.Infrastructure;
 using DiscordChatExporter.Cli.Commands.Base;
 using DiscordChatExporter.Core.Discord;
 using DiscordChatExporter.Core.Discord.Data;
-using DiscordChatExporter.Core.Utils.Extensions;
 
 namespace DiscordChatExporter.Cli.Commands;
 
@@ -52,21 +50,31 @@ public class ExportGuildCommand : ExportCommandBase
         }
 
         var cancellationToken = console.RegisterCancellationHandler();
+        var channels = new List<Channel>();
 
         await console.Output.WriteLineAsync("Fetching channels...");
 
-        var channels = (await Discord.GetGuildChannelsAsync(GuildId, cancellationToken))
-            .Where(c => c.Kind != ChannelKind.GuildCategory)
-            .Where(c => IncludeVoiceChannels || !c.Kind.IsVoice())
-            .ToArray();
+        // Regular channels
+        await foreach (var channel in Discord.GetGuildChannelsAsync(GuildId, cancellationToken))
+        {
+            if (channel.Kind == ChannelKind.GuildCategory)
+                continue;
+
+            if (!IncludeVoiceChannels && channel.Kind.IsVoice())
+                continue;
+
+            channels.Add(channel);
+        }
 
-        var threads = IncludeThreads
-            ? await Discord.GetGuildThreadsAsync(GuildId, IncludeArchivedThreads, cancellationToken)
-            : Array.Empty<Channel>();
+        // Threads
+        if (IncludeThreads)
+        {
+            await foreach (var thread in Discord.GetGuildThreadsAsync(GuildId, IncludeArchivedThreads, cancellationToken))
+            {
+                channels.Add(thread);
+            }
+        }
 
-        await ExportAsync(
-            console,
-            channels.Concat(threads).ToArray()
-        );
+        await ExportAsync(console, channels);
     }
 }
\ No newline at end of file
diff --git a/DiscordChatExporter.Core/Discord/Data/ChannelKind.cs b/DiscordChatExporter.Core/Discord/Data/ChannelKind.cs
index 9bc4a0e..6794695 100644
--- a/DiscordChatExporter.Core/Discord/Data/ChannelKind.cs
+++ b/DiscordChatExporter.Core/Discord/Data/ChannelKind.cs
@@ -27,4 +27,7 @@ public static class ChannelKindExtensions
 
     public static bool IsVoice(this ChannelKind kind) =>
         kind is ChannelKind.GuildVoiceChat or ChannelKind.GuildStageVoice;
+
+    public static bool IsThread(this ChannelKind kind) =>
+        kind is ChannelKind.GuildNewsThread or ChannelKind.GuildPublicThread or ChannelKind.GuildPrivateThread;
 }
\ No newline at end of file