add custom parse function to AtatCmd derive

dragonnn · Jan 29, 2024 · 6a5bca5 · 6a5bca5
1 parent aeac149
commit 6a5bca5
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 11 deletions.
diff --git a/atat/src/derive.rs b/atat/src/derive.rs
@@ -302,4 +302,30 @@ mod tests {
             from_str::<MixedEnum<'_>>("6,\"abc\"")
         );
     }
+
+    fn custom_parse(response: &[u8]) -> Result<CustomResponseParse, atat::Error> {
+        Ok(CustomResponseParse {
+            arg1: core::str::from_utf8(&response[6..])
+                .unwrap()
+                .parse()
+                .unwrap(),
+        })
+    }
+
+    #[derive(Debug, PartialEq, AtatResp)]
+    struct CustomResponseParse {
+        arg1: u8,
+    }
+
+    #[derive(Debug, PartialEq, AtatCmd)]
+    #[at_cmd("+CFUN", CustomResponseParse, parse = custom_parse)]
+    struct RequestWithCustomResponseParse;
+
+    #[test]
+    fn test_custom_parse() {
+        assert_eq!(
+            RequestWithCustomResponseParse.parse(Ok(b"ignore123")),
+            Ok(CustomResponseParse { arg1: 123 })
+        );
+    }
 }
diff --git a/atat_derive/src/cmd.rs b/atat_derive/src/cmd.rs
@@ -17,6 +17,7 @@ pub fn atat_cmd(input: TokenStream) -> TokenStream {
     let CmdAttributes {
         cmd,
         resp,
+        parse,
         timeout_ms,
         attempts,
         reattempt_on_parse_err,
@@ -89,6 +90,32 @@ pub fn atat_cmd(input: TokenStream) -> TokenStream {
 
     let ident_len = format_ident!("ATAT_{}_LEN", ident.to_string().to_uppercase());
 
+    let parse = if let Some(parse) = parse {
+        quote! {
+            #[inline]
+            fn parse(&self, res: Result<&[u8], atat::InternalError>) -> core::result::Result<Self::Response, atat::Error> {
+                match res {
+                    Ok(resp) => #parse(resp).map_err(|e| {
+                        atat::Error::Parse
+                    }),
+                    Err(e) => Err(e.into())
+                }
+            }
+        }
+    } else {
+        quote! {
+            #[inline]
+           fn parse(&self, res: Result<&[u8], atat::InternalError>) -> core::result::Result<Self::Response, atat::Error> {
+               match res {
+                   Ok(resp) => atat::serde_at::from_slice::<#resp>(resp).map_err(|e| {
+                       atat::Error::Parse
+                   }),
+                   Err(e) => Err(e.into())
+               }
+           }
+        }
+    };
+
     TokenStream::from(quote! {
         #[automatically_derived]
         impl #impl_generics atat::AtatLen for #ident #ty_generics #where_clause {
@@ -124,15 +151,7 @@ pub fn atat_cmd(input: TokenStream) -> TokenStream {
                 }
             }
 
-            #[inline]
-            fn parse(&self, res: Result<&[u8], atat::InternalError>) -> core::result::Result<Self::Response, atat::Error> {
-                match res {
-                    Ok(resp) => atat::serde_at::from_slice::<#resp>(resp).map_err(|e| {
-                        atat::Error::Parse
-                    }),
-                    Err(e) => Err(e.into())
-                }
-            }
+            #parse
         }
 
         #[automatically_derived]

diff --git a/atat_derive/src/parse.rs b/atat_derive/src/parse.rs
@@ -1,8 +1,8 @@
 use proc_macro2::Span;
 use syn::parse::{Error, Parse, ParseStream, Result};
 use syn::{
-    Attribute, Data, DataEnum, DataStruct, DeriveInput, Expr, ExprLit, Fields, Generics, Ident,
-    Lit, LitByteStr, Path, Type,
+    Attribute, Data, DataEnum, DataStruct, DeriveInput, Expr, ExprLit, ExprPath, Fields, Generics,
+    Ident, Lit, LitByteStr, Path, Type,
 };
 
 #[derive(Clone)]
@@ -19,6 +19,7 @@ pub struct ParseInput {
 pub struct CmdAttributes {
     pub cmd: String,
     pub resp: Path,
+    pub parse: Option<Path>,
     pub timeout_ms: Option<u32>,
     pub attempts: Option<u8>,
     pub abortable: Option<bool>,
@@ -251,6 +252,7 @@ impl Parse for CmdAttributes {
         let mut at_cmd = Self {
             cmd: cmd.value(),
             resp: response_ident,
+            parse: None,
             timeout_ms: None,
             attempts: None,
             abortable: None,
@@ -291,6 +293,13 @@ impl Parse for CmdAttributes {
                         ))
                     }
                 }
+            } else if optional.path.is_ident("parse") {
+                match optional.value {
+                    Expr::Path(ExprPath { path, .. }) => {
+                        at_cmd.parse = Some(path);
+                    }
+                    _ => return Err(Error::new(call_site, "expected function for 'parse'")),
+                }
             } else if optional.path.is_ident("reattempt_on_parse_err") {
                 match optional.value {
                     Expr::Lit(ExprLit {