add automatic host guess in download_data

Author: zoghbi-a

CHANGES.rst

@@ -75,6 +75,7 @@ heasarc
 - Add ``query_by_column`` to allow querying of different catalog columns. [#3403]
 - Add support for uploading tables when using TAP directly through ``query_tap``. [#3403]
 - Improve how maxrec works. If it is bigger than the default server limit, add a TOP statement. [#3403]
+- Add automatic guessing for the data host in ``download_data``. [#3403]
 
 alma
 ^^^^

astroquery/heasarc/core.py

@@ -750,6 +750,36 @@ def enable_cloud(self, provider='aws', profile=None):
 
         self.s3_client = self.s3_resource.meta.client
 
+    def _guess_host(self, host):
+        """Guess the host to use for downloading data
+
+        Parameters
+        ----------
+        host : str
+            The host provided by the user
+
+        Returns
+        -------
+        host : str
+            The guessed host
+
+        """
+        if host in ['heasarc', 'sciserver', 'aws']:
+            return host
+        elif host is not None:
+            raise ValueError(
+                'host has to be one of heasarc, sciserver, aws or None')
+
+        # host is None, so we guess
+        if os.environ['HOME'] == '/home/idies' and os.path.exists('/FTP/'):
+            # we are on idies, so we can use sciserver
+            return 'sciserver'
+
+        for var in ['AWS_REGION', 'AWS_DEFAULT_REGION', 'AWS_ROLE_ARN']:
+            if var in os.environ:
+                return 'aws'
+        return 'heasarc'
+
     def download_data(self, links, host='heasarc', location='.'):
         """Download data products in links with a choice of getting the
         data from either the heasarc server, sciserver, or the cloud in AWS.
@@ -781,8 +811,8 @@ def download_data(self, links, host='heasarc', location='.'):
         if isinstance(links, Row):
             links = links.table[[links.index]]
 
-        if host not in ['heasarc', 'sciserver', 'aws']:
-            raise ValueError('host has to be one of heasarc, sciserver, aws')
+        # guess the host if not provided
+        host = self._guess_host(host)
 
         host_column = 'access_url' if host == 'heasarc' else host
         if host_column not in links.colnames:

astroquery/heasarc/tests/test_heasarc.py

@@ -506,6 +506,29 @@ def test_locate_data_row():
         Heasarc.locate_data(table[0:2], catalog_name="xray")
 
 
+def test__guess_host_default():
+    # Use a new HeasarcClass object
+    assert Heasarc._guess_host(host=None) == 'heasarc'
+
+
+@pytest.mark.parametrize("host", ["heasarc", "sciserver", "aws"])
+def test__guess_host_know(host):
+    # Use a new HeasarcClass object
+    assert Heasarc._guess_host(host=host) == host
+
+
+def test__guess_host_sciserver(monkeypatch):
+    monkeypatch.setenv("HOME", "/home/idies")
+    monkeypatch.setattr("os.path.exists", lambda path: path.startswith('/FTP'))
+    assert Heasarc._guess_host(host=None) == 'sciserver'
+
+
+@pytest.mark.parametrize("var", ["AWS_REGION", "AWS_REGION_DEFAULT", "AWS_ROLE_ARN"])
+def test__guess_host_aws(monkeypatch, var):
+    monkeypatch.setenv("AWS_REGION", var)
+    assert Heasarc._guess_host(host=None) == 'aws'
+
+
 def test_download_data__empty():
     with pytest.raises(ValueError, match="Input links table is empty"):
         Heasarc.download_data(Table())

docs/heasarc/heasarc.rst

@@ -247,6 +247,7 @@ You can specify where the data are to be downloaded using the ``location`` param
 
 To download the data, you can pass ``links`` table (or row) to `~astroquery.heasarc.HeasarcClass.download_data`,
 specifying from where you want the data to be fetched by specifying the ``host`` parameter. By default,
+the function will try to guess the best host based on your environment. If it cannot guess, then
 the data is fetched from the main HEASARC servers.
 The recommendation is to use different hosts depending on where your code is running:
 * ``host='sciserver'``: Use this option if you running you analysis on Sciserver. Because