From 25a21844cb0adb8586bb9151aee607fa2e79060f Mon Sep 17 00:00:00 2001
From: Sadie Louise Bartholomew <sadie.bartholomew@ncas.ac.uk>
Date: Tue, 22 Nov 2022 19:43:15 +0000
Subject: [PATCH] Query.set_condition_units: consolidate w/ inner helper
 function

---
 cf/query.py | 74 +++++++++++++++++++++++++++--------------------------
 1 file changed, 38 insertions(+), 36 deletions(-)

diff --git a/cf/query.py b/cf/query.py
index 68ae24fbd3..9ae6eda78b 100644
--- a/cf/query.py
+++ b/cf/query.py
@@ -858,13 +858,40 @@ def set_condition_units(self, units):
         <CF Query: (ge 3000 m)>
 
         """
+        # def try_applying_units(u, val, q_units):
+        #     try:
+        #         val = val.copy()
+        #         val.Units = u
+        #     except ValueError:
+        #         raise ValueError(
+        #             f"Units {u!r} are not equivalent to "
+        #             f"query condition units {q_units!r}"
+        #         )
+        #     return val
+
+        def get_and_set_value_units(v, u):
+            """Helper method to simplify logic to set specified units."""
+            v_units = getattr(v, "Units", None)
+            if v_units is None:  # Value 'v' has no units
+                v = Data(v, units=u)
+            else:  # Value 'v' already has units
+                try:
+                    v = v.copy()
+                    v.Units = u
+                except ValueError:
+                    raise ValueError(
+                        f"Units {u!r} are not equivalent to "
+                        f"query condition units {v_units!r}"
+                    )
+
+            return v
+
         units = Units(units)
 
         compound = self._compound
         if compound:
             for r in compound:
                 r.set_condition_units(units)
-
             return
 
         value = self._value
@@ -872,45 +899,20 @@ def set_condition_units(self, units):
             return
 
         value_units = getattr(value, "Units", None)
-        if value_units is None:
-            # Value has no units
-            if self.operator in ("wi", "wo", "set"):
-                # value is a sequence of things that may or may not
-                # already have units
-                new = []
-                for v in value:
-                    v_units = getattr(v, "Units", None)
-                    if v_units is None:
-                        v = Data(v, units=units)
-                    else:
-                        try:
-                            v = v.copy()
-                            v.Units = units
-                        except ValueError:
-                            raise ValueError(
-                                f"Units {units!r} are not equivalent to "
-                                f"query condition units {v_units!r}"
-                            )
-
-                    new.append(v)
-
-                value = new
-            else:
-                value = Data(value, units=units)
+        if self.operator in ("wi", "wo", "set"):
+            # Value is a sequence of things that may or may not
+            # already have units
+            new_values = []
+            for v in value:
+                v = get_and_set_value_units(v, units)
+                new_values.append(v)
+
+            value = new_values
         else:
-            # Value already has units
-            try:
-                value = value.copy()
-                value.Units = units
-            except ValueError:
-                raise ValueError(
-                    f"Units {units!r} are not equivalent to "
-                    f"query condition units {value_units!r}"
-                )
+            value = get_and_set_value_units(value, units)
 
         self._value = value
 
-        self._value = value
 
     # ----------------------------------------------------------------
     # Deprecated attributes and methods