Allow custom labels in field2bytes

birm · birm · commit 4fcf27165c62 · 2022-06-30T10:27:03.000-04:00
diff --git a/wfdb/io/annotation.py b/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,17 @@ 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)
+        print(label_table)
+
     # 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
