Merge pull request #467 from keflavich/issue467

astroquery.vizier not retrieving selected columns

Author: keflavich

.travis.yml

@@ -72,7 +72,7 @@ install:
     - source activate test
 
     # CORE DEPENDENCIES
-    - if [[ $SETUP_CMD != egg_info ]]; then $CONDA_INSTALL numpy=$NUMPY_VERSION pytest pip Cython; fi
+    - if [[ $SETUP_CMD != egg_info ]]; then $CONDA_INSTALL numpy=$NUMPY_VERSION pytest pip Cython jinja2; fi
     - if [[ $SETUP_CMD != egg_info ]]; then $PIP_INSTALL pytest-xdist; fi
 
     # ASTROPY

astropy_helpers

@@ -4,6 +4,7 @@
 import os
 import warnings
 import json
+import copy
 
 from astropy.extern import six
 import astropy.units as u
@@ -56,7 +57,8 @@ def __init__(self, columns=["*"], column_filters={}, catalog=None, keywords=None
     def columns(self):
         """ Columns to include.  The special keyword 'all' will return ALL
         columns from ALL retrieved tables. """
-        return self._columns
+        # columns need to be immutable but still need to be a list
+        return list(tuple(self._columns))
 
     @columns.setter
     def columns(self, values):
@@ -459,12 +461,16 @@ def _args_to_payload(self, *args, **kwargs):
         # process: columns
         columns = kwargs.get('columns')
         if columns is None:
-            columns = self.columns
+            columns = copy.copy(self.columns)
         else:
             columns = self.columns + columns
 
-        if 'all' in columns:
-            columns.remove('all')
+        # keyword names that can mean 'all' need to be treated separately
+        alls = ['all','*']
+        if any(x in columns for x in alls):
+            for x in alls:
+                if x in columns:
+                    columns.remove(x)
             body['-out.all'] = 2
 
         # process: columns - always request computed positions in degrees
@@ -485,6 +491,7 @@ def _args_to_payload(self, *args, **kwargs):
             else:
                 columns_out += [column]
         body['-out.add'] = ','.join(columns_out)
+        body['-out'] = columns_out
         if len(sorts_out) > 0:
             body['-sort'] = ','.join(sorts_out)
         # process: maximum rows returned

astroquery/vizier/core.py

@@ -51,6 +51,25 @@ def test_query_Vizier_instance(self):
                                 catalog=["NOMAD", "UCAC"])
         assert isinstance(result, commons.TableList)
 
+    def test_vizier_column_restriction(self):
+        # Check that the column restriction worked.  At least some of these
+        # catalogs include Bmag's
+        v = vizier.core.Vizier(columns=['_RAJ2000', 'DEJ2000', 'B-V', 'Vmag', 'Plx'],
+                               column_filters={"Vmag":">10"}, keywords=["optical", "xry"])
+        result = v.query_object("HD 226868", catalog=["NOMAD", "UCAC"])
+        for table in result:
+            assert 'Bmag' not in table.columns
+
+    @pytest.mark.parametrize('all',('all','*'))
+    def test_alls_withaddition(self, all):
+        # Check that all the expected columns are there plus the _r
+        # (radius from target) that we've added
+        v = vizier.core.Vizier(columns=[all, "+_r"], catalog="II/246")
+        result = v.query_region("HD 226868", radius="20s")
+        table = result['II/246/out']
+        assert 'Jmag' in table.columns
+        assert '_r' in table.columns
+
     def test_get_catalogs(self):
         result = vizier.core.Vizier.get_catalogs('J/ApJ/706/83')
         assert isinstance(result, commons.TableList)
@@ -73,12 +92,16 @@ def test_multicoord(self):
         targets = commons.ICRSCoordGenerator(ra=[299.590, 299.90],
                                              dec=[35.201, 35.201],
                                              unit=(u.deg, u.deg))
+        # Regression test: the columns of the default should never
+        # be modified from default
+        assert vizier.core.Vizier.columns == ['*']
         result = vizier.core.Vizier.query_region(targets,
                                                  radius=10 * u.arcsec,
                                                  catalog=["HIP", "NOMAD", "UCAC"])
 
         assert len(result) >= 5
         assert 'I/239/hip_main' in result.keys()
+        assert 'HIP' in result['I/239/hip_main'].columns
         assert result['I/239/hip_main']['HIP'] == 98298
 
     def test_findcatalog_maxcatalog(self):

astroquery/vizier/tests/test_vizier_remote.py

@@ -207,23 +207,25 @@ this Vizier instance:
    
     >>> result = v.query_object("HD 226868", catalog=["NOMAD", "UCAC"])              
     >>> print(result)
-    TableList with 2 tables:
-       '0:I/289/out' with 3 column(s) and 18 row(s) 
-       '1:I/322A/out' with 4 column(s) and 10 row(s) 
-
+    TableList with 3 tables:
+        '0:I/297/out' with 3 column(s) and 50 row(s)
+        '1:I/289/out' with 3 column(s) and 18 row(s)
+        '2:I/322A/out' with 3 column(s) and 10 row(s)
+    
     >>> print(result['I/322A/out'])
-     _RAJ2000    DEJ2000    Vmag   _DEJ2000 
-    ---------- ----------- ------ ----------
-    299.572419  35.1942342 15.986  35.194234
-    299.580291  35.1768889 13.274  35.176889
-    299.582571  35.1852253 14.863  35.185225
-    299.594172  35.1799948 14.690  35.179995
-    299.601402  35.1981078 14.644  35.198108
-    299.617669  35.1869987 14.394  35.186999
-    299.561498  35.2016928 15.687  35.201693
-    299.570217  35.2256634 14.878  35.225663
-    299.601081  35.2333378 13.170  35.233338
-    299.617995  35.2058637 13.946  35.205864
+     _RAJ2000   _DEJ2000   Vmag
+       deg        deg      mag
+    ---------- ---------- ------
+    299.572419  35.194234 15.986
+    299.580291  35.176889 13.274
+    299.582571  35.185225 14.863
+    299.594172  35.179995 14.690
+    299.601402  35.198108 14.644
+    299.617669  35.186999 14.394
+    299.561498  35.201693 15.687
+    299.570217  35.225663 14.878
+    299.601081  35.233338 13.170
+    299.617995  35.205864 13.946
 
 When specifying the columns of the query, sorting of the returned table can be
 requested by adding ``+`` (or ``-`` for reverse sorting order) in front of the column