Simplifying code, and making sure both pixel and angular size works

jaymedina · jaymedina · commit 64157c9e4554 · 2023-04-04T17:05:12.000-04:00
diff --git a/astroquery/mast/cutouts.py b/astroquery/mast/cutouts.py
@@ -35,7 +35,7 @@
 __all__ = ["TesscutClass", "Tesscut", "ZcutClass", "Zcut"]
 
 
-def _parse_cutout_size(size, timeout_add=None, mission=None):
+def _parse_cutout_size(size, timeout=None, mission=None):
     """
     Take a user input cutout size and parse it into the regular format
     [ny,nx] where nx/ny are quantities with units either pixels or degrees.
@@ -53,8 +53,8 @@ def _parse_cutout_size(size, timeout_add=None, mission=None):
         The mission for which the size parsing is being done. This parameter
         is mainly meant to trigger a cutout size warning specifically for TESSCut
         requests. Default is None.
-    timeout_add : int or float, optional
-        The amount (in seconds) by which the request processing time upper limit will be changed.
+    timeout : int or float, optional
+        The modified request timeout limit.
         The request processing time by default is 600 seconds, meaning an attempt at communicating
         with the API will take 600 seconds before timing out. In the context of this function, this
         parameter is meant to keep track of whether or not the timeout limit has been modified, which
@@ -72,7 +72,7 @@ def _parse_cutout_size(size, timeout_add=None, mission=None):
     if np.isscalar(size):
         size = np.repeat(size, 2)
 
-        limit_reached = size[0] > 30 or size[1] > 30
+        limit_reached = (size > 30).any()
 
     if isinstance(size, u.Quantity):
         size = np.atleast_1d(size)
@@ -83,9 +83,8 @@ def _parse_cutout_size(size, timeout_add=None, mission=None):
         # Based on the literature, TESS resolution is approx. 21 arcseconds per pixel.
         # We will convert the recommended upper limit for a dimension from pixels
         # to the unit being passed.
-        unit = size[0].unit
-        upper_limit = (30 * 21*u.arcsec).to(unit).value
-        limit_reached = size[0].value > upper_limit or size[1].value > upper_limit
+        with u.set_enabled_equivalencies(u.pixel_scale(21 * u.arcsec / u.pixel)):
+            limit_reached = (size > 30 * u.pixel).any()
 
     if len(size) > 2:
         warnings.warn("Too many dimensions in cutout size, only the first two will be used.",
@@ -111,10 +110,10 @@ def _parse_cutout_size(size, timeout_add=None, mission=None):
     else:
         raise InvalidQueryError("Cutout size must be in pixels or angular quantity.")
 
-    if (mission == 'TESS') & (limit_reached) & (not timeout_add):
+    if (mission == 'TESS') & (limit_reached) & (not timeout):
         warnings.warn("You have selected a large cutout size that may result in a timeout error. We suggest limiting"
                       " the size of your requested cutout, or changing the request timeout limit from its"
-                      " default 600 seconds to something higher, using the timeout_add argument.", InputWarning)
+                      " default 600 seconds to something higher, using the timeout argument.", InputWarning)
 
     return {"x": x, "y": y, "units": units}
 
@@ -137,6 +136,11 @@ def __init__(self):
                     }
         self._service_api_connection.set_service_params(services, "tesscut")
 
+    def _get_default_timeout(self):
+        """ Gets the default request timeout limit. """
+
+        return self._service_api_connection.TIMEOUT
+
     def get_sectors(self, *, coordinates=None, radius=0*u.deg, product='SPOC', objectname=None,
                     moving_target=False, mt_type=None):
 
@@ -256,7 +260,7 @@ def get_sectors(self, *, coordinates=None, radius=0*u.deg, product='SPOC', objec
 
     def download_cutouts(self, *, coordinates=None, size=5, sector=None, product='SPOC', path=".",
                          inflate=True, objectname=None, moving_target=False, mt_type=None, verbose=False,
-                         timeout_add=None):
+                         timeout=None):
         """
         Download cutout target pixel file(s) around the given coordinates with indicated size.
 
