From f4a2e4004c904b84028816f7fed862fecfa9727d Mon Sep 17 00:00:00 2001
From: Florian <me@fgardt.dev>
Date: Tue, 14 May 2024 02:46:21 +0200
Subject: [PATCH] Add User Apps

---
 macros/src/command/mod.rs | 40 +++++++++++++++++++++++++++++++++++++++
 src/structs/command.rs    | 28 +++++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/macros/src/command/mod.rs b/macros/src/command/mod.rs
index 8c7f8bbd6136..80682fec57f4 100644
--- a/macros/src/command/mod.rs
+++ b/macros/src/command/mod.rs
@@ -49,6 +49,9 @@ pub struct CommandArgs {
     category: Option<String>,
     custom_data: Option<syn::Expr>,
 
+    install_context: Option<syn::punctuated::Punctuated<syn::Ident, syn::Token![|]>>,
+    interaction_context: Option<syn::punctuated::Punctuated<syn::Ident, syn::Token![|]>>,
+
     // In seconds
     global_cooldown: Option<u64>,
     user_cooldown: Option<u64>,
@@ -97,6 +100,8 @@ pub struct Invocation {
     default_member_permissions: syn::Expr,
     required_permissions: syn::Expr,
     required_bot_permissions: syn::Expr,
+    install_context: syn::Expr,
+    interaction_context: syn::Expr,
     args: CommandArgs,
 }
 
@@ -216,6 +221,34 @@ pub fn command(
     let required_permissions = permissions_to_tokens(&args.required_permissions);
     let required_bot_permissions = permissions_to_tokens(&args.required_bot_permissions);
 
+    fn build_install_context(
+        contexts: &Option<syn::punctuated::Punctuated<syn::Ident, syn::Token![|]>>,
+    ) -> syn::Expr {
+        match contexts {
+            Some(contexts) => {
+                let contexts = contexts.iter();
+                syn::parse_quote! { Some(vec![ #(poise::serenity_prelude::InstallationContext::#contexts),* ]) }
+            }
+            None => syn::parse_quote! { None },
+        }
+    }
+
+    let install_context = build_install_context(&args.install_context);
+
+    fn build_interaction_context(
+        contexts: &Option<syn::punctuated::Punctuated<syn::Ident, syn::Token![|]>>,
+    ) -> syn::Expr {
+        match contexts {
+            Some(contexts) => {
+                let contexts = contexts.iter();
+                syn::parse_quote! { Some(vec![ #(poise::serenity_prelude::InteractionContext::#contexts),* ]) }
+            }
+            None => syn::parse_quote! { None },
+        }
+    }
+
+    let interaction_context = build_interaction_context(&args.interaction_context);
+
     let inv = Invocation {
         parameters,
         description,
@@ -225,6 +258,8 @@ pub fn command(
         default_member_permissions,
         required_permissions,
         required_bot_permissions,
+        install_context,
+        interaction_context,
     };
 
     Ok(TokenStream::from(generate_command(inv)?))
@@ -291,6 +326,9 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
     let dm_only = inv.args.dm_only;
     let nsfw_only = inv.args.nsfw_only;
 
+    let install_context = &inv.install_context;
+    let interaction_context = &inv.interaction_context;
+
     let help_text = match &inv.args.help_text_fn {
         Some(help_text_fn) => quote::quote! { Some(#help_text_fn()) },
         None => match &inv.help_text {
@@ -364,6 +402,8 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
                 guild_only: #guild_only,
                 dm_only: #dm_only,
                 nsfw_only: #nsfw_only,
+                install_context: #install_context,
+                interaction_context: #interaction_context,
                 checks: vec![ #( |ctx| Box::pin(#checks(ctx)) ),* ],
                 on_error: #on_error,
                 parameters: vec![ #( #parameters ),* ],
diff --git a/src/structs/command.rs b/src/structs/command.rs
index 453a4840151f..647e175e6ee1 100644
--- a/src/structs/command.rs
+++ b/src/structs/command.rs
@@ -122,6 +122,10 @@ pub struct Command<U, E> {
     pub context_menu_name: Option<String>,
     /// Whether responses to this command should be ephemeral by default (application-only)
     pub ephemeral: bool,
+    /// List of installation contexts for this command (application-only)
+    pub install_context: Option<Vec<serenity::InstallationContext>>,
+    /// List of interaction contexts for this command (application-only)
+    pub interaction_context: Option<Vec<serenity::InteractionContext>>,
 
     // Like #[non_exhaustive], but #[poise::command] still needs to be able to create an instance
     #[doc(hidden)]
@@ -196,7 +200,17 @@ impl<U, E> Command<U, E> {
         }
 
         if self.guild_only {
-            builder = builder.dm_permission(false);
+            builder = builder.contexts(vec![serenity::InteractionContext::Guild]);
+        } else if self.dm_only {
+            builder = builder.contexts(vec![serenity::InteractionContext::BotDm]);
+        }
+
+        if let Some(install_context) = self.install_context.clone() {
+            builder = builder.integration_types(install_context);
+        }
+
+        if let Some(interaction_context) = self.interaction_context.clone() {
+            builder = builder.contexts(interaction_context);
         }
 
         if self.subcommands.is_empty() {
@@ -230,7 +244,17 @@ impl<U, E> Command<U, E> {
         });
 
         if self.guild_only {
-            builder = builder.dm_permission(false);
+            builder = builder.contexts(vec![serenity::InteractionContext::Guild]);
+        } else if self.dm_only {
+            builder = builder.contexts(vec![serenity::InteractionContext::BotDm]);
+        }
+
+        if let Some(install_context) = self.install_context.clone() {
+            builder = builder.integration_types(install_context);
+        }
+
+        if let Some(interaction_context) = self.interaction_context.clone() {
+            builder = builder.contexts(interaction_context);
         }
 
         Some(builder)