diff --git a/datalab/datalab_session/analysis/get_tif.py b/datalab/datalab_session/analysis/get_tif.py index f96866e..c6f6293 100644 --- a/datalab/datalab_session/analysis/get_tif.py +++ b/datalab/datalab_session/analysis/get_tif.py @@ -19,7 +19,7 @@ def get_tif(input: dict): else: # If tif file doesn't exist, generate a new tif file fits_path = get_fits(basename) - tif_path = create_tif(basename, fits_path) - tif_url = add_file_to_bucket(file_key, tif_path) + with create_tif(basename, fits_path) as tif_path: + tif_url = add_file_to_bucket(file_key, tif_path) return {"tif_url": tif_url} diff --git a/datalab/datalab_session/analysis/line_profile.py b/datalab/datalab_session/analysis/line_profile.py index 45bffff..745fe77 100644 --- a/datalab/datalab_session/analysis/line_profile.py +++ b/datalab/datalab_session/analysis/line_profile.py @@ -3,6 +3,7 @@ from astropy.wcs import WcsError from astropy import coordinates +from datalab.datalab_session.exceptions import ClientAlertException from datalab.datalab_session.utils.file_utils import scale_points, get_hdu from datalab.datalab_session.utils.s3_utils import get_fits @@ -22,7 +23,10 @@ def line_profile(input: dict): """ fits_path = get_fits(input['basename'], input['source']) - sci_hdu = get_hdu(fits_path, 'SCI') + try: + sci_hdu = get_hdu(fits_path, 'SCI') + except TypeError as e: + raise ClientAlertException(f'Error: {e}') x_points, y_points = scale_points(input["height"], input["width"], sci_hdu.data.shape[0], sci_hdu.data.shape[1], x_points=[input["x1"], input["x2"]], y_points=[input["y1"], input["y2"]]) diff --git a/datalab/datalab_session/data_operations/fits_output_handler.py b/datalab/datalab_session/data_operations/fits_output_handler.py index 1e2bbca..1e8a323 100644 --- a/datalab/datalab_session/data_operations/fits_output_handler.py +++ b/datalab/datalab_session/data_operations/fits_output_handler.py @@ -1,7 +1,9 @@ import tempfile +import os import numpy as np from astropy.io import fits +from datalab import settings from datalab.datalab_session.utils.file_utils import create_jpgs from datalab.datalab_session.utils.s3_utils import save_fits_and_thumbnails @@ -52,11 +54,14 @@ def create_and_save_data_products(self, index: int=None, large_jpg_path: str=Non small_jpg (str): Optionally add a path to a small jpg to save, will not create a new jpg. """ hdu_list = fits.HDUList([self.primary_hdu, self.image_hdu]) - fits_output_path = tempfile.NamedTemporaryFile(suffix=f'{self.datalab_id}.fits').name - hdu_list.writeto(fits_output_path, overwrite=True) - - # allow for operations to pregenerate the jpgs, ex. RGB stacking - if not large_jpg_path or not small_jpg_path: - large_jpg_path, small_jpg_path = create_jpgs(self.datalab_id, fits_output_path) + + with tempfile.NamedTemporaryFile(suffix=f'{self.datalab_id}.fits', dir=settings.TEMP_FITS_DIR) as fits_output_file: + fits_output_path = fits_output_file.name + hdu_list.writeto(fits_output_path, overwrite=True) - return save_fits_and_thumbnails(self.datalab_id, fits_output_path, large_jpg_path, small_jpg_path, index) + # allow for operations to pregenerate the jpgs, ex. RGB stacking + if not large_jpg_path or not small_jpg_path: + with create_jpgs(self.datalab_id, fits_output_path) as (large_jpg_path, small_jpg_path): + return save_fits_and_thumbnails(self.datalab_id, fits_output_path, large_jpg_path, small_jpg_path, index) + else: + return save_fits_and_thumbnails(self.datalab_id, fits_output_path, large_jpg_path, small_jpg_path, index) diff --git a/datalab/datalab_session/data_operations/rgb_stack.py b/datalab/datalab_session/data_operations/rgb_stack.py index a26c488..8cbc570 100644 --- a/datalab/datalab_session/data_operations/rgb_stack.py +++ b/datalab/datalab_session/data_operations/rgb_stack.py @@ -133,19 +133,18 @@ def operate(self): aligned_images = self._align_images(fits_files) - large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, aligned_images, color=True, zmin=zmin_list, zmax=zmax_list) - - stacked_ndarray = self._create_3d_array(input_handlers) - - rgb_comment = f'Datalab RGB Stack on files {", ".join(input["basename"] for input in rgb_inputs)}' - output = FITSOutputHandler( - self.cache_key, - stacked_ndarray, - rgb_comment - ).create_and_save_data_products( - large_jpg_path=large_jpg_path, - small_jpg_path=small_jpg_path - ) + with create_jpgs(self.cache_key, aligned_images, color=True, zmin=zmin_list, zmax=zmax_list) as (large_jpg_path, small_jpg_path): + stacked_ndarray = self._create_3d_array(input_handlers) + + rgb_comment = f'Datalab RGB Stack on files {", ".join(input["basename"] for input in rgb_inputs)}' + output = FITSOutputHandler( + self.cache_key, + stacked_ndarray, + rgb_comment + ).create_and_save_data_products( + large_jpg_path=large_jpg_path, + small_jpg_path=small_jpg_path + ) log.info(f'RGB Stack output: {output}') self.set_output(output) diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py index 2743de1..f540a42 100644 --- a/datalab/datalab_session/tests/test_operations.py +++ b/datalab/datalab_session/tests/test_operations.py @@ -177,9 +177,9 @@ def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get # return the test fits paths in order of the input_files instead of aws fetch mock_get_fits.side_effect = [self.test_fits_1_path, self.test_fits_2_path] # save temp output to a known path so we can test it - mock_named_tempfile.return_value.name = self.temp_median_path + mock_named_tempfile.return_value.__enter__.return_value.name = self.temp_median_path # avoids overwriting our output - mock_create_jpgs.return_value = ('test_path', 'test_path') + mock_create_jpgs.return_value.__enter__.return_value = ('test_path', 'test_path') # don't save to s3 mock_save_fits_and_thumbnails.return_value = self.temp_median_path @@ -230,9 +230,9 @@ def test_operate(self, mock_get_fits, mock_named_tempfile, mock_create_jpgs, moc # return the test fits paths in order of the input_files instead of aws fetch mock_get_fits.side_effect = [self.test_red_path, self.test_green_path, self.test_blue_path] # save temp output to a known path so we can test - mock_named_tempfile.return_value.name = self.temp_rgb_path + mock_named_tempfile.return_value.__enter__.return_value.name = self.temp_rgb_path # avoids overwriting our output - mock_create_jpgs.return_value = ('test_path', 'test_path') + mock_create_jpgs.return_value.__enter__.return_value = ('test_path', 'test_path') # don't save to s3 mock_save_fits_and_thumbnails.return_value = self.temp_rgb_path @@ -300,9 +300,9 @@ def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get mock_get_fits.side_effect = [self.test_fits_1_path, self.test_fits_2_path, self.temp_fits_1_negative_path, self.temp_fits_2_negative_path] # save temp output to a known path so we can test it - mock_named_tempfile.return_value.name = self.temp_stacked_path + mock_named_tempfile.return_value.__enter__.return_value.name = self.temp_stacked_path # avoids overwriting our output - mock_create_jpgs.return_value = ('test_path', 'test_path') + mock_create_jpgs.return_value.__enter__.return_value = ('test_path', 'test_path') # don't save to s3 mock_save_fits_and_thumbnails.return_value = self.temp_stacked_path diff --git a/datalab/datalab_session/tests/test_utils.py b/datalab/datalab_session/tests/test_utils.py index 375d13a..404f381 100644 --- a/datalab/datalab_session/tests/test_utils.py +++ b/datalab/datalab_session/tests/test_utils.py @@ -17,37 +17,33 @@ def test_get_fits_dimensions(self): def test_create_fits(self): test_2d_ndarray = np.zeros((10, 10)) - path = create_fits('create_fits_test', test_2d_ndarray) + with create_fits('create_fits_test', test_2d_ndarray) as path: + # test the file was written out to a path + self.assertIsInstance(path, str) + self.assertIsFile(path) + + # test the file has the right data + hdu = fits.open(path) + self.assertEqual(hdu[0].header['KEY'], 'create_fits_test') + self.assertEqual(hdu[1].data.tolist(), test_2d_ndarray.tolist()) - # test the file was written out to a path - self.assertIsInstance(path, str) - self.assertIsFile(path) - - # test the file has the right data - hdu = fits.open(path) - self.assertEqual(hdu[0].header['KEY'], 'create_fits_test') - self.assertEqual(hdu[1].data.tolist(), test_2d_ndarray.tolist()) - def test_create_tif(self): fits_path = self.test_fits_path - tif_path = create_tif('create_tif_test', fits_path) - - # test the file was written out to a path - self.assertIsInstance(tif_path, str) - self.assertIsFile(tif_path) + with create_tif('create_tif_test', fits_path) as tif_path: + # test the file was written out to a path + self.assertIsInstance(tif_path, str) + self.assertIsFile(tif_path) + self.assertFilesEqual(tif_path, self.test_tif_path) - self.assertFilesEqual(tif_path, self.test_tif_path) - def test_create_jpgs(self): fits_path = self.test_fits_path - jpg_paths = create_jpgs('create_jpgs_test', fits_path) - - # test the files were written out to a path - self.assertEqual(len(jpg_paths), 2) - self.assertIsFile(jpg_paths[0]) - self.assertIsFile(jpg_paths[1]) - self.assertFilesEqual(jpg_paths[0], self.test_large_jpg_path) - self.assertFilesEqual(jpg_paths[1], self.test_small_jpg_path) + with create_jpgs('create_jpgs_test', fits_path) as jpg_paths: + # test the files were written out to a path + self.assertEqual(len(jpg_paths), 2) + self.assertIsFile(jpg_paths[0]) + self.assertIsFile(jpg_paths[1]) + self.assertFilesEqual(jpg_paths[0], self.test_large_jpg_path) + self.assertFilesEqual(jpg_paths[1], self.test_small_jpg_path) def test_stack_arrays(self): test_array_1 = np.zeros((10, 20)) diff --git a/datalab/datalab_session/utils/file_utils.py b/datalab/datalab_session/utils/file_utils.py index 32d5ffe..00c3392 100644 --- a/datalab/datalab_session/utils/file_utils.py +++ b/datalab/datalab_session/utils/file_utils.py @@ -1,12 +1,14 @@ import tempfile import logging +import os +from contextlib import contextmanager from astropy.io import fits import numpy as np from fits2image.conversions import fits_to_jpg, fits_to_tif +from datalab import settings from datalab.datalab_session.exceptions import ClientAlertException -from datalab.datalab_session.utils.s3_utils import save_fits_and_thumbnails log = logging.getLogger() log.setLevel(logging.INFO) @@ -29,6 +31,7 @@ def get_fits_dimensions(fits_file, extension: str = 'SCI') -> tuple: hdu_shape = hdu[extension].shape return hdu_shape +@contextmanager def create_fits(key: str, image_arr: np.ndarray, comment=None) -> str: """ Creates a fits file with the given key and image array @@ -41,22 +44,30 @@ def create_fits(key: str, image_arr: np.ndarray, comment=None) -> str: image_hdu = fits.ImageHDU(data=image_arr, name='SCI') hdu_list = fits.HDUList([primary_hdu, image_hdu]) - fits_path = tempfile.NamedTemporaryFile(suffix=f'{key}.fits').name + fits_path = tempfile.NamedTemporaryFile(suffix=f'{key}.fits', dir=settings.TEMP_FITS_DIR).name hdu_list.writeto(fits_path, overwrite=True) - return fits_path + try: + yield fits_path + finally: + os.remove(fits_path) +@contextmanager def create_tif(key: str, fits_path: np.ndarray) -> str: """ Creates a full sized TIFF file from a FITs Returns the path to the TIFF file """ height, width = get_fits_dimensions(fits_path) - tif_path = tempfile.NamedTemporaryFile(suffix=f'{key}.tif').name + tif_path = tempfile.NamedTemporaryFile(suffix=f'{key}.tif', dir=settings.TEMP_FITS_DIR).name fits_to_tif(fits_path, tif_path, width=width, height=height) - return tif_path + try: + yield tif_path + finally: + os.remove(tif_path) +@contextmanager def create_jpgs(cache_key, fits_paths: str, color=False, zmin=None, zmax=None) -> list: """ Create jpgs from fits files and save them to S3 @@ -68,8 +79,8 @@ def create_jpgs(cache_key, fits_paths: str, color=False, zmin=None, zmax=None) - fits_paths = [fits_paths] # create the jpgs from the fits files - large_jpg_path = tempfile.NamedTemporaryFile(suffix=f'{cache_key}-large.jpg').name - thumbnail_jpg_path = tempfile.NamedTemporaryFile(suffix=f'{cache_key}-small.jpg').name + large_jpg_path = tempfile.NamedTemporaryFile(suffix=f'{cache_key}-large.jpg', dir=settings.TEMP_FITS_DIR).name + thumbnail_jpg_path = tempfile.NamedTemporaryFile(suffix=f'{cache_key}-small.jpg', dir=settings.TEMP_FITS_DIR).name max_height = 0 max_width = 0 @@ -81,7 +92,14 @@ def create_jpgs(cache_key, fits_paths: str, color=False, zmin=None, zmax=None) - fits_to_jpg(fits_paths, large_jpg_path, width=max_width, height=max_height, color=color, zmin=zmin, zmax=zmax) fits_to_jpg(fits_paths, thumbnail_jpg_path, color=color, zmin=zmin, zmax=zmax) - return large_jpg_path, thumbnail_jpg_path + try: + yield large_jpg_path, thumbnail_jpg_path + finally: + try: + os.remove(large_jpg_path) + os.remove(thumbnail_jpg_path) + except FileNotFoundError: + pass def crop_arrays(array_list: list): """ @@ -116,16 +134,3 @@ def scale_points(height_1: int, width_1: int, height_2: int, width_2: int, x_poi x_points = width_2 - x_points return x_points, y_points - -def create_output(cache_key, np_array, large_jpg=None, small_jpg=None, index=None, comment=None): - """ - A more automated way of creating output for a dev - Dev can specify just a cache_key and np array and the function will create the fits and jpgs - or the dev can pass the fits_file or jpgs and the function will save them - """ - fits_file = create_fits(cache_key, np_array, comment) - - if not large_jpg or not small_jpg: - large_jpg, small_jpg = create_jpgs(cache_key, fits_file) - - return save_fits_and_thumbnails(cache_key, fits_file, large_jpg, small_jpg, index)