Port read and write functions

src/finch/__init__.py

@@ -130,6 +130,10 @@
     iinfo,
     can_cast,
 )
+from .io import (
+    read,
+    write,
+)
 
 __all__ = [
     "Tensor",
@@ -252,6 +256,8 @@
     "real",
     "imag",
     "conj",
+    "read",
+    "write",
 ]
 
 __array_api_version__: str = "2023.12"

src/finch/io.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+from .julia import add_package, jl
+from .tensor import Tensor
+
+
+def _import_deps(filename: str) -> None:
+    fn = filename
+    if fn.endswith(".mtx") or fn.endswith(".ttx") or fn.endswith(".tns"):
+        add_package("TensorMarket", hash="8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77", version="0.2.0")
+        jl.seval("using TensorMarket")
+    elif fn.endswith(".bspnpy"):
+        add_package("NPZ", hash="15e1cf62-19b3-5cfa-8e77-841668bca605", version="0.4.3")
+        jl.seval("using NPZ")
+    elif fn.endswith(".bsp.h5"):
+        add_package("HDF5", hash="f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f", version="0.17.2")
+        jl.seval("using HDF5")
+    else:
+        raise ValueError(
+            f"Unsupported file extension. Supported extensions are "
+            "`.mtx`, `.ttx`, `.tns`, `.bspnpy`, and `.bsp.h5`."
+        )
+
+
+def read(filename: Path | str) -> Tensor:
+    fn = str(filename)
+    _import_deps(fn)
+    julia_obj = jl.fread(fn)
+    return Tensor(julia_obj)
+
+
+def write(filename: Path | str, tns: Tensor) -> None:
+    fn = str(filename)
+    _import_deps(fn)
+    jl.fwrite(fn, tns._obj)

src/finch/julia.py

@@ -2,6 +2,14 @@
 
 import juliapkg
 
+
+def add_package(name: str, hash: str, version: str) -> None:
+    deps = juliapkg.deps.load_cur_deps()
+    if (deps.get("packages", {}).get(name, {}).get("version", None) != version):
+        juliapkg.add(name, hash, version=version)
+        juliapkg.resolve()
+
+
 _FINCH_NAME = "Finch"
 _FINCH_VERSION = "0.6.31"
 _FINCH_HASH = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce"
@@ -16,12 +24,7 @@
 elif _FINCH_REPO_URL:
     juliapkg.add(_FINCH_NAME, _FINCH_HASH, url=_FINCH_REPO_URL, dev=True)
 else:
-    deps = juliapkg.deps.load_cur_deps()
-    if (
-        deps.get("packages", {}).get(_FINCH_NAME, {}).get("version", None)
-        != _FINCH_VERSION
-    ):
-        juliapkg.add(_FINCH_NAME, _FINCH_HASH, version=_FINCH_VERSION)
+    add_package(_FINCH_NAME, _FINCH_HASH, _FINCH_VERSION)
 
 import juliacall as jc  # noqa
 

tests/data/matrix_1.ttx

@@ -0,0 +1,17 @@
+%%MatrixMarket matrix coordinate integer general
+3 5 15
+1 1 0
+2 1 1
+3 1 0
+1 2 0
+2 2 0
+3 2 5
+1 3 3
+2 3 0
+3 3 0
+1 4 2
+2 4 1
+3 4 0
+1 5 0
+2 5 0
+3 5 0

tests/test_io.py

@@ -0,0 +1,23 @@
+import os
+
+from numpy.testing import assert_equal
+
+import finch
+
+base_path = "tests/data"
+
+
+def test_read(arr2d):
+    tns = finch.read(f"{base_path}/matrix_1.ttx")
+
+    assert_equal(tns.todense(), arr2d)
+
+
+def test_write(tmp_path, arr2d):
+    tns = finch.asarray(arr2d)
+    finch.write(tmp_path / "tmp.ttx", tns)
+
+    expected = open(f"{base_path}/matrix_1.ttx").read()
+    actual = open(tmp_path / "tmp.ttx").read()
+
+    assert actual == expected