From 41b0fba09434062f926ee94eb4a73d3a4cb5ccd4 Mon Sep 17 00:00:00 2001
From: gibsongreen <gibsongreen99@outlook.com>
Date: Thu, 5 Dec 2024 00:59:24 -0500
Subject: [PATCH] handle 1 source edge case, implement max sources for all
 catalog types, test coverage

---
 .../imviz/plugins/catalogs/catalogs.py        | 25 +++++++++++-----
 .../imviz/plugins/catalogs/catalogs.vue       |  9 ++++--
 jdaviz/configs/imviz/tests/test_catalogs.py   | 30 +++++++++++++++----
 3 files changed, 49 insertions(+), 15 deletions(-)

diff --git a/jdaviz/configs/imviz/plugins/catalogs/catalogs.py b/jdaviz/configs/imviz/plugins/catalogs/catalogs.py
index caac1ad56c..d723e8f61f 100644
--- a/jdaviz/configs/imviz/plugins/catalogs/catalogs.py
+++ b/jdaviz/configs/imviz/plugins/catalogs/catalogs.py
@@ -36,7 +36,7 @@ class Catalogs(PluginTemplateMixin, ViewerSelectMixin, HasFileImportSelect, Tabl
     catalog_selected = Unicode("").tag(sync=True)
     results_available = Bool(False).tag(sync=True)
     number_of_results = Int(0).tag(sync=True)
-    max_gaia_sources = IntHandleEmpty(1000).tag(sync=True)
+    max_sources = IntHandleEmpty(1000).tag(sync=True)
 
     # setting the default table headers and values
     _default_table_values = {
@@ -159,6 +159,8 @@ def search(self, error_on_fail=False):
                     zoom_radius = r_max
                 query_region_result = SDSS.query_region(skycoord_center, radius=zoom_radius,
                                                         data_release=17)
+                if len(query_region_result) > self.max_sources:
+                    query_region_result = query_region_result[:self.max_sources]
             except Exception as e:  # nosec
                 errmsg = (f"Failed to query {self.catalog_selected} with c={skycoord_center} and "
                           f"r={zoom_radius}: {repr(e)}")
@@ -184,7 +186,8 @@ def search(self, error_on_fail=False):
         elif self.catalog_selected == 'Gaia':
             from astroquery.gaia import Gaia, conf
 
-            with conf.set_temp("ROW_LIMIT", self.max_gaia_sources):
+            with conf.set_temp("ROW_LIMIT", self.max_sources):
+                Gaia.ROW_LIMIT = self.max_sources
                 sources = Gaia.query_object(skycoord_center, radius=zoom_radius,
                                             columns=('source_id', 'ra', 'dec'))
             self.app._catalog_source_table = sources
@@ -195,8 +198,10 @@ def search(self, error_on_fail=False):
             # but this exceptions might be raised here if setting from_file from the UI
             table = self.catalog.selected_obj
             self.app._catalog_source_table = table
-            skycoord_table = table['sky_centroid']
-
+            if len(table['sky_centroid']) > self.max_sources:
+                skycoord_table = table['sky_centroid'][:self.max_sources]
+            else:
+                skycoord_table = table['sky_centroid']
         else:
             self.results_available = False
             self.number_of_results = 0
@@ -243,6 +248,12 @@ def search(self, error_on_fail=False):
         # NOTE: If performance becomes a problem, see
         # https://docs.astropy.org/en/stable/table/index.html#performance-tips
         if self.catalog_selected == 'From File...':
+            # for single source convert table information to lists for zipping
+            if len(self.app._catalog_source_table) == 1 or self.max_sources == 1:
+                table = [self.app._catalog_source_table]
+                x_coordinates = [x_coordinates]
+                y_coordinates = [y_coordinates]
+
             for row, x_coord, y_coord in zip(self.app._catalog_source_table,
                                              x_coordinates, y_coordinates):
                 # Check if the row contains the required keys
@@ -250,9 +261,9 @@ def search(self, error_on_fail=False):
                             'Declination (degrees)': row['sky_centroid'].dec.deg,
                             'Object ID': str(row.get('label', 'N/A')),
                             'id': len(self.table),
-                            'x_coord': x_coord,
-                            'y_coord': y_coord}
-
+                            'x_coord': x_coord.item() if x_coord.size == 1 else x_coord,
+                            'y_coord': y_coord.item() if y_coord.size == 1 else y_coord
+                            }
                 self.table.add_item(row_info)
 
         filtered_skycoord_table = viewer.state.reference_data.coords.pixel_to_world(x_coordinates,
diff --git a/jdaviz/configs/imviz/plugins/catalogs/catalogs.vue b/jdaviz/configs/imviz/plugins/catalogs/catalogs.vue
index ef72aa77d9..830b02672d 100644
--- a/jdaviz/configs/imviz/plugins/catalogs/catalogs.vue
+++ b/jdaviz/configs/imviz/plugins/catalogs/catalogs.vue
@@ -37,12 +37,15 @@
       </j-docs-link>
     </v-row>
 
-    <v-row v-if="catalog_selected === 'Gaia'">
+    <v-row>
       <v-text-field
-        v-model.number="max_gaia_sources"
+        v-model.number="max_sources"
         type="number"
         step="10"
-        :rules="[() => max_gaia_sources!=='' || 'This field is required']"
+        :rules="[
+          () => max_sources!=='' || 'This field is required',
+          max_sources => max_sources > 0 || 'Max sources must be a postive integer.',
+          ]"
         label="Max sources"
         hint="Maximum number of sources."
         persistent-hint
