Repository: ImperialCollegeLondon/pycsvy

Issue:
Add support for `polars`
## Problem

[`polars`](https://pola.rs/) has become recently a very attractive - and faster - alternative to `pandas`, so we should support it, as well. 

## Solution

Largely copying what's in place for `pandas`, implement support for reading data into a `polars` dataframe with [`read_csv`](https://docs.pola.rs/api/python/stable/reference/api/polars.read_csv.html) and from a `polars` dataframe with [`write_csv`](https://docs.pola.rs/api/python/stable/reference/api/polars.DataFrame.write_csv.html). 

Discussion:

Patch:
diff --git a/csvy/__init__.py b/csvy/__init__.py
index 9955d89..b053cda 100644
--- a/csvy/__init__.py
+++ b/csvy/__init__.py
@@ -6,5 +6,6 @@ from .readers import (  # noqa: F401
     read_metadata,
     read_to_array,
     read_to_dataframe,
+    read_to_polars,
 )
 from .writers import Writer, write, write_header  # noqa: F401
diff --git a/csvy/readers.py b/csvy/readers.py
index ba9803e..7cec3f5 100644
--- a/csvy/readers.py
+++ b/csvy/readers.py
@@ -1,5 +1,7 @@
 """A collection of functions for parsing CSVY files."""
 
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 from typing import Any
@@ -19,7 +21,17 @@ try:
 except ModuleNotFoundError:
     DataFrame = None  # type: ignore
     logging.getLogger().debug(
-        "Pandas is not installed. Reading into a DataFrame will not work."
+        "Pandas is not installed. Reading into a pd.DataFrame will not work."
+    )
+
+try:
+    from polars import DataFrame as PolarsDataFrame
+    from polars import LazyFrame
+except ModuleNotFoundError:
+    LazyFrame = None  # type: ignore
+    PolarsDataFrame = None  # type: ignore
+    logging.getLogger().debug(
+        "Polars is not installed. Reading into a pl.DataFrame will not work."
     )
 
 
@@ -170,6 +182,54 @@ def read_to_dataframe(
     return pd.read_csv(filename, **options), header
 
 
+def read_to_polars(
+    filename: Path | str,
+    marker: str = "---",
+    csv_options: dict[str, Any] | None = None,
+    yaml_options: dict[str, Any] | None = None,
+    eager: bool = False,
+) -> tuple[LazyFrame | PolarsDataFrame, dict[str, Any]]:
+    """Reads a CSVY file into dict with the header and a Polars LazyFrame with the data.
+
+    This uses the `scan_csv` method from Polars to read the data. This returns a polars
+    LazyFrame, which means the data is not loaded into memory until it is needed. To
+    load the data into memory, set the `eager` parameter to `True`.
+
+    Possible 'skip_rows' and 'comment_prefix' argument provided in the 'csv_options'
+    dictionary will be ignored.
+
+    Args:
+        filename:  Name of the file to read.
+        marker: The marker characters that indicate the yaml header.
+        csv_options: Options to pass to pl.scan_csv.
+        yaml_options: Options to pass to yaml.safe_load.
+        eager: Whether to load the data into memory.
+
+    Raises:
+        ModuleNotFoundError: If polars is not found.
+
+    Returns:
+        Tuple containing: The polars LazyFrame and the header as a dictionary.
+    """
+    if LazyFrame is None:
+        raise ModuleNotFoundError(
+            "Module polars is not present. Install it to read data into DataFrame."
+        )
+    import polars as pl
+
+    yaml_options = yaml_options if yaml_options is not None else {}
+    header, nlines, comment = read_header(filename, marker=marker, **yaml_options)
+
+    options = csv_options.copy() if csv_options is not None else {}
+    options["skip_rows"] = nlines
+    options["comment_prefix"] = comment[0] if len(comment) >= 1 else None
+
+    lf = pl.scan_csv(filename, **options)
+    if eager:
+        return lf.collect(), header
+    return lf, header
+
+
 def read_to_list(
     filename: Path | str,
     marker: str = "---",
diff --git a/csvy/writers.py b/csvy/writers.py
index 60584ec..068971d 100644
--- a/csvy/writers.py
+++ b/csvy/writers.py
@@ -235,6 +235,42 @@ def write_pandas(
     return False
 
 
+@register_writer
+def write_polars(
+    filename: Path | str, data: Any, comment: str = "", **kwargs: Any
+) -> bool:
+    """Writes the polars dataframe to the chosen file, adding it after the header.
+
+    Args:
+        filename: Name of the file to save the data into. The data will be added to the
+            end of the file.
+        data: The data. If it is a polars DataFrame or LazyFrame, it will be saved,
+            otherwise nothing is done.
+        comment: String to use to mark the header lines as comments.
+        **kwargs: Arguments to be passed to the underlaying saving method.
+
+    Returns:
+        True if the writer worked, False otherwise.
+    """
+    try:
+        import polars as pl
+
+        if isinstance(data, pl.LazyFrame):
+            # Streaming mode (saving with `LazyFrame.sink_csv`) is unstable, so we
+            # collect the data into a DataFrame first
+            data = data.collect()
+        if isinstance(data, pl.DataFrame):
+            with open(filename, "a", newline="") as f:
+                data.write_csv(f, **kwargs)
+
+            return True
+
+    except ModuleNotFoundError:
+        logging.getLogger().debug("Polars is not installed, so not using 'write_csv'.")
+
+    return False
+
+
 def write_csv(
     filename: Path | str, data: Any, comment: str = "", **kwargs: Any
 ) -> bool:
diff --git a/poetry.lock b/poetry.lock
index 4b7bdb1..702ae28 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -323,13 +323,13 @@ dev = ["flake8", "markdown", "twine", "wheel"]
 
 [[package]]
 name = "griffe"
-version = "1.3.1"
+version = "1.3.2"
 description = "Signatures for entire Python programs. Extract the structure, the frame, the skeleton of your project, to generate API documentation or find breaking changes in your API."
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "griffe-1.3.1-py3-none-any.whl", hash = "sha256:940aeb630bc3054b4369567f150b6365be6f11eef46b0ed8623aea96e6d17b19"},
-    {file = "griffe-1.3.1.tar.gz", hash = "sha256:3f86a716b631a4c0f96a43cb75d05d3c85975003c20540426c0eba3b0581c56a"},
+    {file = "griffe-1.3.2-py3-none-any.whl", hash = "sha256:2e34b5e46507d615915c8e6288bb1a2234bd35dee44d01e40a2bc2f25bd4d10c"},
+    {file = "griffe-1.3.2.tar.gz", hash = "sha256:1ec50335aa507ed2445f2dd45a15c9fa3a45f52c9527e880571dfc61912fd60c"},
 ]
 
 [package.dependencies]
@@ -942,6 +942,47 @@ files = [
 dev = ["pre-commit", "tox"]
 testing = ["pytest", "pytest-benchmark"]
 
+[[package]]
+name = "polars"
+version = "1.8.2"
+description = "Blazingly fast DataFrame library"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "polars-1.8.2-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:114be1ebfb051b794fb9e1f15999430c79cc0824595e237d3f45632be3e56d73"},
+    {file = "polars-1.8.2-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:e4fc36cfe48972d4c5be21a7cb119d6378fb7af0bb3eeb61456b66a1f43228e3"},
+    {file = "polars-1.8.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:67c1e448d6e38697650b22dd359f13c40b567c0b66686c8602e4367400e87801"},
+    {file = "polars-1.8.2-cp38-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:570ee86b033dc5a6dbe2cb0df48522301642f304dda3da48f53d7488899a2206"},
+    {file = "polars-1.8.2-cp38-abi3-win_amd64.whl", hash = "sha256:ce1a1c1e2150ffcc44a5f1c461d738e1dcd95abbd0f210af0271c7ac0c9f7ef9"},
+    {file = "polars-1.8.2.tar.gz", hash = "sha256:42f69277d5be2833b0b826af5e75dcf430222d65c9633872856e176a0bed27a0"},
+]
+
+[package.extras]
+adbc = ["adbc-driver-manager[dbapi]", "adbc-driver-sqlite[dbapi]"]
+all = ["polars[async,cloudpickle,database,deltalake,excel,fsspec,graph,iceberg,numpy,pandas,plot,pyarrow,pydantic,style,timezone]"]
+async = ["gevent"]
+calamine = ["fastexcel (>=0.9)"]
+cloudpickle = ["cloudpickle"]
+connectorx = ["connectorx (>=0.3.2)"]
+database = ["nest-asyncio", "polars[adbc,connectorx,sqlalchemy]"]
+deltalake = ["deltalake (>=0.15.0)"]
+excel = ["polars[calamine,openpyxl,xlsx2csv,xlsxwriter]"]
+fsspec = ["fsspec"]
+gpu = ["cudf-polars-cu12"]
+graph = ["matplotlib"]
+iceberg = ["pyiceberg (>=0.5.0)"]
+numpy = ["numpy (>=1.16.0)"]
+openpyxl = ["openpyxl (>=3.0.0)"]
+pandas = ["pandas", "polars[pyarrow]"]
+plot = ["altair (>=5.4.0)"]
+pyarrow = ["pyarrow (>=7.0.0)"]
+pydantic = ["pydantic"]
+sqlalchemy = ["polars[pandas]", "sqlalchemy"]
+style = ["great-tables (>=0.8.0)"]
+timezone = ["backports-zoneinfo", "tzdata"]
+xlsx2csv = ["xlsx2csv (>=0.8.0)"]
+xlsxwriter = ["xlsxwriter"]
+
 [[package]]
 name = "pre-commit"
 version = "2.21.0"
@@ -1374,4 +1415,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "b0446d787719bdd9d5e2eb5c7589cee89dda137be65b03d969a9c8640313bcfe"
+content-hash = "13364fb2c65525dd9c7718d630b55e5fb263798b47947fba272fdeca42105460"
diff --git a/pyproject.toml b/pyproject.toml
index 9925993..5d1ee07 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,9 +19,7 @@ classifiers = [
 	"Operating System :: OS Independent",
 ]
 
-packages = [
-    { include = "csvy" },
-]
+packages = [{ include = "csvy" }]
 
 [tool.poetry.dependencies]
 python = "^3.9"
@@ -41,6 +39,7 @@ pandas = "^2.0.0"
 types-PyYAML = "^6.0.7"
 bump2version = "^1.0.1"
 coverage = "^7.1.0"
+polars = "^1.8.2"
 ruff = "^0.6.8"
 
 [tool.poetry.group.docs]
@@ -48,7 +47,7 @@ optional = true
 
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.4.1"
-mkdocstrings = {extras = ["python"], version = "^0.19.0"}
+mkdocstrings = { extras = ["python"], version = "^0.19.0" }
 mkdocs-material = "^8.5.7"
 mkdocs-gen-files = "^0.4.0"
 mkdocs-literate-nav = "^0.5.0"
@@ -57,12 +56,12 @@ mkdocs-section-index = "^0.3.4"
 [tool.ruff]
 target-version = "py39"
 lint.select = [
-    "D",   # pydocstyle
-    "E",   # pycodestyle
-    "F",   # Pyflakes
-    "I",   # isort
-    "UP",  # pyupgrade
-    "RUF", # ruff
+	"D",   # pydocstyle
+	"E",   # pycodestyle
+	"F",   # Pyflakes
+	"I",   # isort
+	"UP",  # pyupgrade
+	"RUF", # ruff
 ]
 lint.pydocstyle.convention = "google"
 


Test patch:
diff --git a/tests/test_read.py b/tests/test_read.py
index ed95ca7..9d2041f 100644
--- a/tests/test_read.py
+++ b/tests/test_read.py
@@ -91,6 +91,30 @@ def test_read_to_dataframe(data_path):
         read_to_dataframe(data_path)
 
 
+def test_read_to_polars(data_path):
+    """Test the read_to_polars function."""
+    import polars as pl
+    from polars.testing import assert_frame_equal
+
+    from csvy.readers import read_to_polars
+
+    lazy_data, header = read_to_polars(data_path)
+    assert isinstance(lazy_data, pl.LazyFrame)
+    assert tuple(lazy_data.columns) == ("Date", "WTI")
+    assert isinstance(header, dict)
+    assert len(header) > 0
+
+    eager_data, _ = read_to_polars(data_path, eager=True)
+    assert_frame_equal(lazy_data.collect(), eager_data)
+
+    import csvy.readers as readers
+
+    readers.LazyFrame = None
+
+    with pytest.raises(ModuleNotFoundError):
+        read_to_polars(data_path)
+
+
 def test_read_to_list(array_data_path):
     """Test the read_to_list function."""
     from csvy.readers import read_to_list
diff --git a/tests/test_write.py b/tests/test_write.py
index 365104a..08962bf 100644
--- a/tests/test_write.py
+++ b/tests/test_write.py
@@ -69,7 +69,7 @@ def test_write_numpy(mock_save, tmpdir):
 
 @patch("pandas.DataFrame.to_csv")
 def test_write_pandas(mock_save, tmpdir):
-    """Test the save_header function."""
+    """Test the write_pandas function."""
     import pandas as pd
 
     from csvy.writers import write_pandas
@@ -84,6 +84,30 @@ def test_write_pandas(mock_save, tmpdir):
     mock_save.assert_called_once()
 
 
+@patch("polars.DataFrame.write_csv")
+def test_write_polars(mock_save, tmpdir, mocker):
+    """Test the write_polars function."""
+    import polars as pl
+
+    from csvy.writers import write_polars
+
+    filename = tmpdir / "some_file.csv"
+
+    data = []
+    assert not write_polars(filename, data)
+
+    data = pl.DataFrame()
+    assert write_polars(filename, data)
+    mock_save.assert_called_once()
+
+    data = pl.LazyFrame()
+    collect_spy = mocker.spy(data, "collect")
+    mock_save.reset_mock()
+    assert write_polars(filename, data)
+    collect_spy.assert_called_once()
+    mock_save.assert_called_once()
+
+
 @patch("csv.writer")
 def test_write_csv(mock_save, tmpdir):
     """Test the write_csv function."""