Make polars an optional dependency

Author: amas0

cmdstanpy/utils/stancsv.py

@@ -7,6 +7,7 @@
 import math
 import os
 import re
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
 from typing import (
@@ -26,7 +27,6 @@
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
-import polars as pl
 
 from cmdstanpy import _CMDSTAN_SAMPLING, _CMDSTAN_THIN, _CMDSTAN_WARMUP
 
@@ -152,13 +152,34 @@ def csv_bytes_list_to_numpy(
     """Efficiently converts a list of bytes representing whose concatenation
     represents a CSV file into a numpy array. Includes header specifies
     whether the bytes contains an initial header line."""
-    out = (
-        pl.read_csv(
-            io.BytesIO(b"".join(csv_bytes_list)), has_header=includes_header
-        )
-        .to_numpy()
-        .astype(np.float32)
-    )
+    try:
+        import polars as pl
+
+        try:
+            out = (
+                pl.read_csv(
+                    io.BytesIO(b"".join(csv_bytes_list)),
+                    has_header=includes_header,
+                )
+                .to_numpy()
+                .astype(np.float32)
+            )
+            if out.shape[0] == 0:
+                raise ValueError("No data found to parse")
+        except pl.exceptions.NoDataError as exc:
+            raise ValueError("No data found to parse") from exc
+    except ImportError as exc:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore")
+            out = np.loadtxt(
+                csv_bytes_list,
+                skiprows=int(includes_header),
+                delimiter=",",
+                dtype=np.float32,
+            )
+        if out.shape == (0,):
+            raise ValueError("No data found to parse") from exc
+
     # Telling the type checker we know the type is correct
     return cast(npt.NDArray[np.float32], out)
 

pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 license = { text = "BSD-3-Clause" }
 authors = [{ name = "Stan Dev Team" }]
 requires-python = ">=3.8"
-dependencies = ["pandas", "numpy>=1.21", "tqdm", "stanio>=0.4.0,<2.0.0", "polars>=1.8.2"]
+dependencies = ["pandas", "numpy>=1.21", "tqdm", "stanio>=0.4.0,<2.0.0"]
 dynamic = ["version"]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -40,7 +40,7 @@ packages = ["cmdstanpy", "cmdstanpy.stanfit", "cmdstanpy.utils"]
 "cmdstanpy" = ["py.typed"]
 
 [project.optional-dependencies]
-all = ["xarray"]
+all = ["xarray", "polars>=1.8.2"]
 test = [
     "flake8",
     "pylint",

test/test_stancsv.py

@@ -1,9 +1,9 @@
 """testing stancsv parsing"""
 
 from typing import List
+from unittest import mock
 
 import numpy as np
-import polars as pl
 import pytest
 
 from cmdstanpy.utils import stancsv
@@ -30,6 +30,28 @@ def test_csv_bytes_to_numpy_no_header():
     assert arr_out[0].dtype == np.float32
 
 
+def test_csv_bytes_to_numpy_no_header_no_polars():
+    lines = [
+        b"-6.76206,1,0.787025,1,1,0,6.81411,0.229458\n",
+        b"-6.81411,0.983499,0.787025,1,1,0,6.8147,0.20649\n",
+        b"-6.85511,0.994945,0.787025,2,3,0,6.85536,0.310589\n",
+        b"-6.85511,0.812189,0.787025,1,1,0,7.16517,0.310589\n",
+    ]
+    expected = np.array(
+        [
+            [-6.76206, 1, 0.787025, 1, 1, 0, 6.81411, 0.229458],
+            [-6.81411, 0.983499, 0.787025, 1, 1, 0, 6.8147, 0.20649],
+            [-6.85511, 0.994945, 0.787025, 2, 3, 0, 6.85536, 0.310589],
+            [-6.85511, 0.812189, 0.787025, 1, 1, 0, 7.16517, 0.310589],
+        ],
+        dtype=np.float32,
+    )
+    with mock.patch.dict("sys.modules", {"polars": None}):
+        arr_out = stancsv.csv_bytes_list_to_numpy(lines, includes_header=False)
+        assert np.array_equiv(arr_out, expected)
+        assert arr_out[0].dtype == np.float32
+
+
 def test_csv_bytes_to_numpy_with_header():
     lines = [
         (
@@ -54,21 +76,65 @@ def test_csv_bytes_to_numpy_with_header():
     assert np.array_equiv(arr_out, expected)
 
 
+def test_csv_bytes_to_numpy_with_header_no_polars():
+    lines = [
+        (
+            b"lp__,accept_stat__,stepsize__,treedepth__,"
+            b"n_leapfrog__,divergent__,energy__,theta\n"
+        ),
+        b"-6.76206,1,0.787025,1,1,0,6.81411,0.229458\n",
+        b"-6.81411,0.983499,0.787025,1,1,0,6.8147,0.20649\n",
+        b"-6.85511,0.994945,0.787025,2,3,0,6.85536,0.310589\n",
+        b"-6.85511,0.812189,0.787025,1,1,0,7.16517,0.310589\n",
+    ]
+    expected = np.array(
+        [
+            [-6.76206, 1, 0.787025, 1, 1, 0, 6.81411, 0.229458],
+            [-6.81411, 0.983499, 0.787025, 1, 1, 0, 6.8147, 0.20649],
+            [-6.85511, 0.994945, 0.787025, 2, 3, 0, 6.85536, 0.310589],
+            [-6.85511, 0.812189, 0.787025, 1, 1, 0, 7.16517, 0.310589],
+        ],
+        dtype=np.float32,
+    )
+    with mock.patch.dict("sys.modules", {"polars": None}):
+        arr_out = stancsv.csv_bytes_list_to_numpy(lines, includes_header=True)
+        assert np.array_equiv(arr_out, expected)
+
+
 def test_csv_bytes_to_numpy_empty():
     lines = [b""]
-    with pytest.raises(pl.exceptions.NoDataError):
+    with pytest.raises(ValueError):
         stancsv.csv_bytes_list_to_numpy(lines)
 
 
+def test_csv_bytes_to_numpy_empty_no_polars():
+    lines = [b""]
+    with mock.patch.dict("sys.modules", {"polars": None}):
+        with pytest.raises(ValueError):
+            stancsv.csv_bytes_list_to_numpy(lines)
+
+
 def test_csv_bytes_to_numpy_header_no_draws():
     lines = [
         (
             b"lp__,accept_stat__,stepsize__,treedepth__,"
             b"n_leapfrog__,divergent__,energy__,theta\n"
         ),
     ]
-    arr_out = stancsv.csv_bytes_list_to_numpy(lines)
-    assert arr_out.shape == (0, 8)
+    with pytest.raises(ValueError):
+        stancsv.csv_bytes_list_to_numpy(lines)
+
+
+def test_csv_bytes_to_numpy_header_no_draws_no_polars():
+    lines = [
+        (
+            b"lp__,accept_stat__,stepsize__,treedepth__,"
+            b"n_leapfrog__,divergent__,energy__,theta\n"
+        ),
+    ]
+    with mock.patch.dict("sys.modules", {"polars": None}):
+        with pytest.raises(ValueError):
+            stancsv.csv_bytes_list_to_numpy(lines)
 
 
 def test_parsing_with_rules():