add test and support pick and exclude as integers

mne/channels/layout.py

@@ -32,6 +32,7 @@
     _ensure_int,
     fill_doc,
     logger,
+    verbose,
     warn,
 )
 from ..viz.topomap import plot_layout
@@ -52,9 +53,9 @@ class Layout:
     pos : array, shape=(n_channels, 4)
         The unit-normalized positions of the channels in 2d
         (x, y, width, height).
-    names : list
+    names : list of str
         The channel names.
-    ids : list
+    ids : array-like of int
         The channel ids.
     kind : str
         The type of Layout (e.g. 'Vectorview-all').
@@ -70,7 +71,13 @@ def __init__(self, box, pos, names, ids, kind):
         self.kind = kind
 
     def copy(self):
-        """Return a copy of the layout."""
+        """Return a copy of the layout.
+
+        Returns
+        -------
+        layout : instance of Layout
+            A deepcopy of the layout.
+        """
         return deepcopy(self)
 
     def save(self, fname, overwrite=False):
@@ -143,7 +150,7 @@ def plot(self, picks=None, show_axes=False, show=True):
         """
         return plot_layout(self, picks=picks, show_axes=show_axes, show=show)
 
-    @fill_doc
+    @verbose
     def pick(self, picks=None, exclude=(), *, verbose=None):
         """Pick a subset of channels.
 
@@ -174,10 +181,19 @@ def pick(self, picks=None, exclude=(), *, verbose=None):
             picks = np.arange(len(self.names))[picks]
             apply_exclude = False
         else:
-            picks = deepcopy(picks)
+            try:
+                picks = [_ensure_int(picks)]
+            except TypeError:
+                picks = list(deepcopy(picks))
             apply_exclude = False
         if apply_exclude:
-            exclude = [exclude] if isinstance(exclude, str) else deepcopy(exclude)
+            if isinstance(exclude, str):
+                exclude = [exclude]
+            else:
+                try:
+                    exclude = [_ensure_int(exclude)]
+                except TypeError:
+                    exclude = list(deepcopy(exclude))
         for var, var_name in ((picks, "picks"), (exclude, "exclude")):
             if var_name == "exclude" and not apply_exclude:
                 continue

mne/channels/tests/test_layout.py

@@ -23,6 +23,7 @@
 from mne._fiff.constants import FIFF
 from mne._fiff.meas_info import _empty_info
 from mne.channels import (
+    Layout,
     find_layout,
     make_eeg_layout,
     make_grid_layout,
@@ -94,6 +95,18 @@ def _get_test_info():
     return test_info
 
 
+@pytest.fixture(scope="module")
+def layout():
+    """Get a layout."""
+    return Layout(
+        (0.1, 0.2, 0.1, 1.2),
+        pos=np.array([[0, 0, 0.1, 0.1], [0.2, 0.2, 0.1, 0.1], [0.4, 0.4, 0.1, 0.1]]),
+        names=["0", "1", "2"],
+        ids=[0, 1, 2],
+        kind="test",
+    )
+
+
 def test_io_layout_lout(tmp_path):
     """Test IO with .lout files."""
     layout = read_layout(fname="Vectorview-all", scale=False)
@@ -398,3 +411,65 @@ def test_generate_2d_layout():
     # Make sure background image normalizing is correct
     lt_bg = generate_2d_layout(xy, bg_image=bg_image)
     assert_allclose(lt_bg.pos[:, :2].max(), xy.max() / float(sbg))
+
+
+def test_layout_copy(layout):
+    """Test copying a layout."""
+    layout2 = layout.copy()
+    assert_allclose(layout.pos, layout2.pos)
+    assert layout.names == layout2.names
+    layout2.names[0] = "foo"
+    layout2.pos[0, 0] = 0.8
+    assert layout.names != layout2.names
+    assert layout.pos[0, 0] != layout2.pos[0, 0]
+
+
+@pytest.mark.parametrize(
+    "picks, exclude",
+    [
+        ([0, 1], ()),
+        (["0", 1], ()),
+        (None, ["2"]),
+        (None, "2"),
+        (None, [2]),
+        (None, 2),
+        ("all", 2),
+        ("all", "2"),
+        (slice(0, 2), ()),
+        (("0", "1"), ("0", "1")),
+        (("0", 1), ("0", "1")),
+        (("0", 1), (0, "1")),
+    ],
+)
+def test_layout_pick(layout, picks, exclude):
+    """Test selection of channels in a layout."""
+    layout2 = layout.copy()
+    layout2.pick(picks, exclude)
+    assert layout2.names == layout.names[:2]
+    assert_allclose(layout2.pos, layout.pos[:2, :])
+
+
+def test_layout_pick_more(layout):
+    """Test more channel selection in a layout."""
+    layout2 = layout.copy()
+    layout2.pick(0)
+    assert len(layout2.names) == 1
+    assert layout2.names[0] == layout.names[0]
+    assert_allclose(layout2.pos, layout.pos[:1, :])
+
+    layout2 = layout.copy()
+    layout2.pick("all", exclude=("0", "1"))
+    assert len(layout2.names) == 1
+    assert layout2.names[0] == layout.names[2]
+    assert_allclose(layout2.pos, layout.pos[2:, :])
+
+    layout2 = layout.copy()
+    layout2.pick("all", exclude=("0", 1))
+    assert len(layout2.names) == 1
+    assert layout2.names[0] == layout.names[2]
+    assert_allclose(layout2.pos, layout.pos[2:, :])
+
+
+def test_layout_pick_errors(layout):
+    """Test validation of layout.pick."""
+    pass