Remove unused elementpath check in SHHS sleep staging task

docs/api/tasks/pyhealth.tasks.sleep_staging.rst

@@ -3,4 +3,9 @@
 
 .. autofunction:: pyhealth.tasks.sleep_staging.sleep_staging_isruc_fn
 .. autofunction:: pyhealth.tasks.sleep_staging.sleep_staging_sleepedf_fn
-.. autofunction:: pyhealth.tasks.sleep_staging.sleep_staging_shhs_fn
+.. autofunction:: pyhealth.tasks.sleep_staging.sleep_staging_shhs_fn
+
+.. note::
+   ``sleep_staging_shhs_fn`` parses SHHS annotation files with the
+   standard-library :mod:`xml.etree.ElementTree` module only; no additional
+   packages are required.

examples/eeg/sleep_staging/sleep_staging_shhs_contrawr.py

@@ -1,8 +1,7 @@
-from pyhealth.datasets import split_by_patient, get_dataloader
-from pyhealth.trainer import Trainer
-from pyhealth.datasets import SHHSDataset
+from pyhealth.datasets import SHHSDataset, get_dataloader, split_by_patient
+from pyhealth.models import SparcNet
 from pyhealth.tasks import sleep_staging_shhs_fn
-from pyhealth.models import ContraWR, SparcNet
+from pyhealth.trainer import Trainer
 
 # step 1: load signal data
 dataset = SHHSDataset(

pyhealth/tasks/sleep_staging.py

@@ -1,8 +1,10 @@
 import os
 import pickle
+import xml.etree.ElementTree as ET
+
 import mne
-import pandas as pd
 import numpy as np
+import pandas as pd
 
 
 def sleep_staging_isruc_fn(record, epoch_seconds=10, label_id=1):
@@ -254,20 +256,6 @@ def sleep_staging_shhs_fn(record, epoch_seconds=30):
             'label': '0'
         }
     """
-
-    # test whether the ogb and torch_scatter packages are ready
-    dependencies = ["elementpath"]
-    try:
-        from importlib.metadata import version
-        version(dependencies)
-        import xml.etree.ElementTree as ET
-    except Exception as e:
-        print(e)
-        print ('-----------')
-        print(
-            "Please follow the error message and install the ['elementpath'] packages first."
-        )
-
     SAMPLE_RATE = 125
 
     root, signal_file, label_file, save_path = (

tests/core/test_sleep_staging_shhs.py

@@ -0,0 +1,91 @@
+import os
+import pickle
+import tempfile
+import unittest
+import xml.etree.ElementTree as ET
+from unittest import mock
+
+import numpy as np
+
+from pyhealth.tasks import sleep_staging
+from pyhealth.tasks.sleep_staging import sleep_staging_shhs_fn
+
+SHHS_XML = """<?xml version="1.0" encoding="UTF-8"?>
+<PSGAnnotation>
+    <EpochLength>30</EpochLength>
+    <SleepStages>
+        <SleepStage>0</SleepStage>
+        <SleepStage>4</SleepStage>
+        <SleepStage>2</SleepStage>
+        <SleepStage>1</SleepStage>
+    </SleepStages>
+</PSGAnnotation>
+"""
+
+
+class _FakeRawEDF:
+    """Minimal stand-in for the object returned by mne.io.read_raw_edf."""
+
+    def __init__(self, data):
+        self._data = data
+
+    def get_data(self):
+        return self._data
+
+
+class TestSleepStagingSHHS(unittest.TestCase):
+    """Test cases for sleep_staging_shhs_fn."""
+
+    def setUp(self):
+        """Set up a fake SHHS recording (signal path + XML label file)."""
+        self._tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self._tmp.cleanup)
+        self.root = os.path.join(self._tmp.name, "shhs")
+        self.save_path = os.path.join(self._tmp.name, "epochs")
+        os.makedirs(self.root)
+        os.makedirs(self.save_path)
+
+        self.signal_file = "shhs1-200001.edf"
+        self.label_file = "shhs1-200001.xml"
+        with open(os.path.join(self.root, self.label_file), "w") as f:
+            f.write(SHHS_XML)
+
+        self.record = [
+            {
+                "load_from_path": self.root,
+                "signal_file": self.signal_file,
+                "label_file": self.label_file,
+                "save_to_path": self.save_path,
+            }
+        ]
+
+    def test_xml_parser_available_at_module_level(self):
+        """ET must be bound unconditionally so label parsing cannot NameError."""
+        self.assertIs(sleep_staging.ET, ET)
+
+    def test_parses_sleep_stage_labels_from_xml(self):
+        """The task reads SleepStage labels via ElementTree, with no extra deps."""
+        num_epochs = 4
+        # 14 channels: skips both missing-channel branches in the task fn.
+        signal = np.zeros((14, num_epochs * 125))
+        with mock.patch.object(
+            sleep_staging.mne.io,
+            "read_raw_edf",
+            return_value=_FakeRawEDF(signal),
+        ):
+            samples = sleep_staging_shhs_fn(self.record, epoch_seconds=1)
+
+        self.assertEqual(len(samples), num_epochs)
+        for i, sample in enumerate(samples):
+            self.assertEqual(sample["record_id"], f"shhs1-200001-{i}")
+            self.assertEqual(sample["patient_id"], "shhs1-200001")
+            self.assertEqual(sample["label"], ["0", "4", "2", "1"][i])
+            self.assertTrue(os.path.exists(sample["epoch_path"]))
+            with open(sample["epoch_path"], "rb") as f:
+                epoch = pickle.load(f)
+            self.assertEqual(epoch["signal"].shape, (2, 125))
+            self.assertEqual(epoch["label"], sample["label"])
+
+
+if __name__ == "__main__":
+    unittest.main()