From 3c44df5f3832b057bd1d0b53a41a67b36dc7e267 Mon Sep 17 00:00:00 2001
From: Louise Deconinck <louise.deconinck@gmail.com>
Date: Fri, 8 Nov 2024 10:07:01 +0100
Subject: [PATCH] Fixes #2 and generates empty uns if asked for (#3)

* Fixes #2 and generates empty uns if asked for

* Fix empty nested uns

* add Iterable import
---
 src/dummy_anndata/generate_dataset.py | 13 ++++++++++++-
 src/dummy_anndata/generate_dict.py    | 22 +++++++++++++++++-----
 tests/test_basic.py                   | 10 +++++++---
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/src/dummy_anndata/generate_dataset.py b/src/dummy_anndata/generate_dataset.py
index e4ebbc2..244068c 100644
--- a/src/dummy_anndata/generate_dataset.py
+++ b/src/dummy_anndata/generate_dataset.py
@@ -20,6 +20,7 @@ def generate_dataset(
     obsp_types=None,
     varp_types=None,
     uns_types=None,
+    nested_uns_types=None,
 ):
     """
     Generate a synthetic AnnData dataset with specified dimensions and data types.
@@ -48,6 +49,9 @@ def generate_dataset(
         Types of matrices to generate for `varp`. Each type must be a key in `matrix_generators`.
     uns_types : list of str, optional
         Types of data to generate for `uns`. Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators`.
+    nested_uns_types : list of str, optional
+        Types of data to generate for the nested `uns` dictionary. They will be a new dictionary at the key `nested`.
+        Each type must be a key in `vector_generators`, `matrix_generators`, or `scalar_generators`.
 
     Returns:
     --------
@@ -70,6 +74,7 @@ def generate_dataset(
     check_iterable_types(obsp_types, "obsp_types")
     check_iterable_types(varp_types, "varp_types")
     check_iterable_types(uns_types, "uns_types")
+    check_iterable_types(nested_uns_types, "nested_uns_types")
 
     assert layer_types is None or all(
         t in matrix_generators.keys() for t in layer_types
@@ -120,6 +125,12 @@ def generate_dataset(
             + list(matrix_generators.keys())
             + list(scalar_generators.keys())
         )
+    if nested_uns_types is None:
+        nested_uns_types = (
+            list(vector_generators.keys())
+            + list(matrix_generators.keys())
+            + list(scalar_generators.keys())
+        )
 
     X = matrix_generators[x_type](n_obs, n_vars)
     layers = {t: matrix_generators[t](n_obs, n_vars) for t in layer_types}
@@ -149,7 +160,7 @@ def generate_dataset(
     obsp = {t: matrix_generators[t](n_obs, n_obs) for t in obsp_types}
     varp = {t: matrix_generators[t](n_vars, n_vars) for t in varp_types}
 
-    uns = generate_dict(n_obs, n_vars, uns_types)
+    uns = generate_dict(n_obs, n_vars, uns_types, nested_uns_types)
 
     return ad.AnnData(
         X,
diff --git a/src/dummy_anndata/generate_dict.py b/src/dummy_anndata/generate_dict.py
index ca633a3..d2e7fa2 100644
--- a/src/dummy_anndata/generate_dict.py
+++ b/src/dummy_anndata/generate_dict.py
@@ -31,15 +31,27 @@ def generate_type(type, n_rows, n_cols):
     return None
 
 
-def generate_dict(n_rows, n_cols, types=None, nested=True):
+def generate_dict(n_rows, n_cols, types=None, nested_uns_types=None):
     if types is None:  # types are all vectors and all matrices
-        scalar_types = list(scalar_generators.keys()) + [f"scalar_{t}" for t in vector_generators.keys()]
-        types = scalar_types + list(vector_generators.keys()) + list(matrix_generators.keys())
+        types = (
+            list(scalar_generators.keys())
+            + [f"scalar_{t}" for t in vector_generators.keys()]
+            + list(vector_generators.keys())
+            + list(matrix_generators.keys())
+        )
+
+    if nested_uns_types is None:
+        nested_uns_types = (
+            list(scalar_generators.keys())
+            + [f"scalar_{t}" for t in vector_generators.keys()]
+            + list(vector_generators.keys())
+            + list(matrix_generators.keys())
+        )
 
     data = {}
     if types:  # types is not empty
         data = {t: generate_type(t, n_rows, n_cols) for t in types}
-        if nested:
-            data["nested"] = generate_dict(n_rows, n_cols, types, False)
+    if nested_uns_types:
+        data["nested"] = generate_dict(n_rows, n_cols, types = nested_uns_types, nested_uns_types=[])
 
     return data
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 1a82fb7..181c83f 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -12,8 +12,12 @@ def test_generating_dataset(tmp_path):
     filename = tmp_path / "dummy.h5ad"
     dummy.write_h5ad(filename)
 
+def test_uns():
+    dummy_empty = dummy_anndata.generate_dataset(uns_types=[], nested_uns_types=[])
+    assert dummy_empty.uns == {}
 
-def test_empty_uns():
-    dummy = dummy_anndata.generate_dataset(uns_types=[])
+    dummy_nested = dummy_anndata.generate_dataset(uns_types=[])
+    assert "nested" in dummy_nested.uns and dummy_nested.uns["nested"] != {}
 
-    assert dummy.uns == {}
+    dummy_no_nested = dummy_anndata.generate_dataset(nested_uns_types=[])
+    assert "nested" not in dummy_no_nested.uns