Optimization of wfdb.io.annotation.field2bytes function

tests/test_record.py

@@ -589,7 +589,10 @@ def test_to_dataframe(self):
         self.assertEqual(record.sig_name, list(df.columns))
         self.assertEqual(len(df), record.sig_len)
         self.assertEqual(df.index[0], pd.Timedelta(0))
-        self.assertEqual(df.index[-1], pd.Timedelta(seconds=1 / record.fs * (record.sig_len - 1)))
+        self.assertEqual(
+            df.index[-1],
+            pd.Timedelta(seconds=1 / record.fs * (record.sig_len - 1)),
+        )
         assert np.array_equal(record.p_signal, df.values)
 
     def test_header_with_non_utf8(self):

wfdb/io/_signal.py

@@ -2505,4 +2505,3 @@ def _infer_sig_len(
     sig_len = int(data_size / (BYTES_PER_SAMPLE[fmt] * tsamps_per_frame))
 
     return sig_len
-

wfdb/io/annotation.py

@@ -1627,10 +1627,7 @@ def field2bytes(field, value, custom_labels=None):
     # samp and sym bytes come together
     if field == "samptype":
         # Numerical value encoding annotation symbol
-        typecode = label_table.loc[
-            label_table["symbol"] == value[1], "label_store"
-        ].values[0]
-
+        typecode = typecodes[value[1]]
         # sample difference
         sd = value[0]
 
@@ -3208,3 +3205,7 @@ def __str__(self):
 )
 ann_label_table.set_index(ann_label_table["label_store"].values, inplace=True)
 ann_label_table = ann_label_table[["label_store", "symbol", "description"]]
+typecodes = {
+    ann_label_table.iloc[i]["symbol"]: ann_label_table.iloc[i]["label_store"]
+    for i in range(len(ann_label_table))
+}

wfdb/io/record.py

@@ -1021,11 +1021,7 @@ def to_dataframe(self) -> pd.DataFrame:
         else:
             raise ValueError("No signal in record.")
 
-        return pd.DataFrame(
-            data=data,
-            index=index,
-            columns=self.sig_name
-        )
+        return pd.DataFrame(data=data, index=index, columns=self.sig_name)
 
 
 class MultiRecord(BaseRecord, _header.MultiHeaderMixin):