From cce8cc64ff109963246f51c34218afee0b3bafd9 Mon Sep 17 00:00:00 2001
From: GnomedDev <david2005thomas@gmail.com>
Date: Fri, 1 Nov 2024 02:01:09 +0000
Subject: [PATCH] Add poise::StrArg for FromStr arguments

---
 macros/src/lib.rs                     |  2 +-
 src/argument.rs                       | 61 +++++++++++++++++++++++++++
 src/lib.rs                            | 11 ++---
 src/prefix_argument/argument_trait.rs |  2 +-
 4 files changed, 69 insertions(+), 7 deletions(-)
 create mode 100644 src/argument.rs

diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index 76158deb785..8e37f5d506d 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -88,7 +88,7 @@ for example for command-specific help (i.e. `~help command_name`). Escape newlin
 SlashContext, which contain a variety of context data each. Context provides some utility methods to
 access data present in both PrefixContext and SlashContext, like `author()` or `created_at()`.
 
-All following parameters are inputs to the command. You can use all types that implement `poise::PopArgument`, `serenity::ArgumentConvert` or `std::str::FromStr`.
+All following parameters are inputs to the command. You can use all types that implement `PopArgument` (for prefix) or `SlashArgument` (for slash).
 You can also wrap types in `Option` or `Vec` to make them optional or variadic. In addition, there
 are multiple attributes you can use on parameters:
 
diff --git a/src/argument.rs b/src/argument.rs
new file mode 100644
index 00000000000..443286bacb1
--- /dev/null
+++ b/src/argument.rs
@@ -0,0 +1,61 @@
+use std::str::FromStr;
+
+use crate::{
+    serenity_prelude as serenity, PopArgument, PopArgumentResult, SlashArgError, SlashArgument,
+};
+
+/// A wrapper for `T` to implement [`SlashArgument`] and [`PopArgument`] via [`FromStr`].
+///
+/// This is useful if you need to take an argument via a string, but immediately convert it via [`FromStr`].
+pub struct StrArg<T>(pub T);
+
+#[async_trait::async_trait]
+impl<T> SlashArgument for StrArg<T>
+where
+    T: FromStr,
+    T::Err: std::error::Error + Send + Sync + 'static,
+{
+    fn create(builder: serenity::CreateCommandOption) -> serenity::CreateCommandOption {
+        builder.kind(serenity::CommandOptionType::String)
+    }
+
+    async fn extract(
+        _: &serenity::Context,
+        _: &serenity::CommandInteraction,
+        value: &serenity::ResolvedValue<'_>,
+    ) -> Result<Self, SlashArgError> {
+        let serenity::ResolvedValue::String(value) = value else {
+            return Err(SlashArgError::new_command_structure_mismatch(
+                "expected a String",
+            ));
+        };
+
+        match T::from_str(value) {
+            Ok(value) => Ok(Self(value)),
+            Err(err) => Err(SlashArgError::Parse {
+                error: err.into(),
+                input: String::from(*value),
+            }),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<'a, T> PopArgument<'a> for StrArg<T>
+where
+    T: FromStr,
+    T::Err: std::error::Error + Send + Sync + 'static,
+{
+    async fn pop_from(
+        args: &'a str,
+        attachment_index: usize,
+        ctx: &serenity::Context,
+        msg: &serenity::Message,
+    ) -> PopArgumentResult<'a, Self> {
+        let (args, attach_idx, value) = String::pop_from(args, attachment_index, ctx, msg).await?;
+        match T::from_str(&value) {
+            Ok(value) => Ok((args, attach_idx, Self(value))),
+            Err(err) => Err((Box::new(err), Some(value))),
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 032a9ef40ca..c7d4de4b595 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -173,7 +173,7 @@ For another example of subcommands, see `examples/feature_showcase/subcommands.r
 Also see the [`command`] macro docs
 
 ```rust
-use poise::serenity_prelude as serenity;
+use poise::{StrArg, serenity_prelude as serenity};
 type Data = ();
 type Error = Box<dyn std::error::Error + Send + Sync>;
 type Context<'a> = poise::Context<'a, Data, Error>;
@@ -190,9 +190,9 @@ type Context<'a> = poise::Context<'a, Data, Error>;
 )]
 async fn my_huge_ass_command(
     ctx: Context<'_>,
-    #[description = "Consectetur"] ip_addr: std::net::IpAddr, // implements FromStr
-    #[description = "Amet"] user: serenity::Member, // implements ArgumentConvert
-    #[description = "Sit"] code_block: poise::CodeBlock, // implements PopArgument
+    #[description = "Amet"] user: serenity::Member,
+    #[description = "Consectetur"] #[rename = "ip_addr"] StrArg(ip_addr): StrArg<std::net::IpAddr>,
+    #[description = "Sit"] code_block: poise::CodeBlock,
     #[description = "Dolor"] #[flag] my_flag: bool,
     #[description = "Ipsum"] #[lazy] always_none: Option<String>,
     #[description = "Lorem"] #[rest] rest: String,
@@ -381,6 +381,7 @@ underlying this framework, so that's what I chose.
 Also, poise is a stat in Dark Souls
 */
 
+mod argument;
 pub mod builtins;
 pub mod choice_parameter;
 pub mod cooldown;
@@ -400,7 +401,7 @@ pub mod macros {
 
 #[doc(no_inline)]
 pub use {
-    choice_parameter::*, cooldown::*, dispatch::*, framework::*, macros::*, modal::*,
+    argument::*, choice_parameter::*, cooldown::*, dispatch::*, framework::*, macros::*, modal::*,
     prefix_argument::*, reply::*, slash_argument::*, structs::*, track_edits::*,
 };
 
diff --git a/src/prefix_argument/argument_trait.rs b/src/prefix_argument/argument_trait.rs
index 77633842a42..0ca690c2166 100644
--- a/src/prefix_argument/argument_trait.rs
+++ b/src/prefix_argument/argument_trait.rs
@@ -127,7 +127,7 @@ impl_popargument_via_argumentconvert!(
     f32, f64,
     u8, u16, u32, u64,
     i8, i16, i32, i64,
-    serenity::UserId, serenity::User,
+    serenity::UserId, serenity::User, serenity::Member,
     serenity::MessageId, serenity::Message,
     serenity::ChannelId, serenity::Channel, serenity::GuildChannel,
     serenity::EmojiId, serenity::Emoji,