Merge pull request #76 from astrofrog/csv-fix

FIx CSV output (missing comma) and include number of processes in CSV output
astrofrog · Apr 26, 2024 · b76d59f · b76d59f
2 parents 34ef8d6 + 7ac8cd9
commit b76d59f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git 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
@@ -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")