fix: allow ident name start with _

cloudwego · Oct 25, 2024 · b4a9311 · b4a9311
1 parent 5a06bf8
commit b4a9311
Show file tree

Hide file tree

Showing 7 changed files with 154 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pilota-build/Cargo.toml b/pilota-build/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pilota-build"
-version = "0.11.21"
+version = "0.11.22"
 edition = "2021"
 description = "Compile thrift and protobuf idl into rust code at compile-time."
 documentation = "https://docs.rs/pilota-build"

diff --git a/pilota-build/src/symbol.rs b/pilota-build/src/symbol.rs
@@ -151,26 +151,37 @@ pub trait IdentName {
 
 impl IdentName for &str {
     fn upper_camel_ident(&self) -> FastStr {
-        let s = self.to_upper_camel_case();
-        s.into()
+        if let Some(index) = self.find(|c: char| c != '_') {
+            let s = self[index..].to_upper_camel_case();
+            return format!("{}{}", &self[0..index], s).into();
+        }
+        self.to_string().into()
     }
 
     fn snake_ident(&self) -> FastStr {
-        if is_common_initialism(self) {
-            to_snake_case(self)
-        } else {
-            self.to_snake_case()
+        if let Some(index) = self.find(|c: char| c != '_') {
+            let s = &self[index..];
+            let s = if is_common_initialism(s) {
+                to_snake_case(s)
+            } else {
+                s.to_snake_case()
+            };
+            return format!("{}{}", &self[0..index], s).into();
         }
-        .into()
+        self.to_string().into()
     }
 
     fn shouty_snake_case(&self) -> FastStr {
-        if is_common_initialism(self) {
-            to_snake_case(self).to_uppercase()
-        } else {
-            self.to_shouty_snake_case()
+        if let Some(index) = self.find(|c: char| c != '_') {
+            let s = &self[index..];
+            let s = if is_common_initialism(s) {
+                to_snake_case(s).to_uppercase()
+            } else {
+                s.to_shouty_snake_case()
+            };
+            return format!("{}{}", &self[0..index], s).into();
         }
-        .into()
+        self.to_string().into()
     }
 }
 

diff --git a/pilota-build/test_data/thrift/enum_test.rs b/pilota-build/test_data/thrift/enum_test.rs
@@ -1086,6 +1086,8 @@ pub mod enum_test {
             pub Index: Index,
 
             pub index: ::std::option::Option<Index>,
+
+            pub _enum: ::std::option::Option<_Enum>,
         }
         impl ::pilota::thrift::Message for Request {
             fn encode<T: ::pilota::thrift::TOutputProtocol>(
@@ -1101,6 +1103,9 @@ pub mod enum_test {
                 if let Some(value) = self.index.as_ref() {
                     __protocol.write_i32_field(2, (value).inner())?;
                 }
+                if let Some(value) = self._enum.as_ref() {
+                    __protocol.write_i32_field(3, (value).inner())?;
+                }
                 __protocol.write_field_stop()?;
                 __protocol.write_struct_end()?;
                 ::std::result::Result::Ok(())
@@ -1114,6 +1119,7 @@ pub mod enum_test {
 
                 let mut var_1 = None;
                 let mut var_2 = None;
+                let mut var_3 = None;
 
                 let mut __pilota_decoding_field_id = None;
 
@@ -1135,6 +1141,9 @@ pub mod enum_test {
                             Some(2) if field_ident.field_type == ::pilota::thrift::TType::I32 => {
                                 var_2 = Some(::pilota::thrift::Message::decode(__protocol)?);
                             }
+                            Some(3) if field_ident.field_type == ::pilota::thrift::TType::I32 => {
+                                var_3 = Some(::pilota::thrift::Message::decode(__protocol)?);
+                            }
                             _ => {
                                 __protocol.skip(field_ident.field_type)?;
                             }
@@ -1165,6 +1174,7 @@ pub mod enum_test {
                 let data = Self {
                     Index: var_1,
                     index: var_2,
+                    _enum: var_3,
                 };
                 ::std::result::Result::Ok(data)
             }
@@ -1182,6 +1192,7 @@ pub mod enum_test {
                 ::std::boxed::Box::pin(async move {
                     let mut var_1 = None;
                     let mut var_2 = None;
+                    let mut var_3 = None;
 
                     let mut __pilota_decoding_field_id = None;
 
@@ -1215,6 +1226,16 @@ pub mod enum_test {
                                         .await?,
                                     );
                                 }
+                                Some(3)
+                                    if field_ident.field_type == ::pilota::thrift::TType::I32 =>
+                                {
+                                    var_3 = Some(
+                                        <_Enum as ::pilota::thrift::Message>::decode_async(
+                                            __protocol,
+                                        )
+                                        .await?,
+                                    );
+                                }
                                 _ => {
                                     __protocol.skip(field_ident.field_type).await?;
                                 }
@@ -1248,6 +1269,7 @@ pub mod enum_test {
                     let data = Self {
                         Index: var_1,
                         index: var_2,
+                        _enum: var_3,
                     };
                     ::std::result::Result::Ok(data)
                 })
@@ -1262,10 +1284,104 @@ pub mod enum_test {
                     + self.index.as_ref().map_or(0, |value| {
                         __protocol.i32_field_len(Some(2), (value).inner())
                     })
+                    + self._enum.as_ref().map_or(0, |value| {
+                        __protocol.i32_field_len(Some(3), (value).inner())
+                    })
                     + __protocol.field_stop_len()
                     + __protocol.struct_end_len()
             }
         }
+        #[derive(PartialOrd, Hash, Eq, Ord, Debug, ::pilota::derivative::Derivative)]
+        #[derivative(Default)]
+        #[derive(Clone, PartialEq, Copy)]
+        #[repr(transparent)]
+        pub struct _Enum(i32);
+
+        impl _Enum {
+            pub const _1: Self = Self(1);
+            pub const _2: Self = Self(2);
+
+            pub fn inner(&self) -> i32 {
+                self.0
+            }
+
+            pub fn to_string(&self) -> ::std::string::String {
+                match self {
+                    Self(1) => ::std::string::String::from("_1"),
+                    Self(2) => ::std::string::String::from("_2"),
+                    Self(val) => val.to_string(),
+                }
+            }
+        }
+
+        impl ::std::convert::From<i32> for _Enum {
+            fn from(value: i32) -> Self {
+                Self(value)
+            }
+        }
+
+        impl ::std::convert::From<_Enum> for i32 {
+            fn from(value: _Enum) -> i32 {
+                value.0
+            }
+        }
+
+        impl ::pilota::thrift::Message for _Enum {
+            fn encode<T: ::pilota::thrift::TOutputProtocol>(
+                &self,
+                __protocol: &mut T,
+            ) -> ::std::result::Result<(), ::pilota::thrift::ThriftException> {
+                #[allow(unused_imports)]
+                use ::pilota::thrift::TOutputProtocolExt;
+                __protocol.write_i32(self.inner())?;
+                ::std::result::Result::Ok(())
+            }
+
+            fn decode<T: ::pilota::thrift::TInputProtocol>(
+                __protocol: &mut T,
+            ) -> ::std::result::Result<Self, ::pilota::thrift::ThriftException> {
+                #[allow(unused_imports)]
+                use ::pilota::{thrift::TLengthProtocolExt, Buf};
+                let value = __protocol.read_i32()?;
+                ::std::result::Result::Ok(::std::convert::TryFrom::try_from(value).map_err(
+                    |err| {
+                        ::pilota::thrift::new_protocol_exception(
+                            ::pilota::thrift::ProtocolExceptionKind::InvalidData,
+                            format!("invalid enum value for _Enum, value: {}", value),
+                        )
+                    },
+                )?)
+            }
+
+            fn decode_async<'a, T: ::pilota::thrift::TAsyncInputProtocol>(
+                __protocol: &'a mut T,
+            ) -> ::std::pin::Pin<
+                ::std::boxed::Box<
+                    dyn ::std::future::Future<
+                            Output = ::std::result::Result<Self, ::pilota::thrift::ThriftException>,
+                        > + Send
+                        + 'a,
+                >,
+            > {
+                ::std::boxed::Box::pin(async move {
+                    let value = __protocol.read_i32().await?;
+                    ::std::result::Result::Ok(::std::convert::TryFrom::try_from(value).map_err(
+                        |err| {
+                            ::pilota::thrift::new_protocol_exception(
+                                ::pilota::thrift::ProtocolExceptionKind::InvalidData,
+                                format!("invalid enum value for _Enum, value: {}", value),
+                            )
+                        },
+                    )?)
+                })
+            }
+
+            fn size<T: ::pilota::thrift::TLengthProtocol>(&self, __protocol: &mut T) -> usize {
+                #[allow(unused_imports)]
+                use ::pilota::thrift::TLengthProtocolExt;
+                __protocol.i32_len(self.inner())
+            }
+        }
         pub trait Test {}
         #[derive(PartialOrd, Hash, Eq, Ord, Debug, ::pilota::derivative::Derivative)]
         #[derivative(Default)]

diff --git a/pilota-build/test_data/thrift/enum_test.thrift b/pilota-build/test_data/thrift/enum_test.thrift
@@ -10,9 +10,15 @@ enum Err {
 enum Ok {
 }
 
+enum _Enum {
+    _1 = 1,
+    _2 = 2
+}
+
 struct Request {
     1: required Index Index,
     2: Index index,
+    3: _Enum _enum,
 }
 service Test {
    Err test_enum(1: Ok req);

diff --git a/pilota-thrift-parser/Cargo.toml b/pilota-thrift-parser/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pilota-thrift-parser"
-version = "0.11.1"
+version = "0.11.2"
 edition = "2021"
 description = "Pilota thrift Parser."
 documentation = "https://docs.rs/pilota"

diff --git a/pilota-thrift-parser/src/parser/identifier.rs b/pilota-thrift-parser/src/parser/identifier.rs
@@ -1,9 +1,6 @@
 use nom::{
-    bytes::complete::take_while,
-    character::complete::{char as cchar, satisfy},
-    combinator::recognize,
-    sequence::tuple,
-    IResult,
+    bytes::complete::take_while, character::complete::satisfy, combinator::recognize,
+    sequence::tuple, IResult,
 };
 
 use super::super::{descriptor::Ident, parser::*};
@@ -15,8 +12,7 @@ impl Parser for Ident {
     fn parse(input: &str) -> IResult<&str, Ident> {
         map(
             recognize(tuple((
-                many0(cchar('_')),
-                satisfy(|c| c.is_ascii_alphabetic()),
+                satisfy(|c| c.is_ascii_alphabetic() || c == '_'),
                 take_while(|c: char| c.is_ascii_alphanumeric() || c == '_'),
             ))),
             |ident: &str| -> Ident { Ident(ident.into()) },
@@ -42,8 +38,8 @@ mod test {
 
         assert_eq!(Ident::parse("_ihciah,").unwrap().1, "_ihciah");
         assert_eq!(Ident::parse("ihciah,").unwrap().1, "ihciah");
-        assert!(Ident::parse("_123").is_err());
-        assert!(Ident::parse("_").is_err());
+        assert_eq!(Ident::parse("_123").unwrap().1, "_123");
+        assert_eq!(Ident::parse("_").unwrap().1, "_");
         assert!(Ident::parse("123").is_err());
     }