Merge pull request #2494 from eerovaher/simbad-refractor

Simplify `simbad` code

Author: bsipocz

CHANGES.rst

@@ -85,6 +85,13 @@ oac
 - Fix bug in parsing events that contain html tags (e.g. in their alias
   field). [#2423]
 
+simbad
+^^^^^^
+
+- It is now possible to specify multiple coordinates together with a single
+  radius as a string in ``query_region()`` and ``query_region_async()``.
+  [#2494]
+
 svo_fps
 ^^^^^^^
 

astroquery/simbad/core.py

@@ -3,7 +3,6 @@
 Simbad query class for accessing the Simbad Service
 """
 
-import copy
 import re
 import requests
 import json
@@ -12,6 +11,7 @@
 from io import BytesIO
 import warnings
 import astropy.units as u
+from astropy.utils import isiterable
 from astropy.utils.data import get_pkg_data_filename
 import astropy.coordinates as coord
 from astropy.table import Table
@@ -202,25 +202,20 @@ class SimbadBibcodeResult(SimbadResult):
     """Bibliography-type Simbad result"""
     @property
     def table(self):
-        bibcode_match = bibcode_regex.search(self.script)
-        splitter = bibcode_match.group(2)
-        ref_list = [splitter + ref for ref in self.data.split(splitter)][1:]
-        max_len = max([len(r) for r in ref_list])
-        table = Table(names=['References'], dtype=['U%i' % max_len])
-        for ref in ref_list:
-            table.add_row([ref])
-        return table
+        splitter = bibcode_regex.search(self.script).group(2)
+        ref_list = [[splitter + ref] for ref in self.data.split(splitter)[1:]]
+        max_len = max(len(r[0]) for r in ref_list)
+        return Table(rows=ref_list, names=['References'], dtype=[f"U{max_len}"])
 
 
 class SimbadObjectIDsResult(SimbadResult):
     """Object identifier list Simbad result"""
     @property
     def table(self):
-        max_len = max([len(i) for i in self.data.splitlines()])
-        table = Table(names=['ID'], dtype=['S%i' % max_len])
-        for id in self.data.splitlines():
-            table.add_row([id.strip()])
-        return table
+        split_lines = self.data.splitlines()
+        ids = [[id.strip()] for id in split_lines]
+        max_len = max(map(len, split_lines))
+        return Table(rows=ids, names=['ID'], dtype=[f"S{max_len}"])
 
 
 class SimbadBaseQuery(BaseQuery):
@@ -281,8 +276,6 @@ class SimbadClass(SimbadBaseQuery):
                   ),
         '[^0-9]': 'Any (one) character not in the list.'}
 
-    _ORDERED_WILDCARDS = ['*', '?', '[abc]', '[^0-9]']
-
     # query around not included since this is a subcase of query_region
     _function_to_command = {
         'query_object_async': 'query id',
@@ -303,7 +296,7 @@ class SimbadClass(SimbadBaseQuery):
 
     def __init__(self):
         super().__init__()
-        self._VOTABLE_FIELDS = copy.copy(self._VOTABLE_FIELDS)
+        self._VOTABLE_FIELDS = self._VOTABLE_FIELDS.copy()
 
     def list_wildcards(self):
         """
@@ -323,10 +316,7 @@ def list_wildcards(self):
         [abc] : Exactly one character taken in the list.
                 Can also be defined by a range of characters: [A-Z]
         """
-        for key in self._ORDERED_WILDCARDS:
-            print("{key} : {value}\n".format(key=key,
-                                             value=self.WILDCARDS[key]))
-        return
+        print("\n\n".join(f"{k} : {v}" for k, v in self.WILDCARDS.items()))
 
     def list_votable_fields(self):
         """
@@ -353,8 +343,8 @@ def list_votable_fields(self):
             fields_dict = json.load(f)
 
         print("Available VOTABLE fields:\n")
-        for field in list(sorted(fields_dict.keys())):
-            print("{}".format(field))
+        for field in sorted(fields_dict.keys()):
+            print(str(field))
         print("For more information on a field:\n"
               "Simbad.get_field_description ('field_name') \n"
               "Currently active VOTABLE fields:\n {0}"
@@ -415,10 +405,7 @@ def add_votable_fields(self, *args):
             os.path.join('data', 'votable_fields_dict.json'))
 
         with open(dict_file, "r") as f:
-            fields_dict = json.load(f)
-            fields_dict = dict(
-                ((strip_field(ff), fields_dict[ff])
-                 for ff in fields_dict))
+            fields_dict = {strip_field(k): v for k, v in json.load(f).items()}
 
         for field in args:
             sf = strip_field(field)
@@ -427,34 +414,30 @@ def add_votable_fields(self, *args):
             else:
                 self._VOTABLE_FIELDS.append(field)
 
-    def remove_votable_fields(self, *args, **kwargs):
+    def remove_votable_fields(self, *args, strip_params=False):
         """
         Removes the specified field names from ``SimbadClass._VOTABLE_FIELDS``
 
         Parameters
         ----------
         list of field_names to be removed
-        strip_params: bool
+        strip_params: bool, optional
             If true, strip the specified keywords before removing them:
             e.g., ra(foo) would remove ra(bar) if this is True
         """
-        strip_params = kwargs.pop('strip_params', False)
-
         if strip_params:
-            sargs = [strip_field(a) for a in args]
+            sargs = {strip_field(a) for a in args}
             sfields = [strip_field(a) for a in self._VOTABLE_FIELDS]
         else:
-            sargs = args
+            sargs = set(args)
             sfields = self._VOTABLE_FIELDS
-        absent_fields = set(sargs) - set(sfields)
-
-        for b, f in list(zip(sfields, self._VOTABLE_FIELDS)):
-            if b in sargs:
-                self._VOTABLE_FIELDS.remove(f)
 
-        for field in absent_fields:
+        for field in sargs.difference(sfields):
             warnings.warn("{field}: this field is not set".format(field=field))
 
+        zipped_fields = zip(sfields, self._VOTABLE_FIELDS)
+        self._VOTABLE_FIELDS = [f for b, f in zipped_fields if b not in sargs]
+
         # check if all fields are removed
         if not self._VOTABLE_FIELDS:
             warnings.warn("All fields have been removed. "
@@ -678,8 +661,6 @@ def query_region_async(self, coordinates, radius=2*u.arcmin,
 
         # handle the vector case
         if isinstance(ra, list):
-            vector = True
-
             if len(ra) > 10000:
                 warnings.warn("For very large queries, you may receive a "
                               "timeout error.  SIMBAD suggests splitting "
@@ -689,21 +670,19 @@ def query_region_async(self, coordinates, radius=2*u.arcmin,
             if len(set(frame)) > 1:
                 raise ValueError("Coordinates have different frames")
             else:
-                frame = set(frame).pop()
+                frame = frame[0]
 
-            if vector and _has_length(radius) and len(radius) == len(ra):
-                # all good, continue
-                pass
-            elif vector and _has_length(radius) and len(radius) != len(ra):
-                raise ValueError("Mismatch between radii and coordinates")
-            elif vector and not _has_length(radius):
+            # `radius` as `str` is iterable, but contains only one value.
+            if isiterable(radius) and not isinstance(radius, str):
+                if len(radius) != len(ra):
+                    raise ValueError("Mismatch between radii and coordinates")
+            else:
                 radius = [_parse_radius(radius)] * len(ra)
 
-            if vector:
-                query_str = "\n".join([base_query_str
-                                      .format(ra=ra_, dec=dec_, rad=rad_,
-                                              frame=frame, equinox=equinox)
-                                      for ra_, dec_, rad_ in zip(ra, dec, radius)])
+            query_str = "\n".join(base_query_str
+                                  .format(ra=ra_, dec=dec_, rad=rad_,
+                                          frame=frame, equinox=equinox)
+                                  for ra_, dec_, rad_ in zip(ra, dec, radius))
 
         else:
             radius = _parse_radius(radius)
@@ -956,20 +935,13 @@ def query_objectids_async(self, object_name, cache=True,
         return response
 
     def _get_query_header(self, get_raw=False):
-        votable_fields = ','.join(self.get_votable_fields())
         # if get_raw is set then don't fetch as votable
         if get_raw:
             return ""
-        votable_def = "votable {" + votable_fields + "}"
-        votable_open = "votable open"
-        return "\n".join([votable_def, votable_open])
+        return "votable {" + ','.join(self.get_votable_fields()) + "}\nvotable open"
 
     def _get_query_footer(self, get_raw=False):
-        if get_raw:
-            return ""
-        votable_close = "votable close"
-
-        return votable_close
+        return "" if get_raw else "votable close"
 
     @validate_epoch_decorator
     @validate_equinox_decorator
@@ -1004,25 +976,19 @@ def _args_to_payload(self, *args, **kwargs):
             kwargs['equi'] = kwargs['equinox']
             del kwargs['equinox']
         # remove default None from kwargs
-        # be compatible with python3
-        for key in list(kwargs):
-            if not kwargs[key]:
-                del kwargs[key]
+        kwargs = {key: value for key, value in kwargs.items() if value is not None}
         # join in the order specified otherwise results in error
         all_keys = ['radius', 'frame', 'equi', 'epoch']
         present_keys = [key for key in all_keys if key in kwargs]
         if caller == 'query_criteria_async':
-            for k in kwargs:
-                present_keys.append(k)
+            present_keys.extend(kwargs)
             # need ampersands to join args
             args_str = '&'.join([str(val) for val in args])
-            if len(args) > 0 and len(present_keys) > 0:
+            if args and present_keys:
                 args_str += " & "
         else:
             args_str = ' '.join([str(val) for val in args])
-        kwargs_str = ' '.join("{key}={value}".format(key=key,
-                                                     value=kwargs[key])
-                              for key in present_keys)
+        kwargs_str = ' '.join(f"{key}={kwargs[key]}" for key in present_keys)
 
         # For the record, I feel dirty for writing this wildcard-case hack.
         # This entire function should be refactored when someone has time.
@@ -1081,17 +1047,8 @@ def _parse_coordinates(coordinates):
         raise ValueError("Coordinates not specified correctly")
 
 
-def _has_length(x):
-    # some objects have '__len__' attributes but have no len()
-    try:
-        len(x)
-        return True
-    except (TypeError, AttributeError):
-        return False
-
-
 def _get_frame_coords(coordinates):
-    if _has_length(coordinates):
+    if isiterable(coordinates):
         # deal with vectors differently
         parsed = [_get_frame_coords(cc) for cc in coordinates]
         return ([ra for ra, dec, frame in parsed],

astroquery/simbad/tests/test_simbad.py

@@ -243,6 +243,7 @@ def test_query_catalog(patch_post):
                           (FK4_COORDS, '5d0m0s', 2000.0, 'J2000'),
                           (FK5_COORDS, None, 2000.0, 'J2000'),
                           (multicoords, 0.5*u.arcsec, 2000.0, 'J2000'),
+                          (multicoords, "0.5s", 2000.0, 'J2000'),
                           ])
 def test_query_region_async(patch_post, coordinates, radius, equinox, epoch):
     response1 = simbad.core.Simbad.query_region_async(
@@ -283,6 +284,11 @@ def test_query_region_radius_error(patch_post, coordinates, radius,
             coordinates, radius=radius, equinox=equinox, epoch=epoch)
 
 
+def test_query_region_coord_radius_mismatch():
+    with pytest.raises(ValueError, match="^Mismatch between radii and coordinates$"):
+        simbad.SimbadClass().query_region(multicoords, radius=[1, 2, 3] * u.deg)
+
+
 @pytest.mark.parametrize(('coordinates', 'radius', 'equinox', 'epoch'),
                          [(ICRS_COORDS, "0d", 2000.0, 'J2000'),
                           (GALACTIC_COORDS, 1.0 * u.marcsec, 2000.0, 'J2000')

astroquery/simbad/tests/test_simbad_remote.py

@@ -107,10 +107,11 @@ def test_query_region_async(self, temp_dir):
 
         assert response is not None
 
-    def test_query_region_async_vector(self, temp_dir):
+    @pytest.mark.parametrize("radius", (0.5 * u.arcsec, "0.5s"))
+    def test_query_region_async_vector(self, temp_dir, radius):
         simbad = Simbad()
         simbad.cache_location = temp_dir
-        response1 = simbad.query_region_async(multicoords, radius=0.5*u.arcsec)
+        response1 = simbad.query_region_async(multicoords, radius=radius)
         assert response1.request.body == 'script=votable+%7Bmain_id%2Ccoordinates%7D%0Avotable+open%0Aquery+coo+5%3A35%3A17.3+-80%3A52%3A00+radius%3D0.5s+frame%3DICRS+equi%3D2000.0%0Aquery+coo+17%3A47%3A20.4+-28%3A23%3A07.008+radius%3D0.5s+frame%3DICRS+equi%3D2000.0%0Avotable+close'   # noqa
 
     def test_query_region(self, temp_dir):