address reviewer comments

LSSTDESC · May 9, 2024 · 2635398 · 2635398
1 parent aaa953d
commit 2635398
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 24 deletions.
diff --git a/skycatalogs/objects/base_object.py b/skycatalogs/objects/base_object.py
@@ -600,17 +600,10 @@ def __getitem__(self, key):
         elif type(key) == slice:
             if key.start is None:
                 key.start = 0
-            ixdata = [i for i in range(min(key.stop, len(self._ra)))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i)
-                    for i in ixes]
+            return [self.__getitem__(i) for i in range(self.__len__())[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
-            #  check it's a list of int-like?
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i)
-                    for i in key[0]]
+            return[self.__getitem__(i) for i in key[0]]
 
     def get_partition_id(self):
         return self._partition_id
@@ -725,8 +718,8 @@ def __getitem__(self, key):
         If key is a slice return a list of object
         '''
         one_only = isinstance(key, int) or isinstance(key, np.int64)
-        is_slice = type(key) == slice
-        is_list = type(key) == tuple and isinstance(key[0], Iterable)
+        is_slice = isinstance(key, slice)
+        is_list = isinstance(key, tuple) and isinstance(key[0], Iterable)
 
         if one_only:
             start = key

diff --git a/skycatalogs/objects/gaia_object.py b/skycatalogs/objects/gaia_object.py
@@ -377,10 +377,11 @@ def __getitem__(self, key):
             row = {col: self.df[col][key] for col in cols}
             return GaiaObject(row, self, key)
 
-        elif type(key) == slice:
-            ixdata = [i for i in range(min(key.stop, len(self.df['id'])))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self.__getitem__(i) for i in ixes]
+        elif isinstance(key, slice):
+            # ixdata = [i for i in range(min(key.stop, len(self.df['id'])))]
+            # ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
+            # return [self.__getitem__(i) for i in ixes]
+            return [self.__getitem__(i) for i in range(len(self.df))[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
             return [self.__getitem__(i) for i in key[0]]

diff --git a/skycatalogs/objects/sso_object.py b/skycatalogs/objects/sso_object.py
@@ -163,14 +163,7 @@ def __getitem__(self, key):
         elif type(key) == slice:
             if key.start is None:
                 key.start = 0
-            ixdata = [i for i in range(min(key.stop, len(self._ra)))]
-            ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i, self._mjds[i])
-                    for i in ixes]
+            return [self.__getitem__(i) for i in range(self.__len__())[key]]
 
         elif type(key) == tuple and isinstance(key[0], Iterable):
-            #  check it's a list of int-like?
-            return [self._object_class(self._ra[i], self._dec[i], self._id[i],
-                                       object_type, self, i, self._mjds[i])
-                    for i in key[0]]
+            return [self.__getitem__(i) for i in key[0]]