FEAT/FIX: added support for proper stack(Session) (closes #1057)

I consider this as a bug fix because stack(Session) *seemed* to work but did not keep labels
(because the Session was considered as a Sequence), which is surprising and caused issue in production code

Author: gdementen

doc/source/changes/version_0_34_1.rst.inc

@@ -20,6 +20,8 @@ New features
 
 * added support for Python 3.11.
 
+* added support for stacking all arrays of a Session by simply doing: `stack(my_session)` instead of
+  `stack(my_session.items())` (closes :issue:`1057`).
 
 .. _misc:
 

larray/core/array.py

@@ -9586,13 +9586,12 @@ def stack(elements=None, axes=None, title=None, meta=None, dtype=None, res_axes=
 
     Parameters
     ----------
-    elements : tuple, list or dict.
+    elements : tuple, list, dict or Session.
         Elements to stack. Elements can be scalars, arrays, sessions, (label, value) pairs or a {label: value} mapping.
-        In the later case, axis must be defined and cannot be a name only, because we need to have labels order,
-        which the mapping does not provide.
 
-        Stacking sessions will return a new session containing the arrays of all sessions stacked together. An array
-        missing in a session will be replaced by NaN.
+        Stacking a single session will stack all its arrays in a single array.
+        Stacking several sessions will take the corresponding arrays in all the sessions and stack them, returning a
+        new session. An array missing in a session will be replaced by NaN.
     axes : str, Axis, Group or sequence of Axis, optional
         Axes to create. If None, defaults to a range() axis.
     title : str, optional
@@ -9733,14 +9732,18 @@ def stack(elements=None, axes=None, title=None, meta=None, dtype=None, res_axes=
     if kwargs:
         elements = kwargs.items()
 
-    if isinstance(elements, dict):
+    if isinstance(elements, (dict, Session)):
         elements = elements.items()
 
     if isinstance(elements, Array):
         if axes is None:
             axes = -1
         axes = elements.axes[axes]
         items = elements.items(axes)
+    elif isinstance(elements, Session):
+        if axes is None:
+            axes = 'array'
+        items = elements.items()
     elif isinstance(elements, Iterable):
         if not isinstance(elements, Sequence):
             elements = list(elements)

larray/tests/test_session.py

@@ -657,6 +657,27 @@ def test_arrays():
 
 
 def test_stack():
+    # stacking all arrays of a single session
+    # =======================================
+    # a) using explicit axis
+    s = Session(arr1=ndtest(3), arr2=ndtest(3) + 10)
+    axis = Axis("array=arr2,arr1")
+    res = stack(s, axis)
+    expected = stack(arr1=s.arr1, arr2=s.arr2, axes=axis)
+    assert_larray_equal(res, expected)
+
+    # b) using explicit axis name
+    s = Session(arr1=ndtest(3), arr2=ndtest(3) + 10)
+    res = stack(s, "array")
+    expected = stack(arr1=s.arr1, arr2=s.arr2, axes="array")
+    assert_larray_equal(res, expected)
+
+    # c) using not axis information
+    s = Session(arr1=ndtest(3), arr2=ndtest(3) + 10)
+    res = stack(s)
+    expected = stack(arr1=s.arr1, arr2=s.arr2)
+    assert_larray_equal(res, expected)
+
     # stacking two sessions (it will stack all arrays with corresponding names
     s1 = Session(arr1=ndtest(3), arr2=ndtest(3) + 10)
     s2 = Session(arr1=ndtest(3), arr2=ndtest(3) + 30)