Address comments

Signed-off-by: John Howard <[email protected]>
prometheus · Jan 3, 2023 · 9860e05 · 9860e05
1 parent e42a284
commit 9860e05
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/derive-encode/src/lib.rs b/derive-encode/src/lib.rs
@@ -22,19 +22,24 @@ pub fn derive_encode_label_set(input: TokenStream) -> TokenStream {
             syn::Fields::Named(syn::FieldsNamed { named, .. }) => named
                 .into_iter()
                 .map(|f| {
-                    let flatten = f
+                    let attribute = f
                         .attrs
                         .iter()
                         .find(|a| a.path.is_ident("prometheus"))
-                        .map(|a| a.parse_args::<syn::Ident>().unwrap().to_string() == "flatten")
-                        .unwrap_or(false);
+                        .map(|a| a.parse_args::<syn::Ident>().unwrap().to_string());
+                    let flatten = match attribute.as_deref() {
+                        Some("flatten") => true,
+                        Some(other) => {
+                            panic!("Provided attribute '{other}', but only 'flatten' is supported")
+                        }
+                        None => false,
+                    };
+                    let ident = f.ident.unwrap();
                     if flatten {
-                        let ident = f.ident.unwrap();
                         quote! {
-                            self.#ident.encode(encoder)?;
+                             EncodeLabelSet::encode(&self.#ident, encoder)?;
                         }
                     } else {
-                        let ident = f.ident.unwrap();
                         let ident_string = KEYWORD_IDENTIFIERS
                             .iter()
                             .find(|pair| ident == pair.1)

diff --git a/derive-encode/tests/lib.rs b/derive-encode/tests/lib.rs
@@ -156,13 +156,12 @@ fn flatten() {
     registry.register("my_counter", "This is my counter", family.clone());
 
     // Record a single HTTP GET request.
-    family.get_or_create(&Labels {
-        unique: 1,
-        common: CommonLabels{
-            a: 2,
-            b: 3,
-        },
-    }).inc();
+    family
+        .get_or_create(&Labels {
+            unique: 1,
+            common: CommonLabels { a: 2, b: 3 },
+        })
+        .inc();
 
     // Encode all metrics in the registry in the text format.
     let mut buffer = String::new();