@@ -313,12 +317,23 @@ def download_cutouts(self, *, coordinates=None, size=5, sector=None, product='SP
             first majorbody is tried and then smallbody if a matching majorbody is not found.
 
             NOTE: If moving_target is supplied, this argument is ignored.
+        timeout : int or float, optional
+            The modified request timeout limit.
+            The request processing time by default is 600 seconds, meaning an attempt at communicating
+            with the API will take 600 seconds before timing out. The timeout upper limit can be modified
+            using this argument for large cutout requests via TESSCut. Default is None.
 
         Returns
         -------
         response : `~astropy.table.Table`
         """
 
+        # Modify TIMEOUT attribute if necessary (usually this is modified for large requests)
+        if timeout:
+            self._service_api_connection.TIMEOUT = timeout
+            log.info(f"Request timeout upper limit is being changed to {self._service_api_connection.TIMEOUT}"
+                     " seconds.")
+
         if moving_target:
 
             # The Moving Targets service is currently only available for SPOC
@@ -348,7 +363,7 @@ def download_cutouts(self, *, coordinates=None, size=5, sector=None, product='SP
             astrocut_request = f"astrocut?ra={coordinates.ra.deg}&dec={coordinates.dec.deg}"
 
         # Adding the arguments that are common between moving/still astrocut requests
-        size_dict = _parse_cutout_size(size, timeout_add=timeout_add, mission='TESS')
+        size_dict = _parse_cutout_size(size, timeout=timeout, mission='TESS')
         astrocut_request += f"&y={size_dict['y']}&x={size_dict['x']}&units={size_dict['units']}"
 
         # Making sure input product is either SPOC or TICA,
@@ -389,10 +404,14 @@ def download_cutouts(self, *, coordinates=None, size=5, sector=None, product='SP
         os.remove(zipfile_path)
 
         localpath_table['Local Path'] = [path+x for x in cutout_files]
+
+        if timeout:
+            self._service_api_connection.TIMEOUT = self._get_default_timeout()
+
         return localpath_table
 
     def get_cutouts(self, *, coordinates=None, size=5, product='SPOC', sector=None,
-                    objectname=None, moving_target=False, mt_type=None, timeout_add=None):
+                    objectname=None, moving_target=False, mt_type=None, timeout=None):
         """
         Get cutout target pixel file(s) around the given coordinates with indicated size,
         and return them as a list of  `~astropy.io.fits.HDUList` objects.
@@ -441,8 +460,8 @@ def get_cutouts(self, *, coordinates=None, size=5, product='SPOC', sector=None,
             first majorbody is tried and then smallbody if a matching majorbody is not found.
 
             NOTE: If moving_target is supplied, this argument is ignored.
-        timeout_add : int or float, optional
-            The amount (in seconds) by which the request processing time upper limit will be changed.
+        timeout : int or float, optional
+            The modified request timeout limit.
             The request processing time by default is 600 seconds, meaning an attempt at communicating
             with the API will take 600 seconds before timing out. The timeout upper limit can be modified
             using this argument for large cutout requests via TESSCut. Default is None.
@@ -453,13 +472,13 @@ def get_cutouts(self, *, coordinates=None, size=5, product='SPOC', sector=None,
         """
 
         # Modify TIMEOUT attribute if necessary (usually this is modified for large requests)
-        if timeout_add:
-            self._service_api_connection.TIMEOUT = self._service_api_connection.TIMEOUT + timeout_add
+        if timeout:
+            self._service_api_connection.TIMEOUT = timeout
             log.info(f"Request timeout upper limit is being changed to {self._service_api_connection.TIMEOUT}"
                      " seconds.")
 
         # Setting up the cutout size
-        param_dict = _parse_cutout_size(size, timeout_add=timeout_add, mission='TESS')
+        param_dict = _parse_cutout_size(size, timeout=timeout, mission='TESS')
 
         # Add sector if present
         if sector:
@@ -529,6 +548,9 @@ def get_cutouts(self, *, coordinates=None, size=5, product='SPOC', sector=None,
             # preserve the original filename in the fits object
             cutout_hdus_list[-1].filename = name
 
+        if timeout:
+            self._service_api_connection.TIMEOUT = self._get_default_timeout()
+
         return cutout_hdus_list
 
 
@@ -592,7 +614,7 @@ def get_surveys(self, coordinates, *, radius="0d"):
         return survey_json
 
     def download_cutouts(self, coordinates, *, size=5, survey=None, cutout_format="fits", path=".", inflate=True,
-                         verbose=False, timeout_add=None, **img_params):
+                         verbose=False, **img_params):
         """
         Download cutout FITS/image file(s) around the given coordinates with indicated size.
 
@@ -633,24 +655,13 @@ def download_cutouts(self, coordinates, *, size=5, survey=None, cutout_format="f
             The Column Name is the keyword, with the argument being one or more acceptable
             values for that parameter, except for fields with a float datatype where the
             argument should be in the form [minVal, maxVal].
-        timeout_add : int or float, optional
-            The amount (in seconds) by which the request processing time upper limit will be changed.
-            The request processing time by default is 600 seconds, meaning an attempt at communicating
-            with the API will take 600 seconds before timing out. The timeout upper limit can be modified
-            using this argument for large cutout requests via TESSCut. Default is None.
 
         Returns
         -------
         response : `~astropy.table.Table`
             Cutout file(s) for given coordinates
         """
 
-        # Modify TIMEOUT attribute if necessary (usually this is modified for large requests)
-        if timeout_add:
-            self._service_api_connection.TIMEOUT = self._service_api_connection.TIMEOUT + timeout_add
-            log.info(f"Request timeout upper limit is being changed to {self._service_api_connection.TIMEOUT}"
-                     " seconds.")
-
         # Get Skycoord object for coordinates/object
         coordinates = parse_input_location(coordinates)
         size_dict = _parse_cutout_size(size)
diff --git a/docs/mast/mast.rst b/docs/mast/mast.rst
@@ -1060,10 +1060,10 @@ around 34 x 34 pixels-squared.
    >>> hdulist = Tesscut.get_cutouts(coordinates=cutout_coord, size=0.2*u.deg)
    WARNING: InputWarning: You have selected a large cutout size that may result in a timeout error.
    We suggest limiting the size of your requested cutout, or changing the request timeout limit from
-   its default 600 seconds to something higher, using the timeout_add argument. [astroquery.mast.cutouts]
+   its default 600 seconds to something higher, using the timeout argument. [astroquery.mast.cutouts]
    
 At this point, users may choose to decrease their cutout size or extend the request timeout limit from
-600 seconds to something longer. Using the example above, we will use the ``timeout_add`` argument to extend
+600 seconds to something longer. Using the example above, we will use the ``timeout`` argument to change
 the request timeout limit from 600 seconds to 2600 seconds, or approximately 43 minutes.
 
 .. doctest-remote-data::
@@ -1073,7 +1073,7 @@ the request timeout limit from 600 seconds to 2600 seconds, or approximately 43
    >>> from astropy.coordinates import SkyCoord
    ...
    >>> cutout_coord = SkyCoord(107.18696, -70.50919, unit="deg")
-   >>> hdulist = Tesscut.get_cutouts(coordinates=cutout_coord, size=0.2*u.deg, timeout_add=2000)
+   >>> hdulist = Tesscut.get_cutouts(coordinates=cutout_coord, size=0.2*u.deg, timeout=2600)
    INFO: Request timeout upper limit is being changed to 2600 seconds. [astroquery.mast.cutouts]
 
 Sector information
