From 9860e057da675e6771bbcbaeba0c1b9dfe380289 Mon Sep 17 00:00:00 2001
From: John Howard <howardjohn@google.com>
Date: Tue, 3 Jan 2023 11:44:01 -0800
Subject: [PATCH] Address comments

Signed-off-by: John Howard <howardjohn@google.com>
---
 derive-encode/src/lib.rs   | 17 +++++++++++------
 derive-encode/tests/lib.rs | 13 ++++++-------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/derive-encode/src/lib.rs b/derive-encode/src/lib.rs
index 2ed3f062..e126889e 100644
--- 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
index 8cd4df7a..fba8412d 100644
--- 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();