Repository: lincc-frameworks/nested-pandas (language: python)

Issue:
Add a public function to generate test data
**Feature request**

It would be useful to create a `datasets` module with some simple datasets built-in. We may start with a really simple function to generate a simple test time-domain catalog, like we have in `tape`'s tests.

**Before submitting**
Please check the following:

- [x] I have described the purpose of the suggested change, specifying what I need the enhancement to accomplish, i.e. what problem it solves.
- [x] I have included any relevant links, screenshots, environment information, and data relevant to implementing the requested feature, as well as pseudocode for how I want to access the new functionality.
- [x] If I have ideas for how the new feature could be implemented, I have provided explanations and/or pseudocode and/or task lists for the steps.



Discussion / hints:

Fix patch:
diff --git a/src/nested_pandas/datasets/__init__.py b/src/nested_pandas/datasets/__init__.py
new file mode 100644
index 0000000..4b8e827
--- /dev/null
+++ b/src/nested_pandas/datasets/__init__.py
@@ -0,0 +1,1 @@
+from .generation import *  # noqa
diff --git a/src/nested_pandas/datasets/generation.py b/src/nested_pandas/datasets/generation.py
new file mode 100644
index 0000000..f50bfe1
--- /dev/null
+++ b/src/nested_pandas/datasets/generation.py
@@ -0,0 +1,55 @@
+import numpy as np
+
+from nested_pandas import NestedFrame
+
+
+def generate_data(n_base, n_layer, seed=None) -> NestedFrame:
+    """Generates a toy dataset.
+
+    Parameters
+    ----------
+    n_base : int
+        The number of rows to generate for the base layer
+    n_layer : int, or dict
+        The number of rows per n_base row to generate for a nested layer.
+        Alternatively, a dictionary of layer label, layer_size pairs may be
+        specified to created multiple nested columns with custom sizing.
+    seed : int
+        A seed to use for random generation of data
+
+    Returns
+    -------
+    NestedFrame
+        The constructed NestedFrame.
+
+    Examples
+    --------
+    >>> nested_pandas.datasets.generate_data(10,100)
+    >>> nested_pandas.datasets.generate_data(10, {"nested_a": 100, "nested_b": 200})
+    """
+    # use provided seed, "None" acts as if no seed is provided
+    randomstate = np.random.RandomState(seed=seed)
+
+    # Generate base data
+    base_data = {"a": randomstate.random(n_base), "b": randomstate.random(n_base) * 2}
+    base_nf = NestedFrame(data=base_data)
+
+    # In case of int, create a single nested layer called "nested"
+    if isinstance(n_layer, int):
+        n_layer = {"nested": n_layer}
+
+    # It should be a dictionary
+    if isinstance(n_layer, dict):
+        for key in n_layer:
+            layer_size = n_layer[key]
+            layer_data = {
+                "t": randomstate.random(layer_size * n_base) * 20,
+                "flux": randomstate.random(layer_size * n_base) * 100,
+                "band": randomstate.choice(["r", "g"], size=layer_size * n_base),
+                "index": np.arange(layer_size * n_base) % n_base,
+            }
+            layer_nf = NestedFrame(data=layer_data).set_index("index")
+            base_nf = base_nf.add_nested(layer_nf, key)
+        return base_nf
+    else:
+        raise TypeError("Input to n_layer is not an int or dict.")


Test patch:
diff --git a/tests/datasets/test_generation.py b/tests/datasets/test_generation.py
new file mode 100644
index 0000000..5f52aa8
--- /dev/null
+++ b/tests/datasets/test_generation.py
@@ -0,0 +1,24 @@
+import pytest
+from nested_pandas.datasets import generate_data
+
+
+@pytest.mark.parametrize("n_layers", [10, {"nested_a": 10, "nested_b": 20}])
+def test_generate_data(n_layers):
+    """test the data generator function"""
+    nf = generate_data(10, n_layers, seed=1)
+
+    if isinstance(n_layers, int):
+        assert len(nf.nested.nest.to_flat()) == 100
+
+    elif isinstance(n_layers, dict):
+        assert "nested_a" in nf.columns
+        assert "nested_b" in nf.columns
+
+        assert len(nf.nested_a.nest.to_flat()) == 100
+        assert len(nf.nested_b.nest.to_flat()) == 200
+
+
+def test_generate_data_bad_input():
+    """test a poor n_layer input to generate_data"""
+    with pytest.raises(TypeError):
+        generate_data(10, "nested", seed=1)