add support to Row in locate_data

Author: zoghbi-a

CHANGES.rst

@@ -17,7 +17,7 @@ Service fixes and enhancements
 heasarc
 ^^^^^^^
 
-- Add support for astropy.table.Row in Heasarc.download_data. [#3270]
+- Add support for astropy.table.Row in Heasarc.download_data and Heasarc.locate_data. [#3270]
 
 
 Infrastructure, Utility and Other Changes and Additions

astroquery/heasarc/core.py

@@ -470,7 +470,7 @@ def locate_data(self, query_result=None, catalog_name=None):
 
         Parameters
         ----------
-        query_result : `astropy.table.Table`, optional
+        query_result : `astropy.table.Table` or `astropy.table.Row`, optional
             A table that contain the search results. Typically as
             returned by query_region. If None, use the table from the
             most recent query_region call.
@@ -492,8 +492,14 @@ def locate_data(self, query_result=None, catalog_name=None):
             else:
                 query_result = self._last_result
 
+        if isinstance(query_result, Row):
+            query_result = query_result.table[[query_result.index]]
+
         if not isinstance(query_result, Table):
-            raise TypeError('query_result need to be an astropy.table.Table')
+            raise TypeError(
+                'query_result need to be an astropy.table.Table or '
+                'astropy.table.Row'
+            )
 
         # make sure we have a column __row
         if '__row' not in query_result.colnames:

astroquery/heasarc/tests/test_heasarc.py

@@ -274,14 +274,28 @@ def test_locate_data():
         Heasarc.locate_data()
 
     with pytest.raises(
-        TypeError, match="query_result need to be an astropy.table.Table"
+        TypeError, match=(
+            "query_result need to be an astropy.table.Table or astropy.table.Row"
+        )
     ):
         Heasarc.locate_data([1, 2])
 
     with pytest.raises(ValueError, match="No __row column found"):
         Heasarc.locate_data(Table({"id": [1, 2, 3.0]}), catalog_name="xray")
 
 
+def test_locate_data_row():
+    table = Table({"id": [1, 2, 3.0]})
+
+    # we cannot do full call as this is not remote,
+    # but if we check that we error on __row not input type
+    with pytest.raises(ValueError, match="No __row column found"):
+        Heasarc.locate_data(table[0], catalog_name="xray")
+
+    with pytest.raises(ValueError, match="No __row column found"):
+        Heasarc.locate_data(table[0:2], catalog_name="xray")
+
+
 def test_download_data__empty():
     with pytest.raises(ValueError, match="Input links table is empty"):
         Heasarc.download_data(Table())