Add method write_image in HipsDrawResult class

hips/draw/simple.py

@@ -1,6 +1,8 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 """HiPS tile drawing -- simple method."""
 import numpy as np
+from PIL import Image
+from astropy.io import fits
 from typing import List, Tuple
 from astropy.wcs.utils import proj_plane_pixel_scales
 from skimage.transform import ProjectiveTransform, warp
@@ -127,7 +129,7 @@ def tiles(self) -> List[HipsTile]:
 
     @property
     def result(self) -> 'HipsDrawResult':
-        return HipsDrawResult(self.image, self.geometry)
+        return HipsDrawResult(self.image, self.geometry, self.tile_format)
 
     def warp_image(self, tile: HipsTile) -> np.ndarray:
         """Warp a HiPS tile and a sky image."""
@@ -177,12 +179,51 @@ def plot_mpl_hips_tile_grid(self) -> None:
 
 
 class HipsDrawResult:
-    """Container class for reporting information related with fetching / drawing of HiPS tiles."""
+    """Container class for reporting information related with fetching / drawing of HiPS tiles.
 
-    def __init__(self, image: np.ndarray, geometry: WCSGeometry) -> None:
+    Parameters
+    ----------
+    image: `~numpy.ndarray`
+        Container for HiPS tile data
+    geometry : `~hips.utils.WCSGeometry`
+        An object of WCSGeometry
+    tile_format : {'fits', 'jpg', 'png'}
+        Format of HiPS tile
+    """
+
+    def __init__(self, image: np.ndarray, geometry: WCSGeometry, tile_format: str) -> None:
         self.image = image
         self.geometry = geometry
+        self.tile_format = tile_format
+
+    def __str__(self):
+        return (
+            'This class object contains two attributes: image and geometry'
+        )
 
+    def __repr__(self):
+        return (
+            f'width={self.image.shape[0]}, '
+            f'height={self.image.shape[1]}, '
+            f'channels={self.image.ndim}, '
+            f'dtype={self.image.dtype}, '
+            f'format={self.tile_format}'
+        )
+
+    def write_image(self, filename: str) -> None:
+        """Write image to file.
+
+        Parameters
+        ----------
+        filename : str
+            Filename
+        """
+        if self.tile_format == 'fits':
+            hdu = fits.PrimaryHDU(self.image)
+            hdu.writeto(filename)
+        else:
+            image = Image.fromarray(self.image)
+            image.save(filename)
 
 def measure_tile_shape(corners: tuple) -> Tuple[List[float]]:
     """Compute length of tile edges and diagonals."""

hips/draw/tests/test_simple.py

@@ -19,6 +19,7 @@
         data_2=2296,
         data_sum=8756493140,
         dtype='>i2',
+        repr='width=1000, height=2000, channels=2, dtype=>i2, format=fits'
     ),
     dict(
         file_format='jpg',
@@ -28,6 +29,7 @@
         data_2=[137, 116, 114],
         data_sum=828908873,
         dtype='uint8',
+        repr='width=1000, height=2000, channels=3, dtype=uint8, format=jpg'
     ),
     dict(
         file_format='png',
@@ -37,6 +39,7 @@
         data_2=[227, 217, 205, 255],
         data_sum=1635622838,
         dtype='uint8',
+        repr='width=1000, height=2000, channels=3, dtype=uint8, format=png'
     ),
 ]
 
@@ -52,7 +55,8 @@ def test_make_sky_image(pars):
     assert_allclose(np.sum(result.image), pars['data_sum'])
     assert_allclose(result.image[200, 994], pars['data_1'])
     assert_allclose(result.image[200, 995], pars['data_2'])
-
+    # result.write_image('test.' + pars['file_format'])
+    assert repr(result) == pars['repr']
 
 @remote_data
 class TestSimpleTilePainter: