feat: csv and tsv support (#33)

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "cryptography>=41.0.7",
     "numpy>=1.21.1",
     "soundfile>=0.12.1",
+    "pandas>=1.5.3",
 ]
 
 [tool.setuptools.dynamic]

src/iokit/__init__.py

@@ -1,5 +1,6 @@
 __version__ = "0.2.0"
 __all__ = [
+    "Csv",
     "Dat",
     "decrypt",
     "download_file",
@@ -21,6 +22,7 @@
     "SecretState",
     "State",
     "Tar",
+    "Tsv",
     "Txt",
     "Waveform",
     "Wav",
@@ -29,6 +31,7 @@
 ]
 
 from .extensions import (
+    Csv,
     Dat,
     Enc,
     Env,
@@ -41,6 +44,7 @@
     Ogg,
     SecretState,
     Tar,
+    Tsv,
     Txt,
     Wav,
     Waveform,

src/iokit/extensions/__init__.py

@@ -1,4 +1,5 @@
 __all__ = [
+    "Csv",
     "Dat",
     "Enc",
     "Env",
@@ -11,6 +12,7 @@
     "Ogg",
     "SecretState",
     "Tar",
+    "Tsv",
     "Txt",
     "Wav",
     "Waveform",
@@ -28,6 +30,7 @@
 from .json import Json
 from .jsonl import Jsonl
 from .npy import Npy
+from .table import Csv, Tsv
 from .tar import Tar
 from .txt import Txt
 from .yaml import Yaml

src/iokit/extensions/table.py

@@ -0,0 +1,28 @@
+__all__ = ["Csv", "Tsv"]
+
+from io import BytesIO
+from typing import Any
+
+from pandas import DataFrame, read_csv
+
+from iokit.state import State
+
+
+class Csv(State, suffix="json"):
+    def __init__(self, frame: DataFrame, *, index: bool = False, **kwargs: Any):
+        with BytesIO() as buffer:
+            frame.to_csv(buffer, index=index)
+            super().__init__(buffer.getvalue(), **kwargs)
+
+    def load(self) -> DataFrame:
+        return read_csv(self.buffer)
+
+
+class Tsv(State, suffix="tsv"):
+    def __init__(self, frame: DataFrame, *, index: bool = False, **kwargs: Any):
+        with BytesIO() as buffer:
+            frame.to_csv(buffer, sep="\t", index=index)
+            super().__init__(buffer.getvalue(), **kwargs)
+
+    def load(self) -> DataFrame:
+        return read_csv(self.buffer, sep="\t")

tests/test_table.py

@@ -0,0 +1,27 @@
+import pandas as pd
+
+from iokit import Csv, Tsv
+
+
+def _test_teble_data() -> list[dict[str, str | int]]:
+    return [
+        {"name": "Alice", "age": 24},
+        {"name": "Bob", "age": 25},
+        {"name": "Charlie", "age": 26},
+    ]
+
+
+def test_csv() -> None:
+    frame = pd.DataFrame(_test_teble_data())
+    assert frame.equals(Csv(frame).load())
+    assert Csv(frame, index=True).size > Csv(frame).size
+    assert Csv(frame).load().equals(Csv(frame, index=True).load()[["name", "age"]])
+    assert len([1 for ch in Csv(frame).data if ch == b","[0]]) == 4
+
+
+def test_tsv() -> None:
+    frame = pd.DataFrame(_test_teble_data())
+    assert frame.equals(Tsv(frame).load())
+    assert Tsv(frame, index=True).size > Tsv(frame).size
+    assert Tsv(frame).load().equals(Tsv(frame, index=True).load()[["name", "age"]])
+    assert len([1 for ch in Tsv(frame).data if ch == b"\t"[0]]) == 4