Allow stancsv parse function to accept filename/path

Author: amas0

cmdstanpy/stanfit/mcmc.py

@@ -444,8 +444,9 @@ def _assemble_draws(self) -> None:
 
         mass_matrix_per_chain: List[Optional[npt.NDArray[np.float64]]] = []
         for chain in range(self.chains):
-            with open(self.runset.csv_files[chain], "rb") as f:
-                comments, draws = stancsv.parse_stan_csv_comments_and_draws(f)
+            comments, draws = stancsv.parse_stan_csv_comments_and_draws(
+                self.runset.csv_files[chain]
+            )
 
             self._draws[:, chain, :] = stancsv.csv_bytes_list_to_numpy(draws)
 

cmdstanpy/utils/stancsv.py

@@ -7,6 +7,7 @@
 import math
 import re
 import warnings
+from pathlib import Path
 from typing import (
     Any,
     Dict,
@@ -27,21 +28,30 @@
 
 
 def parse_stan_csv_comments_and_draws(
-    lines: Iterator[bytes],
+    stan_csv: Union[str, Path, Iterator[bytes]],
 ) -> Tuple[List[bytes], List[bytes]]:
     """Parses lines of a Stan CSV file into comment lines and draws lines, where
     a draws line is just a non-commented line.
 
     Returns a (comment_lines, draws_lines) tuple.
     """
-    comment_lines, draws_lines = [], []
 
-    for line in lines:
-        if line.startswith(b"#"):  # is comment line
-            comment_lines.append(line)
-        else:
-            draws_lines.append(line)
-    return comment_lines, draws_lines
+    def split_comments_and_draws(
+        lines: Iterator[bytes],
+    ) -> Tuple[List[bytes], List[bytes]]:
+        comment_lines, draws_lines = [], []
+        for line in lines:
+            if line.startswith(b"#"):  # is comment line
+                comment_lines.append(line)
+            else:
+                draws_lines.append(line)
+        return comment_lines, draws_lines
+
+    if isinstance(stan_csv, (str, Path)):
+        with open(stan_csv, "rb") as f:
+            return split_comments_and_draws(f)
+    else:
+        return split_comments_and_draws(stan_csv)
 
 
 def csv_bytes_list_to_numpy(

test/test_stancsv.py

@@ -1,5 +1,7 @@
 """testing stancsv parsing"""
 
+import os
+from pathlib import Path
 from test import without_import
 from typing import List
 
@@ -9,6 +11,9 @@
 import cmdstanpy
 from cmdstanpy.utils import stancsv
 
+HERE = os.path.dirname(os.path.abspath(__file__))
+DATAFILES_PATH = os.path.join(HERE, 'data')
+
 
 def test_csv_bytes_to_numpy_no_header():
     lines = [
@@ -306,3 +311,23 @@ def test_csv_polars_and_numpy_equiv_one_element():
             lines, includes_header=False
         )
     assert np.array_equal(arr_out_polars, arr_out_numpy)
+
+
+def test_parse_stan_csv_from_file():
+    csv_path = os.path.join(DATAFILES_PATH, "bernoulli_output_1.csv")
+
+    comment_lines, draws_lines = stancsv.parse_stan_csv_comments_and_draws(
+        csv_path
+    )
+    assert all(ln.startswith(b"#") for ln in comment_lines)
+    assert all(not ln.startswith(b"#") for ln in draws_lines)
+
+    (
+        comment_lines_path,
+        draws_lines_path,
+    ) = stancsv.parse_stan_csv_comments_and_draws(Path(csv_path))
+    assert all(ln.startswith(b"#") for ln in comment_lines)
+    assert all(not ln.startswith(b"#") for ln in draws_lines)
+
+    assert comment_lines == comment_lines_path
+    assert draws_lines == draws_lines_path