From d6d5aa7ce398c912e1b402ab9f19e1af302cfdc7 Mon Sep 17 00:00:00 2001
From: Tom Pollard <tpollard@mit.edu>
Date: Mon, 1 Jul 2024 15:29:35 -0400
Subject: [PATCH] Add output_dir argument to csv_to_wfdb. Fixes #67.

---
 wfdb/io/convert/csv.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/wfdb/io/convert/csv.py b/wfdb/io/convert/csv.py
index 3cfd25a2..42373401 100644
--- a/wfdb/io/convert/csv.py
+++ b/wfdb/io/convert/csv.py
@@ -33,6 +33,7 @@ def csv_to_wfdb(
     header=True,
     delimiter=",",
     verbose=False,
+    output_dir=None
 ):
     """
     Read a WFDB header file and return either a `Record` object with the
@@ -235,6 +236,9 @@ def csv_to_wfdb(
     verbose : bool, optional
         Whether to print all the information read about the file (True) or
         not (False).
+    output_dir : str, optional
+        The directory where the output files will be saved. If not provided,
+        the output files will be saved in the same directory as the input file.
 
     Returns
     -------
@@ -291,6 +295,7 @@ def csv_to_wfdb(
         df_CSV = pd.read_csv(file_name, delimiter=delimiter, header=None)
     if verbose:
         print("Successfully read CSV")
+
     # Extract the entire signal from the dataframe
     p_signal = df_CSV.values
     # The dataframe should be in (`sig_len`, `n_sig`) dimensions
@@ -300,6 +305,7 @@ def csv_to_wfdb(
     n_sig = p_signal.shape[1]
     if verbose:
         print("Number of signals: {}".format(n_sig))
+
     # Check if signal names are valid and set defaults
     if not sig_name:
         if header:
@@ -318,15 +324,23 @@ def csv_to_wfdb(
             if verbose:
                 print("Signal names: {}".format(sig_name))
 
-    # Set the output header file name to be the same, remove path
-    if os.sep in file_name:
-        file_name = file_name.split(os.sep)[-1]
-    record_name = file_name.replace(".csv", "")
+    # Determine the output directory
+    if output_dir:
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+        output_base = os.path.join(
+            output_dir, os.path.basename(file_name).replace(".csv", "")
+            )
+    else:
+        if os.sep in file_name:
+            file_name = file_name.split(os.sep)[-1]
+        output_base = file_name.replace(".csv", "")
+
     if verbose:
-        print("Output header: {}.hea".format(record_name))
+        print("Output base: {}".format(output_base))
 
     # Replace the CSV file tag with DAT
-    dat_file_name = file_name.replace(".csv", ".dat")
+    dat_file_name = output_base + ".dat"
     dat_file_name = [dat_file_name] * n_sig
     if verbose:
         print("Output record: {}".format(dat_file_name[0]))
@@ -419,7 +433,7 @@ def csv_to_wfdb(
     if record_only:
         # Create the record from the input and generated values
         record = Record(
-            record_name=record_name,
+            record_name=output_base,
             n_sig=n_sig,
             fs=fs,
             samps_per_frame=samps_per_frame,
@@ -454,7 +468,7 @@ def csv_to_wfdb(
     else:
         # Write the information to a record and header file
         wrsamp(
-            record_name=record_name,
+            record_name=output_base,
             fs=fs,
             units=units,
             sig_name=sig_name,