query the NASA exoplanet archive using their API

Author: sashank27

.gitignore

@@ -24,11 +24,23 @@ _build
 docs/api
 docs/_build
 
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
 # Eclipse editor project files
 .project
 .pydevproject
 .settings
 
+# VScode editor project files
+.vscode/*
+
 # Pycharm editor project files
 .idea
 

astroquery/nasa_exoplanet_archive/nasa_exoplanet_archive.py

@@ -8,25 +8,33 @@
 from astropy.table import QTable
 from astropy.coordinates import SkyCoord
 import astropy.units as u
+from urllib.parse import quote
 
 __all__ = ['NasaExoplanetArchive']
 
-EXOPLANETS_CSV_URL = ('http://exoplanetarchive.ipac.caltech.edu/cgi-bin/'
-                      'nstedAPI/nph-nstedAPI?table=exoplanets')
+EXOPLANETS_CSV_URL = ('http://exoplanetarchive.ipac.caltech.edu/cgi-bin/nstedAPI/nph-nstedAPI?table=exoplanets')
 TIME_ATTRS = {'rowupdate': 'iso'}
 BOOL_ATTRS = ('pl_kepflag', 'pl_ttvflag', 'pl_k2flag', 'st_massblend',
               'st_optmagblend', 'st_radblend', 'st_teffblend')
 
 
 class NasaExoplanetArchiveClass(object):
     """
-    Exoplanet Archive querying object. Use the ``get_confirmed_planets_table``
-    or ``query_planet`` methods to get information about exoplanets via the NASA
-    Exoplanet Archive.
+    Querying object for `NExScI Exoplanet Archive Confirmed Planets archive
+    <http://exoplanetarchive.ipac.caltech.edu/index.html>`_
+
+    The Exoplanet Archive table returns lots of columns of data. A full
+    description of the columns can be found `here
+    <https://exoplanetarchive.ipac.caltech.edu/docs/API_exoplanet_columns.html>`_
+
+    Use the ``query_planet`` or ``query_star`` methods to get information
+    about exoplanets via the Archive.
     """
     def __init__(self):
         self._param_units = None
         self._table = None
+        self._cache_planet = None
+        self._cache_star = None
 
     @property
     def param_units(self):
@@ -38,18 +46,16 @@ def param_units(self):
 
         return self._param_units
 
-    def get_confirmed_planets_table(self, cache=True, show_progress=True,
-                                    table_path=None, all_columns=False):
+    def query_planet(self, planet_name, cache=True, show_progress=True,
+                        table_path=None, all_columns=False):
         """
-        Download (and optionally cache) the `NExScI Exoplanet Archive Confirmed
-        Planets table <http://exoplanetarchive.ipac.caltech.edu/index.html>`_.
-
-        The Exoplanet Archive table returns lots of columns of data. A full
-        description of the columns can be found `here
-        <https://exoplanetarchive.ipac.caltech.edu/docs/API_exoplanet_columns.html>`_
+        Download (and optionally cache) the NASA Exoplanet
+        table for a certain planet.
 
         Parameters
         ----------
+        planet_name : str
+            Name of planet
         cache : bool (optional)
             Cache exoplanet table to local astropy cache? Default is `True`.
         show_progress : bool (optional)
@@ -65,27 +71,26 @@ def get_confirmed_planets_table(self, cache=True, show_progress=True,
         Returns
         -------
         table : `~astropy.table.QTable`
-            Table of exoplanet properties.
+            Table of one exoplanet's properties.
         """
-        if self._table is None or not cache:
+
+        planet_name = '\'' + planet_name + '\''
+        EXOPLANET_URL_PLANET = EXOPLANETS_CSV_URL + '&where=pl_name=' + quote(planet_name)
+
+        if self._table is not None and self._cache_planet is planet_name and cache:
+            pass
+
+        else:
             if table_path is None:
-                exoplanets_url = EXOPLANETS_CSV_URL
+                exoplanets_url = EXOPLANET_URL_PLANET
                 if all_columns:
-                    exoplanets_url = EXOPLANETS_CSV_URL + '&select=*'
+                    exoplanets_url = EXOPLANET_URL_PLANET + '&select=*'
 
                 table_path = download_file(exoplanets_url, cache=cache,
                                            show_progress=show_progress,
                                            timeout=120)
             exoplanets_table = ascii.read(table_path)
 
-            # Store column of lowercase names for indexing:
-            lowercase_names = [host_name.lower().replace(' ', '') + letter
-                               for host_name, letter in
-                               zip(exoplanets_table['pl_hostname'].data,
-                                   exoplanets_table['pl_letter'].data)]
-            exoplanets_table['NAME_LOWERCASE'] = lowercase_names
-            exoplanets_table.add_index('NAME_LOWERCASE')
-
             # Create sky coordinate mixin column
             exoplanets_table['sky_coord'] = SkyCoord(ra=exoplanets_table['ra'] * u.deg,
                                                      dec=exoplanets_table['dec'] * u.deg)
@@ -98,28 +103,70 @@ def get_confirmed_planets_table(self, cache=True, show_progress=True,
                         exoplanets_table[col].unit = u.Unit(self.param_units[col])
 
             self._table = QTable(exoplanets_table)
+            self._cache_planet = planet_name
 
         return self._table
 
-    def query_planet(self, planet_name, **kwargs):
+    def query_star(self, host_name, cache=True, show_progress=True,
+                        table_path=None, all_columns=False):
         """
-        Get table of exoplanet properties.
+        Download (and optionally cache) the NASA Exoplanet
+        table for a certain planet.
 
         Parameters
         ----------
-        planet_name : str
-            Name of planet
-        kwargs : dict (optional)
-            Extra keyword arguments passed to ``get_confirmed_planets_table``.
+        host_name : str
+            Name of stellar system
+        cache : bool (optional)
+            Cache exoplanet table to local astropy cache? Default is `True`.
+        show_progress : bool (optional)
+            Show progress of exoplanet table download (if no cached copy is
+            available). Default is `True`.
+        table_path : str (optional)
+            Path to a local table file. Default `None` will trigger a
+            download of the table from the internet.
+        all_columns : bool (optional)
+            Return all available columns. The default returns only the
+            columns in the default category at the link above.
 
         Returns
         -------
         table : `~astropy.table.QTable`
             Table of one exoplanet's properties.
         """
 
-        exoplanet_table = self.get_confirmed_planets_table(**kwargs)
-        return exoplanet_table.loc[planet_name.strip().lower().replace(' ', '')]
+        host_name = '\'' + host_name + '\''
+        EXOPLANET_URL_PLANET = EXOPLANETS_CSV_URL + '&where=pl_hostname=' + quote(host_name)
+
+        if self._table is not None and self._cache_star is host_name and cache:
+            pass
+
+        else:
+            if table_path is None:
+                exoplanets_url = EXOPLANET_URL_PLANET
+                if all_columns:
+                    exoplanets_url = EXOPLANET_URL_PLANET + '&select=*'
+
+                table_path = download_file(exoplanets_url, cache=cache,
+                                           show_progress=show_progress,
+                                           timeout=120)
+            exoplanets_table = ascii.read(table_path)
+
+            # Create sky coordinate mixin column
+            exoplanets_table['sky_coord'] = SkyCoord(ra=exoplanets_table['ra'] * u.deg,
+                                                     dec=exoplanets_table['dec'] * u.deg)
+
+            # Assign units to columns where possible
+            for col in exoplanets_table.colnames:
+                if col in self.param_units:
+                    # Check that unit is implemented in this version of astropy
+                    if hasattr(u, self.param_units[col]):
+                        exoplanets_table[col].unit = u.Unit(self.param_units[col])
+
+            self._table = QTable(exoplanets_table)
+            self._cache_star = host_name
+
+        return self._table
 
 
 NasaExoplanetArchive = NasaExoplanetArchiveClass()

astroquery/nasa_exoplanet_archive/tests/test_nasa_exoplanet_archive_remote.py

@@ -13,20 +13,6 @@
                                 'data', 'nasa_exoplanet_archive.csv')
 
 
-@remote_data
-def test_exoplanet_archive_table():
-    table = NasaExoplanetArchive.get_confirmed_planets_table(cache=False)
-
-    # Check some planets are in the table
-    expected_planets = ['HD 189733 b', 'Kepler-186 f', 'TRAPPIST-1 b',
-                        'HD 209458 b', 'Kepler-62 f', 'GJ 1214 b']
-
-    for planet in expected_planets:
-        assert planet.lower().replace(" ", "") in table['NAME_LOWERCASE']
-
-    assert 'pl_trandep' not in table.colnames
-
-
 @remote_data
 def test_hd209458b_exoplanet_archive():
     # Testing intentionally un-stripped string, no spaced:
@@ -121,17 +107,3 @@ def test_hd209458b_exoplanet_archive_coords():
     sep = params['sky_coord'].separation(simbad_coords)
 
     assert abs(sep) < 5 * u.arcsec
-
-
-@remote_data
-def test_exoplanet_archive_table_all_columns():
-    table = NasaExoplanetArchive.get_confirmed_planets_table(cache=False, all_columns=True)
-
-    # Check some planets are in the table
-    expected_planets = ['HD 189733 b', 'Kepler-186 f', 'TRAPPIST-1 b',
-                        'HD 209458 b', 'Kepler-62 f', 'GJ 1214 b']
-
-    for planet in expected_planets:
-        assert planet.lower().replace(" ", "") in table['NAME_LOWERCASE']
-
-    assert 'pl_trandep' in table.colnames