From ae2363c3824e3efad5357634f211f7a549053d43 Mon Sep 17 00:00:00 2001
From: Sergei Rybakov <koncopd@gmail.com>
Date: Tue, 10 Dec 2024 14:01:58 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20an=20artifact=20loader=20for?=
 =?UTF-8?q?=20.yaml=20(#2270)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 lamindb/core/loaders.py     | 19 ++++++++++++++++++-
 noxfile.py                  |  2 ++
 tests/core/test_artifact.py | 17 ++++++++++++++++-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/lamindb/core/loaders.py b/lamindb/core/loaders.py
index aa9f1940f..409b2b9f8 100644
--- a/lamindb/core/loaders.py
+++ b/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,
diff --git a/noxfile.py b/noxfile.py
index de0280da8..12900e056 100644
--- a/noxfile.py
+++ b/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
diff --git a/tests/core/test_artifact.py b/tests/core/test_artifact.py
index 14ef21cc8..7a1ae16ad 100644
--- a/tests/core/test_artifact.py
+++ b/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)