From 96675d2116158ae7f089bbeb0cbac0c6f68d5865 Mon Sep 17 00:00:00 2001
From: Max Grover <maxi2312@gmail.com>
Date: Thu, 28 Mar 2024 12:45:55 -0500
Subject: [PATCH] ENH: Add a utility for finding sweep number keys (#167)

* ADD: Add a utility for finding sweep number keys

* DOC: Add to history file

* ADD: Add check for sweep in first part

* ADD: Add fake group to test
---
 docs/history.md    |  1 +
 tests/test_util.py | 12 +++++++++++-
 xradar/util.py     | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/docs/history.md b/docs/history.md
index bff2bfd2..26f43af4 100644
--- a/docs/history.md
+++ b/docs/history.md
@@ -8,6 +8,7 @@
 * MNT: fix path for notebook coverage ({pull}`157`) by [@kmuehlbauer](https://github.com/kmuehlbauer).
 * ADD: NEXRAD Level2 structured reader ({pull}`158`) by [@kmuehlbauer](https://github.com/kmuehlbauer) and [@mgrover1](https://github.com/mgrover1).
 * FIX: Add the proper elevation angle to fixed angle ({pull}`162`) by [@mgrover1](https://github.com/mgrover1).
+* ENH: Add a utility for finding sweep number keys ({pull}`167`) by [@mgrover1](https://github.com/mgrover1).
 
 ## 0.4.3 (2024-02-24)
 
diff --git a/tests/test_util.py b/tests/test_util.py
index b3828791..1552f5be 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -9,7 +9,7 @@
 import xarray as xr
 from open_radar_data import DATASETS
 
-from xradar import model, util
+from xradar import io, model, util
 
 
 @pytest.fixture(
@@ -254,3 +254,13 @@ def test_ipol_time2(missing, a1gate, rot):
             UserWarning, match="Rays might miss on beginning and/or end of sweep."
         ):
             dsx.pipe(util.ipol_time, direction=direction)
+
+
+def test_get_sweep_keys():
+    # Test finding sweep keys
+    filename = DATASETS.fetch("sample_sgp_data.nc")
+    dt = io.open_cfradial1_datatree(filename)
+    # set a fake group
+    dt["sneep_1"] = dt["sweep_1"]
+    keys = util.get_sweep_keys(dt)
+    assert keys == ["sweep_0", "sweep_1", "sweep_2", "sweep_3", "sweep_4", "sweep_5"]
diff --git a/xradar/util.py b/xradar/util.py
index 83666528..ea27430f 100644
--- a/xradar/util.py
+++ b/xradar/util.py
@@ -22,6 +22,7 @@
     "extract_angle_parameters",
     "ipol_time",
     "rolling_dim",
+    "get_sweep_keys",
 ]
 
 __doc__ = __doc__.format("\n   ".join(__all__))
@@ -487,3 +488,37 @@ def rolling_dim(data, window):
     shape = data.shape[:-1] + (data.shape[-1] - window + 1, window)
     strides = data.strides + (data.strides[-1],)
     return np.lib.stride_tricks.as_strided(data, shape=shape, strides=strides)
+
+
+def get_sweep_keys(dt):
+    """Return which nodes in the datatree contain sweep variables
+
+    Parameters
+    ----------
+    dt : xarray.DataTree
+        Datatree to check for sweep_n keys
+
+    Returns
+    -------
+    keys : list
+        List of associated keys with sweep_n
+    """
+    sweep_group_keys = []
+    for key in list(dt.children):
+        parts = key.split("_")
+        try:
+            # Try to set the second part of the tree key to an int
+            int(parts[1])
+            # Check for "sweep" in the first part of the key
+            assert "sweep" in parts[0]
+            sweep_group_keys.append(key)
+
+        # This would fail with strings - ex. sweep_group_attrs
+        except ValueError:
+            pass
+
+        # This would fail if "sweep" not in key - ex. radar_parameters
+        except AssertionError:
+            pass
+
+    return sweep_group_keys