Correctly write custom labels to annotation file (#404)

Author: birm

tests/test_annotation.py

@@ -3,6 +3,7 @@
 import unittest
 
 import numpy as np
+import pandas as pd
 
 import wfdb
 
@@ -252,13 +253,31 @@ def test_4(self):
         annotation2 = wfdb.rdann("huge", "qrs")
         self.assertEqual(annotation1, annotation2)
 
+    def test_5(self):
+        """
+        Write and read annotations with custom labels.
+        """
+        ann_idx = np.array([1, 1000, 2000, 3000])
+        ann_chan = np.array([3, 1, 2, 3])
+        # write custom labels
+        ann_label_store = np.array([ 4, 2, 1, 3])
+        ann_custom_labels = {'label_store': [1, 2, 3, 4],
+                             'symbol': ['v','l','r','z'],
+                             'description':['pvc','lbbb','rbbb','pac']}
+        ann_custom_labels = pd.DataFrame(data=ann_custom_labels)
+        wfdb.wrann('CustomLabel', 'atr', ann_idx, chan=ann_chan,
+                   custom_labels=ann_custom_labels, label_store=ann_label_store)
+        ann = wfdb.rdann('CustomLabel', 'atr')
+        self.assertEqual(ann.symbol, ['z', 'l', 'v', 'r'])
+
     @classmethod
     def tearDownClass(cls):
         writefiles = [
             "100.atr",
             "1003.atr",
             "12726.anI",
             "huge.qrs",
+            "CustomLabel.atr"
         ]
         for file in writefiles:
             if os.path.isfile(file):

wfdb/io/annotation.py

@@ -1170,14 +1170,14 @@ def calc_core_bytes(self):
 
             # Process the samp (difference) and sym items
             data_bytes.append(
-                field2bytes("samptype", [sampdiff[i], self.symbol[i]])
+                field2bytes("samptype", [sampdiff[i], self.symbol[i]], self.custom_labels)
             )
 
             # Process the extra optional fields
             for field in extra_write_fields:
                 value = getattr(compact_annotation, field)[i]
                 if value is not None:
-                    data_bytes.append(field2bytes(field, value))
+                    data_bytes.append(field2bytes(field, value, self.custom_labels))
 
         # Flatten and convert to correct format
         data_bytes = np.array(
@@ -1600,7 +1600,7 @@ def compact_carry_field(full_field):
     return compact_field
 
 
-def field2bytes(field, value):
+def field2bytes(field, value, custom_labels=None):
     """
     Convert an annotation field into bytes to write.
 
@@ -1619,11 +1619,16 @@ def field2bytes(field, value):
     """
     data_bytes = []
 
+    # allow use of custom labels
+    label_table = ann_label_table
+    if custom_labels is not None:
+        label_table = pd.concat([label_table, custom_labels], ignore_index=True)
+
     # samp and sym bytes come together
     if field == "samptype":
         # Numerical value encoding annotation symbol
-        typecode = ann_label_table.loc[
-            ann_label_table["symbol"] == value[1], "label_store"
+        typecode = label_table.loc[
+            label_table["symbol"] == value[1], "label_store"
         ].values[0]
 
         # sample difference