From dc7585aabffc15ef9814121af249213b144f4dfc Mon Sep 17 00:00:00 2001
From: Zach Kurtz <zkurtz@protonmail.com>
Date: Tue, 10 Dec 2024 22:02:16 -0500
Subject: [PATCH] support dill/pickle

---
 dummio/__init__.py                   |  7 +++++++
 dummio/dill.py                       | 19 +++++++++++++++++++
 dummio/pickle.py                     | 18 ++++++++++++++++++
 pyproject.toml                       |  3 ++-
 tests/test_all_basic.py              | 18 ++++++++++++++++++
 tests/test_assert_module_protocol.py |  2 ++
 uv.lock                              | 13 ++++++++++++-
 7 files changed, 78 insertions(+), 2 deletions(-)
 create mode 100644 dummio/dill.py
 create mode 100644 dummio/pickle.py

diff --git a/dummio/__init__.py b/dummio/__init__.py
index 33aafed..2917bc3 100644
--- a/dummio/__init__.py
+++ b/dummio/__init__.py
@@ -4,6 +4,7 @@
 
 
 from dummio import json as json
+from dummio import pickle as pickle
 from dummio import text as text
 
 try:
@@ -23,3 +24,9 @@
 except ImportError:
     # this would require the optional dependency pydantic
     pass
+
+try:
+    from dummio import dill as dill
+except ImportError:
+    # this would require the optional dependency dill
+    pass
diff --git a/dummio/dill.py b/dummio/dill.py
new file mode 100644
index 0000000..3a9398b
--- /dev/null
+++ b/dummio/dill.py
@@ -0,0 +1,19 @@
+"""IO for dill."""
+
+from typing import Any
+
+import dill
+
+from dummio.constants import PathType
+
+
+def save(data: Any, *, filepath: PathType) -> None:
+    """Save a dill file."""
+    with open(filepath, "wb") as file:
+        dill.dump(data, file)
+
+
+def load(filepath: PathType) -> Any:
+    """Read a dill file."""
+    with open(filepath, "rb") as file:
+        return dill.load(file)
diff --git a/dummio/pickle.py b/dummio/pickle.py
new file mode 100644
index 0000000..8d42a73
--- /dev/null
+++ b/dummio/pickle.py
@@ -0,0 +1,18 @@
+"""IO for pickle."""
+
+import pickle
+from typing import Any
+
+from dummio.constants import PathType
+
+
+def save(data: Any, *, filepath: PathType) -> None:
+    """Save a pickle file."""
+    with open(filepath, "wb") as file:
+        pickle.dump(data, file)
+
+
+def load(filepath: PathType) -> Any:
+    """Read a pickle file."""
+    with open(filepath, "rb") as file:
+        return pickle.load(file)
diff --git a/pyproject.toml b/pyproject.toml
index 060d49d..8cbd68e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "dummio"
-version = "1.3.0"
+version = "1.4.0"
 description = "Easiest-possible IO for basic file types."
 authors = [{ name = "Zach Kurtz", email = "zkurtz@gmail.com" }]
 readme = "README.md"
@@ -21,6 +21,7 @@ extras = [
   "pyyaml>=6.0.2",
   "onnx>=1.10.1",
   "pydantic>=2.10.2",
+  "dill>=0.2.2",
 ]
 
 [project.urls]
diff --git a/tests/test_all_basic.py b/tests/test_all_basic.py
index ff3ff52..3a06407 100644
--- a/tests/test_all_basic.py
+++ b/tests/test_all_basic.py
@@ -50,3 +50,21 @@ def test_yaml(tmp_path: Path) -> None:
         data=dictionary(),
         module=dummio.yaml,
     )
+
+
+def test_pickle(tmp_path: Path) -> None:
+    """Test the packio package."""
+    _assert_cycle(
+        path=tmp_path / "data.pkl",
+        data=dictionary(),
+        module=dummio.pickle,
+    )
+
+
+def test_dill(tmp_path: Path) -> None:
+    """Test the packio package."""
+    _assert_cycle(
+        path=tmp_path / "data.pkl",
+        data=dictionary(),
+        module=dummio.dill,
+    )
diff --git a/tests/test_assert_module_protocol.py b/tests/test_assert_module_protocol.py
index 35d306d..5c08372 100644
--- a/tests/test_assert_module_protocol.py
+++ b/tests/test_assert_module_protocol.py
@@ -7,8 +7,10 @@
 from dummio.protocol import assert_module_protocol
 
 IO_MODULES = [
+    "dummio.dill",
     "dummio.json",
     "dummio.onnx",
+    "dummio.pickle",
     "dummio.pydantic",
     "dummio.text",
     "dummio.yaml",
diff --git a/uv.lock b/uv.lock
index 9bb5108..abb486e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -102,6 +102,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ea/72/a39471976515f028d3c8823d2a0a43931813251018dea889949d80196fcb/cramjam-2.9.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:65bded20fd2cef17b22246c336ddd67fac842341ee311042b4a70e65dc745aa7", size = 2109333 },
 ]
 
+[[package]]
+name = "dill"
+version = "0.3.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/70/43/86fe3f9e130c4137b0f1b50784dd70a5087b911fe07fa81e53e0c4c47fea/dill-0.3.9.tar.gz", hash = "sha256:81aa267dddf68cbfe8029c42ca9ec6a4ab3b22371d1c450abc54422577b4512c", size = 187000 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/46/d1/e73b6ad76f0b1fb7f23c35c6d95dbc506a9c8804f43dda8cb5b0fa6331fd/dill-0.3.9-py3-none-any.whl", hash = "sha256:468dff3b89520b474c0397703366b7b95eebe6303f108adf9b19da1f702be87a", size = 119418 },
+]
+
 [[package]]
 name = "distlib"
 version = "0.3.9"
@@ -113,7 +122,7 @@ wheels = [
 
 [[package]]
 name = "dummio"
-version = "1.2.0"
+version = "1.3.0"
 source = { editable = "." }
 
 [package.dev-dependencies]
@@ -126,6 +135,7 @@ dev = [
     { name = "skl2onnx" },
 ]
 extras = [
+    { name = "dill" },
     { name = "fastparquet" },
     { name = "onnx" },
     { name = "pandas" },
@@ -145,6 +155,7 @@ dev = [
     { name = "skl2onnx", specifier = ">=1.10.1" },
 ]
 extras = [
+    { name = "dill", specifier = ">=0.2.2" },
     { name = "fastparquet", specifier = ">=2024.11.0" },
     { name = "onnx", specifier = ">=1.10.1" },
     { name = "pandas", specifier = ">=1.5.0" },