diff --git a/pynapple/core/time_series.py b/pynapple/core/time_series.py index 72d33d5a..7aa17e2e 100644 --- a/pynapple/core/time_series.py +++ b/pynapple/core/time_series.py @@ -1172,7 +1172,7 @@ def __getitem__(self, key, *args, **kwargs): if ( (len(index) == 1) and (output.ndim == 1) - and ((len(output) > 1) or isinstance(key[1], (list, np.ndarray))) + and ((len(output) > 1) or isinstance(key, int) or isinstance(key[1], (list, np.ndarray))) ): # reshape output of single index to preserve column axis if there are more than one columns being indexed # or if column key is a list or array diff --git a/tests/test_time_series.py b/tests/test_time_series.py index 0c5e9289..1580a73e 100755 --- a/tests/test_time_series.py +++ b/tests/test_time_series.py @@ -954,6 +954,7 @@ def test_interpolate_with_ep(self, tsd): @pytest.mark.parametrize( "tsdframe", [ + nap.TsdFrame(t=np.arange(100), d=np.random.rand(100, 1), time_units="s"), nap.TsdFrame(t=np.arange(100), d=np.random.rand(100, 3), time_units="s"), nap.TsdFrame( t=np.arange(100),