From 1d495305624ba80f88719a7829369affd20db419 Mon Sep 17 00:00:00 2001
From: Derrick Chambers <d-chambers@users.noreply.github.com>
Date: Fri, 1 Mar 2024 16:55:59 -0700
Subject: [PATCH] Fix 349 (#350)

---
 dascore/proc/units.py              |  3 ++-
 tests/test_proc/test_proc_units.py | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/dascore/proc/units.py b/dascore/proc/units.py
index 4020b44f..d12e9dc7 100644
--- a/dascore/proc/units.py
+++ b/dascore/proc/units.py
@@ -11,7 +11,6 @@
 def _update_attrs_coord_units(patch: dc.Patch, data_units, coords):
     """Update attributes with new units."""
     attrs = patch.attrs
-    data_units = data_units or attrs.get("data_units")
     # set data units
     attrs = attrs.update(data_units=data_units, coords=coords.to_summary_dict())
     return attrs
@@ -46,6 +45,8 @@ def set_units(
     >>> patch_with_units = patch.set_units("km/ms")
     >>> # set the units of the distance coordinate
     >>> patch_feet = patch.set_units(distance='feet')
+    >>> # remove data units
+    >>> patch_removed_units = patch_with_units.set_units(None)
     """
     new_coords = patch.coords.set_units(**kwargs)
     new_attrs = _update_attrs_coord_units(patch, data_units, new_coords)
diff --git a/tests/test_proc/test_proc_units.py b/tests/test_proc/test_proc_units.py
index 4969ddf5..52e96df2 100644
--- a/tests/test_proc/test_proc_units.py
+++ b/tests/test_proc/test_proc_units.py
@@ -64,6 +64,20 @@ def test_set_data_and_coord_units(self, random_patch):
         assert get_quantity(out.attrs.data_units) == get_quantity("m/s")
         assert get_quantity(out.attrs.distance_units) == get_quantity("ft")
 
+    def test_remove_units(self, random_patch):
+        """Ensure set coords can remove units."""
+        patch = random_patch.set_units("m", distance="m")
+        # ensure data units can be removed
+        new_none = patch.set_units(None)
+        assert new_none.attrs.data_units is None
+        new_empty_str = patch.set_units("")
+        assert new_empty_str.attrs.data_units is None
+        # ensure coord units can be removed
+        new_dist_none = random_patch.set_units(distance=None)
+        assert new_dist_none.get_coord("distance").units is None
+        new_dist_empty = random_patch.set_units(distance="")
+        assert new_dist_empty.get_coord("distance").units is None
+
 
 class TestConvertUnits:
     """Tests for converting from one unit to another."""