From 705a8a3632e929b4b1fd84eafa616c015a49574f Mon Sep 17 00:00:00 2001
From: gviejo <guillaume.viejo@gmail.com>
Date: Sun, 3 Mar 2024 06:30:44 -0500
Subject: [PATCH] Fixing IntervalSet loc

---
 docs/HISTORY.md               |  6 +++++
 pynapple/__init__.py          |  2 +-
 pynapple/core/interval_set.py | 11 +++++++++
 pynapple/core/utils.py        | 42 +++++++++++++++++++++++++++++++----
 pyproject.toml                |  2 +-
 setup.py                      |  2 +-
 6 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/docs/HISTORY.md b/docs/HISTORY.md
index db1f1413..318834d6 100644
--- a/docs/HISTORY.md
+++ b/docs/HISTORY.md
@@ -9,6 +9,12 @@ In 2018, Francesco started neuroseries, a Python package built on Pandas. It was
 In 2021, Guillaume and other trainees in Adrien's lab decided to fork from neuroseries and started *pynapple*. The core of pynapple is largely built upon neuroseries. Some of the original changes to TSToolbox made by Luke were included in this package, especially the *time_support* property of all ts/tsd objects.
 
 
+0.6.1 (2024-03-03)
+------------------
+
+- Fixed pynapple `loc` method for new `IntervalSet`
+
+
 0.6.0 (2024-03-02)
 ------------------
 
diff --git a/pynapple/__init__.py b/pynapple/__init__.py
index 989e5ccc..0855144a 100644
--- a/pynapple/__init__.py
+++ b/pynapple/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 from .core import IntervalSet, Ts, Tsd, TsdFrame, TsdTensor, TsGroup, TsIndex, config
 from .io import *
 from .process import *
diff --git a/pynapple/core/interval_set.py b/pynapple/core/interval_set.py
index ec91f311..d3c6cb93 100644
--- a/pynapple/core/interval_set.py
+++ b/pynapple/core/interval_set.py
@@ -539,6 +539,17 @@ def get_intervals_center(self, alpha=0.5):
         t = starts + (ends - starts) * alpha
         return time_series.Ts(t=t, time_support=self)
 
+    def as_dataframe(self):
+        """
+        Convert the `IntervalSet` object to a pandas.DataFrame object.
+
+        Returns
+        -------
+        out: pandas.DataFrame
+            _
+        """
+        return pd.DataFrame(data=self.values, columns=["start", "end"])
+
     def save(self, filename):
         """
         Save IntervalSet object in npz format. The file will contain the starts and ends.
diff --git a/pynapple/core/utils.py b/pynapple/core/utils.py
index 91f88092..bea4c850 100644
--- a/pynapple/core/utils.py
+++ b/pynapple/core/utils.py
@@ -2,14 +2,13 @@
 # @Author: Guillaume Viejo
 # @Date:   2024-02-09 11:45:45
 # @Last Modified by:   gviejo
-# @Last Modified time: 2024-02-21 21:27:04
+# @Last Modified time: 2024-03-03 06:28:59
 
 """
     Utility functions
 """
 
 import warnings
-from numbers import Number
 
 import numpy as np
 from numba import jit
@@ -279,22 +278,57 @@ def __getitem__(self, key):
 
 
 class _IntervalSetSliceHelper:
+    """
+    This class helps `IntervalSet` behaves like pandas.DataFrame for the `loc` function.
+
+    Attributes
+    ----------
+    intervalset : `IntervalSet` to slice
+
+    """
+
     def __init__(self, intervalset):
+        """Class for `loc` slicing function
+
+        Parameters
+        ----------
+        intervalset : IntervalSet
+
+        """
         self.intervalset = intervalset
 
     def __getitem__(self, key):
+        """Getters for `IntervalSet.loc`. Mimics pandas.DataFrame.
+
+        Parameters
+        ----------
+        key : int, list or tuple
+
+        Returns
+        -------
+        IntervalSet or Number or numpy.ndarray
+
+        Raises
+        ------
+        IndexError
+
+        """
         if key in ["start", "end"]:
             return self.intervalset[key]
         elif isinstance(key, list):
             return self.intervalset[key]
-        elif isinstance(key, Number):
+        elif isinstance(key, int):
             return self.intervalset.values[key]
         else:
             if isinstance(key, tuple):
                 if len(key) == 2:
                     if key[1] not in ["start", "end"]:
                         raise IndexError
-                    return self.intervalset[key[0]][key[1]]
+                    out = self.intervalset[key[0]][key[1]]
+                    if len(out) == 1:
+                        return out[0]
+                    else:
+                        return out
                 else:
                     raise IndexError
             else:
diff --git a/pyproject.toml b/pyproject.toml
index 08e570f4..c0cb7ad6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pynapple"
-version = "0.6.0"
+version = "0.6.1"
 description = "PYthon Neural Analysis Package Pour Laboratoires d’Excellence"
 readme = "README.md"
 authors = [{ name = "Guillaume Viejo", email = "guillaume.viejo@gmail.com" }]
diff --git a/setup.py b/setup.py
index ed2e8f42..16747f71 100644
--- a/setup.py
+++ b/setup.py
@@ -59,7 +59,7 @@
     test_suite='tests',
     tests_require=test_requirements,
     url='https://github.com/pynapple-org/pynapple',
-    version='v0.6.0',
+    version='v0.6.1',
     zip_safe=False,
     long_description_content_type='text/markdown',
     download_url='https://github.com/pynapple-org/pynapple/archive/refs/tags/v0.6.0.tar.gz'