Add support for polars

.github/workflows/tests.yml

@@ -44,8 +44,14 @@ jobs:
           python -m pip install --upgrade pip
           pip install '.[dev,examples]'
 
-      - name: Test
+      - name: Test all except polars
         run: pytest
 
+      - name: Install polars
+        run: pip install '.[polars]'
+
+      - name: Test only polars
+        run: pytest -m polars
+
       - name: Examples
         run: python examples/check_timesteps.py tests/data

CHANGELOG.md

@@ -6,6 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+- Add support for polars ([#44](https://github.com/gadomski/pyisd/pull/44))
+
 ## [0.3.0] - 2024-02-02
 
 ### Added

pyproject.toml

@@ -20,6 +20,12 @@ isd = "isd.cli:main"
 dev = ["mypy~=1.2", "pre-commit~=3.2", "pytest~=8.0", "ruff~=0.1.9"]
 examples = ["tqdm~=4.66"]
 docs = ["sphinx~=7.2"]
+polars = ["polars>=0.20.6"]
+
+[tool.pytest.ini_options]
+markers = [
+    "polars: marks tests that require polars",
+]
 
 [build-system]
 requires = ["setuptools>=45", "wheel"]

src/isd/batch.py

@@ -4,13 +4,16 @@
 from io import BytesIO
 from pathlib import Path
 from dataclasses import dataclass
-from typing import List, Union, Optional, Dict, Any, Iterator
+from typing import List, Union, Optional, Dict, Any, Iterator, TYPE_CHECKING
 import datetime
 
 from isd.record import Record
 
 import pandas
 
+if TYPE_CHECKING:
+    import polars
+
 
 @dataclass
 class Batch:
@@ -84,3 +87,9 @@ def to_json(self, indent: int = 4) -> str:
     def to_data_frame(self) -> pandas.DataFrame:
         """Reads a local ISD file into a DataFrame."""
         return pandas.DataFrame([record.to_dict() for record in self.records])
+
+    def to_polars(self) -> "polars.DataFrame":
+        """Reads a local ISD file into a Polars DataFrame."""
+        import polars
+
+        return polars.DataFrame([record.to_dict() for record in self.records])

tests/test_batch.py

@@ -1,6 +1,8 @@
 import datetime
 import json
 
+import pytest
+
 from isd import Batch
 
 
@@ -66,3 +68,13 @@ def test_batch_to_df(batch: Batch) -> None:
     df = batch.to_data_frame()
     df = df[df["datetime"] >= datetime_min]
     assert len(df) == 212
+
+
+@pytest.mark.polars  # type: ignore
+def test_batch_to_polars(batch: Batch) -> None:
+    polars = pytest.importorskip("polars")
+
+    datetime_min = datetime.datetime(2021, 1, 5)
+    df = batch.to_polars()
+    df = df.filter(polars.col("datetime") >= datetime_min)
+    assert len(df) == 212