added support for eye(AxisCollection)

Author: gdementen

doc/source/changes/version_0_33.rst.inc

@@ -55,6 +55,8 @@ Miscellaneous improvements
   - support passing a dict as legend to customize the legend.
   - many tweaks to make several plots look better out of the box.
 
+* :py:obj:`eye()` now supports an AxisCollection as argument, so you can use axes from another array by using
+  ``eye(other_array.axes)``.
 
 Fixes
 ^^^^^

larray/core/array.py

@@ -9192,10 +9192,10 @@ def eye(rows, columns=None, k=0, title=None, dtype=None, meta=None):
 
     Parameters
     ----------
-    rows : int or Axis
-        Rows of the output.
+    rows : int or Axis or tuple or length 2 AxisCollection
+        Rows of the output (if int or Axis) or rows and columns (if tuple or AxisCollection).
     columns : int or Axis, optional
-        Columns of the output. If None, defaults to rows.
+        Columns of the output. Defaults to the value of `rows` if it is an int or Axis.
     k : int, optional
         Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper
         diagonal, and a negative value to a lower diagonal.
@@ -9214,16 +9214,16 @@ def eye(rows, columns=None, k=0, title=None, dtype=None, meta=None):
 
     Examples
     --------
+    >>> eye('sex=M,F')
+    sex\sex    M    F
+          M  1.0  0.0
+          F  0.0  1.0
     >>> eye(2, dtype=int)
     {0}*\{1}*  0  1
             0  1  0
             1  0  1
-    >>> sex = Axis('sex=M,F')
-    >>> eye(sex)
-    sex\sex    M    F
-          M  1.0  0.0
-          F  0.0  1.0
     >>> age = Axis('age=0..2')
+    >>> sex = Axis('sex=M,F')
     >>> eye(age, sex)
     age\sex    M    F
           0  1.0  0.0
@@ -9236,9 +9236,16 @@ def eye(rows, columns=None, k=0, title=None, dtype=None, meta=None):
             2  0.0  0.0  0.0
     """
     meta = _handle_meta(meta, title)
-    if columns is None:
-        columns = rows.copy() if isinstance(rows, Axis) else rows
-    axes = AxisCollection([rows, columns])
+    if isinstance(rows, AxisCollection):
+        assert columns is None
+        axes = rows
+    elif isinstance(rows, (tuple, list)):
+        assert columns is None
+        axes = AxisCollection(rows)
+    else:
+        if columns is None:
+            columns = rows.copy() if isinstance(rows, Axis) else rows
+        axes = AxisCollection([rows, columns])
     shape = axes.shape
     data = np.eye(shape[0], shape[1], k, dtype)
     return Array(data, axes, meta=meta)

larray/tests/test_array.py

@@ -14,9 +14,10 @@
                                  assert_array_equal, assert_array_nan_equal, assert_larray_equiv, assert_larray_equal,
                                  needs_xlwings, needs_pytables, needs_xlsxwriter, needs_openpyxl, needs_python37,
                                  must_warn)
-from larray import (Array, LArray, Axis, LGroup, union, zeros, zeros_like, ndtest, empty, ones, eye, diag, stack,
+from larray import (Array, LArray, Axis, AxisCollection, LGroup, IGroup,
+                    union, zeros, zeros_like, ndtest, empty, ones, eye, diag, stack,
                     clip, exp, where, X, mean, isnan, round, read_hdf, read_csv, read_eurostat, read_excel,
-                    from_lists, from_string, open_excel, from_frame, sequence, nan, IGroup)
+                    from_lists, from_string, open_excel, from_frame, sequence, nan)
 from larray.inout.pandas import from_series
 from larray.core.axis import _to_ticks, _to_key
 from larray.util.misc import LHDFStore
@@ -4508,6 +4509,37 @@ def test_ufuncs(small_array):
     assert_array_equal(rounded, np.round(raw + 0.6))
 
 
+def test_eye():
+    age = Axis('age=0..2')
+    sex = Axis('sex=M,F')
+
+    # using one Axis object
+    res = eye(sex)
+    expected = from_string(r'''
+sex\sex    M    F
+      M  1.0  0.0
+      F  0.0  1.0''')
+    assert_array_equal(res, expected)
+
+    # using an AxisCollection
+    res = eye(AxisCollection([age, sex]))
+    expected = from_string(r'''
+age\sex    M    F
+      0  1.0  0.0
+      1  0.0  1.0
+      2  0.0  0.0''')
+    assert_array_equal(res, expected)
+
+    # using a tuple of axes
+    res = eye((age, sex))
+    expected = from_string(r"""
+age\sex    M    F
+      0  1.0  0.0
+      1  0.0  1.0
+      2  0.0  0.0""")
+    assert_array_equal(res, expected)
+
+
 def test_diag():
     # 2D -> 1D
     a = ndtest((3, 3))