✨ Add an artifact loader for .yaml (#2270)

lamindb/core/loaders.py

@@ -110,8 +110,23 @@ def load_json(path: UPathStr) -> dict:
     return data
 
 
+def load_yaml(path: UPathStr) -> dict | UPathStr:
+    """Load `.yaml` to `dict`."""
+    try:
+        import yaml  # type: ignore
+
+        with open(path) as f:
+            data = yaml.safe_load(f)
+        return data
+    except ImportError:
+        logger.warning(
+            "Please install PyYAML (`pip install PyYAML`) to load `.yaml` files."
+        )
+        return path
+
+
 def load_image(path: UPathStr) -> None | UPathStr:
-    """Display `.svg` in ipython, otherwise return path."""
+    """Display `.jpg`, `.gif` or `.png` in ipython, otherwise return path."""
     if is_run_from_ipython:
         from IPython.display import Image, display
 
@@ -147,7 +162,9 @@ def load_rds(path: UPathStr) -> UPathStr:
     ".zarr": load_anndata_zarr,
     ".html": load_html,
     ".json": load_json,
+    ".yaml": load_yaml,
     ".h5mu": load_h5mu,
+    ".gif": load_image,
     ".jpg": load_image,
     ".png": load_image,
     ".svg": load_svg,

noxfile.py

@@ -85,6 +85,8 @@ def install_ci(session, group):
     extras = ""
     if group == "unit-core":
         extras += "bionty,aws,gcp,zarr,fcs,jupyter"
+        # testing load_to_memory for yaml
+        run(session, "uv pip install --system PyYAML")
         run(session, "uv pip install --system huggingface_hub")
         # pinning 1.15.0rc3 because 1.14.5 is incompatible with anndata>=0.11.0
         # and >1.15.0rc4 has a different append mode API

tests/core/test_artifact.py

@@ -15,6 +15,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import yaml  # type: ignore
 from lamindb import _artifact
 from lamindb._artifact import (
     check_path_is_child_of_root,
@@ -98,6 +99,16 @@ def zip_file():
     filepath.unlink()
 
 
+@pytest.fixture(scope="module")
+def yaml_file():
+    filepath = Path("test.yaml")
+    dct = {"a": 1, "b": 2}
+    with open(filepath, "w") as f:
+        yaml.dump(dct, f)
+    yield filepath
+    filepath.unlink()
+
+
 @pytest.fixture(scope="module")
 def fcs_file():
     return ln.core.datasets.file_fcs()
@@ -719,7 +730,7 @@ def test_serialize_paths():
     assert isinstance(filepath, CloudPath)
 
 
-def test_load_to_memory(tsv_file, zip_file, fcs_file):
+def test_load_to_memory(tsv_file, zip_file, fcs_file, yaml_file):
     # tsv
     df = load_tsv(tsv_file)
     assert isinstance(df, pd.DataFrame)
@@ -730,6 +741,10 @@ def test_load_to_memory(tsv_file, zip_file, fcs_file):
     load_to_memory(zip_file)
     # check that it is a path
     assert isinstance(load_to_memory("./somefile.rds"), UPath)
+    # yaml
+    dct = load_to_memory(yaml_file)
+    assert dct["a"] == 1
+    assert dct["b"] == 2
 
     with pytest.raises(TypeError) as error:
         ln.Artifact(True)