feat: Add trait StyleVariableHolder

2 files changed, 30 insertions, 15 deletions

diff --git a/src/formatter/model.rs b/src/formatter/model.rs
index c0a2269ac..27839400f 100644
--- a/src/formatter/model.rs
+++ b/src/formatter/model.rs
@@ -6,6 +6,11 @@ pub trait VariableHolder<T> {
     fn get_variables(&self) -> BTreeSet<T>;
 }
 
+/// Type that holds a number of style variables of type `T`
+pub trait StyleVariableHolder<T> {
+    fn get_style_variables(&self) -> BTreeSet<T>;
+}
+
 pub struct TextGroup<'a> {
     pub format: Vec<FormatElement<'a>>,
     pub style: Vec<StyleElement<'a>>,
@@ -47,8 +52,8 @@ impl<'a> VariableHolder<Cow<'a, str>> for Vec<FormatElement<'a>> {
     }
 }
 
-impl<'a> VariableHolder<Cow<'a, str>> for StyleElement<'a> {
-    fn get_variables(&self) -> BTreeSet<Cow<'a, str>> {
+impl<'a> StyleVariableHolder<Cow<'a, str>> for StyleElement<'a> {
+    fn get_style_variables(&self) -> BTreeSet<Cow<'a, str>> {
         match self {
             StyleElement::Variable(var) => {
                 let mut variables = BTreeSet::new();
@@ -60,11 +65,23 @@ impl<'a> VariableHolder<Cow<'a, str>> for StyleElement<'a> {
     }
 }
 
-impl<'a> VariableHolder<Cow<'a, str>> for Vec<&StyleElement<'a>> {
-    fn get_variables(&self) -> BTreeSet<Cow<'a, str>> {
+impl<'a> StyleVariableHolder<Cow<'a, str>> for Vec<StyleElement<'a>> {
+    fn get_style_variables(&self) -> BTreeSet<Cow<'a, str>> {
         self.iter().fold(BTreeSet::new(), |mut acc, el| {
-            acc.extend(el.get_variables());
+            acc.extend(el.get_style_variables());
             acc
         })
     }
 }
+
+impl<'a> StyleVariableHolder<Cow<'a, str>> for Vec<FormatElement<'a>> {
+    fn get_style_variables(&self) -> BTreeSet<Cow<'a, str>> {
+        self.iter().fold(BTreeSet::new(), |mut acc, el| match el {
+            FormatElement::TextGroup(textgroup) => {
+                acc.extend(textgroup.style.get_style_variables());
+                acc
+            }
+            _ => acc,
+        })
+    }
+}
diff --git a/src/formatter/string_formatter.rs b/src/formatter/string_formatter.rs
index 830cf554f..682cfcf3b 100644
--- a/src/formatter/string_formatter.rs
+++ b/src/formatter/string_formatter.rs
@@ -44,17 +44,9 @@ impl<'a> StringFormatter<'a> {
                         .map(|key| (key.to_string(), None))
                         .collect::<Vec<(String, Option<_>)>>(),
                 );
-                let style_elements = format
-                    .iter()
-                    .flat_map(|el| match el {
-                        FormatElement::TextGroup(textgroup) => Some(&textgroup.style),
-                        _ => None,
-                    })
-                    .flatten()
-                    .collect::<Vec<&StyleElement>>();
                 let style_variables = StyleVariableMapType::from_iter(
-                    style_elements
-                        .get_variables()
+                    format
+                        .get_style_variables()
                         .into_iter()
                         .map(|key| (key.to_string(), None))
                         .collect::<Vec<(String, Option<_>)>>(),
@@ -204,6 +196,12 @@ impl<'a> VariableHolder<String> for StringFormatter<'a> {
     }
 }
 
+impl<'a> StyleVariableHolder<String> for StringFormatter<'a> {
+    fn get_style_variables(&self) -> BTreeSet<String> {
+        BTreeSet::from_iter(self.style_variables.keys().cloned())
+    }
+}
+
 /// Helper function to create a new segment
 fn _new_segment(name: String, value: String, style: Option<Style>) -> Segment {
     Segment {