From 7ac8cd9e528e797d593240f956f00cf33dfaa997 Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Fri, 26 Apr 2024 14:46:28 +0100
Subject: [PATCH] FIx CSV output (missing comma) and include number of
 processes in CSV output

---
 psrecord/main.py            | 9 ++++++---
 psrecord/tests/test_main.py | 2 +-
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/psrecord/main.py b/psrecord/main.py
index 81bff6a..1e324d7 100644
--- a/psrecord/main.py
+++ b/psrecord/main.py
@@ -176,7 +176,7 @@ def monitor(
                     )
                 )
         elif log_format == "csv":
-            f.write("elapsed_time,cpu,mem_real,mem_virtual")
+            f.write("elapsed_time,nproc,cpu,mem_real,mem_virtual")
             if include_io:
                 f.write(",read_count,write_count,read_bytes,write_bytes")
         else:
@@ -236,6 +236,8 @@ def monitor(
                 read_bytes = counters.read_bytes
                 write_bytes = counters.write_bytes
 
+            n_proc = 1
+
             # Get information for children
             if include_children:
                 for child in all_children(pr):
@@ -250,6 +252,7 @@ def monitor(
                             write_count += counters.write_count
                             read_bytes += counters.read_bytes
                             write_bytes += counters.write_bytes
+                        n_proc += 1
                     except Exception:
                         continue
 
@@ -266,10 +269,10 @@ def monitor(
                         )
                 elif log_format == "csv":
                     f.write(
-                        f"{elapsed_time},{current_cpu},{current_mem_real},{current_mem_virtual}"
+                        f"{elapsed_time},{n_proc},{current_cpu},{current_mem_real},{current_mem_virtual}"
                     )
                     if include_io:
-                        f.write(f"{read_count},{write_count},{read_bytes},{write_bytes}")
+                        f.write(f",{read_count},{write_count},{read_bytes},{write_bytes}")
                 f.write("\n")
                 f.flush()
 
diff --git a/psrecord/tests/test_main.py b/psrecord/tests/test_main.py
index de23c42..9e0037d 100644
--- a/psrecord/tests/test_main.py
+++ b/psrecord/tests/test_main.py
@@ -63,7 +63,7 @@ def test_logfile_csv(self, tmpdir):
         assert len(open(filename).readlines()) > 0
         with open(filename) as csvfile:
             data = csv.reader(csvfile)
-            assert next(data) == ["elapsed_time", "cpu", "mem_real", "mem_virtual"]
+            assert next(data) == ["elapsed_time", "nproc", "cpu", "mem_real", "mem_virtual"]
 
     def test_plot(self, tmpdir):
         pytest.importorskip("matplotlib")