Add assertionchecks all strings

raphaelvallat · Dec 22, 2024 · 82506f5 · 82506f5
1 parent 3f92f8b
commit 82506f5
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/yasa/hypno.py b/yasa/hypno.py
@@ -223,6 +223,10 @@ def __init__(self, values, n_stages=5, *, freq="30s", start=None, scorer=None, p
         assert isinstance(
             values, (list, np.ndarray, pd.Series)
         ), "`values` must be a list, numpy.array or pandas.Series"
+        assert all(isinstance(val, str) for val in values), (
+            "Since v0.7, YASA expects strings to represent sleep stages, e.g. ['WAKE', 'N1', ...]. "
+            "Please refer to the documentation for more details."
+        )
         assert isinstance(n_stages, int), "`n_stages` must be an integer between 2 and 5."
         assert n_stages in [2, 3, 4, 5], "`n_stages` must be an integer between 2 and 5."
         assert isinstance(freq, str), "`freq` must be a pandas frequency string."