diff --git a/jdaviz/configs/imviz/tests/test_catalogs.py b/jdaviz/configs/imviz/tests/test_catalogs.py
index b36c56a7eb..d28544e50c 100644
--- a/jdaviz/configs/imviz/tests/test_catalogs.py
+++ b/jdaviz/configs/imviz/tests/test_catalogs.py
@@ -27,7 +27,7 @@
 from astropy.io import fits
 from astropy.nddata import NDData
 from astropy.coordinates import SkyCoord
-from astropy.table import QTable
+from astropy.table import Table, QTable
 
 
 @pytest.mark.remote_data
@@ -182,7 +182,31 @@ def test_offline_ecsv_catalog(imviz_helper, image_2d_wcs, tmp_path):
     assert catalogs_plugin.number_of_results == n_entries
     assert len(imviz_helper.app.data_collection) == 2  # image + markers
 
+    catalogs_plugin.table.selected_rows = [catalogs_plugin.table.items[0]]
+    assert len(catalogs_plugin.table.selected_rows) == 1
+
+    # test to ensure sources searched for respect the maximum sources traitlet
+    catalogs_plugin.max_sources = 1
+    catalogs_plugin.search(error_on_fail=True)
+    assert catalogs_plugin.number_of_results == catalogs_plugin.max_sources
+
+    catalogs_plugin.clear_table()
+
+    # test single source edge case and docs recommended input file type
+    sky_coord = SkyCoord(ra=337.5202807, dec=-20.83305528, unit='deg')
+    tbl = Table({'sky_centroid': [sky_coord], 'label': ['Source_1']})
+    tbl_file = str(tmp_path / 'sky_centroid1.ecsv')
+    tbl.write(tbl_file, overwrite=True)
+    n_entries = len(tbl)
+
+    catalogs_plugin.from_file = tbl_file
+    out_tbl = catalogs_plugin.search()
+    assert len([out_tbl]) == n_entries
+    assert catalogs_plugin.number_of_results == n_entries
+    assert len(imviz_helper.app.data_collection) == 2  # image + markers
+
     catalogs_plugin.clear()
+
     assert not catalogs_plugin.results_available
     assert len(imviz_helper.app.data_collection) == 2  # markers still there, just hidden
 
@@ -190,10 +214,6 @@ def test_offline_ecsv_catalog(imviz_helper, image_2d_wcs, tmp_path):
     assert not catalogs_plugin.results_available
     assert len(imviz_helper.app.data_collection) == 1  # markers gone for good
 
-    catalogs_plugin.table.selected_rows = [catalogs_plugin.table.items[0]]
-
-    assert len(catalogs_plugin.table.selected_rows) == 1
-
     assert imviz_helper.viewers['imviz-0']._obj.state.x_min == -0.5
     assert imviz_helper.viewers['imviz-0']._obj.state.x_max == 9.5
     assert imviz_helper.viewers['imviz-0']._obj.state.y_min == -0.5