Coverage for python/lsst/afw/display/_write_fits.py: 20%

1# This file is part of afw. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ["writeFitsImage"] 

25 

26import io 

27import os 

28import subprocess 

29 

30import lsst.afw.fits 

31import lsst.afw.geom 

32import lsst.afw.image 

33from lsst.daf.base import PropertyList, PropertySet 

34from lsst.geom import Extent2D 

35 

36 

37def _add_wcs(wcs_name: str, ps: PropertyList, x0: int = 0, y0: int = 0) -> None: 

38 ps.setInt(f"CRVAL1{wcs_name}", x0, "(output) Column pixel of Reference Pixel") 

39 ps.setInt(f"CRVAL2{wcs_name}", y0, "(output) Row pixel of Reference Pixel") 

40 ps.setDouble(f"CRPIX1{wcs_name}", 1.0, "Column Pixel Coordinate of Reference") 

41 ps.setDouble(f"CRPIX2{wcs_name}", 1.0, "Row Pixel Coordinate of Reference") 

42 ps.setString(f"CTYPE1{wcs_name}", "LINEAR", "Type of projection") 

43 ps.setString(f"CTYPE1{wcs_name}", "LINEAR", "Type of projection") 

44 ps.setString(f"CUNIT1{wcs_name}", "PIXEL", "Column unit") 

45 ps.setString(f"CUNIT2{wcs_name}", "PIXEL", "Row unit") 

46 

47 

48def writeFitsImage( 

49 file: str | int | io.BytesIO, 

50 data: lsst.afw.image.Image | lsst.afw.image.Mask, 

51 wcs: lsst.afw.geom.SkyWcs | None = None, 

52 title: str = "", 

53 metadata: PropertySet | None = None, 

54) -> None: 

55 """Write a simple FITS file with no extensions. 

56 

57 Parameters 

58 ---------- 

59 file : `str` or `int` 

60 Path to a file or a file descriptor. 

61 data : `lsst.afw.Image` or `lsst.afw.Mask` 

62 Data to be displayed. 

63 wcs : `lsst.afw.geom.SkyWcs` or `None`, optional 

64 WCS to be written to header to FITS file. 

65 title : `str`, optional 

66 If defined, the value to be stored in the ``OBJECT`` header. 

67 Overrides any value found in ``metadata``. 

68 metadata : `lsst.daf.base.PropertySet` or `None`, optional 

69 Additional information to be written to FITS header. 

70 """ 

71 ps = PropertyList() 

72 

73 # Seed with the external metadata, stripping wcs keywords. 

74 if metadata: 

75 lsst.afw.geom.stripWcsMetadata(metadata) 

76 ps.update(metadata) 

77 

78 # Write WcsB, so that pixel (0,0) is correctly labelled (but ignoring XY0) 

79 _add_wcs("B", ps) 

80 

81 if not wcs: 

82 _add_wcs("", ps) # Works around a ds9 bug that WCSA/B is ignored if no WCS is present. 

83 else: 

84 shift = Extent2D(-data.getX0(), -data.getY0()) 

85 if wcs.hasFitsApproximation(): 

86 wcs = wcs.getFitsApproximation() 

87 new_wcs = wcs.copyAtShiftedPixelOrigin(shift) 

88 wcs_metadata = new_wcs.getFitsMetadata() 

89 ps.update(wcs_metadata) 

90 

91 if title: 

92 ps.set("OBJECT", title, "Image being displayed") 

93 

94 if isinstance(file, str): 

95 data.writeFits(file, metadata=ps) 

96 else: 

97 mem = lsst.afw.fits.MemFileManager() 

98 data.writeFits(manager=mem, metadata=ps) 

99 if isinstance(file, int): 

100 # Duplicate to prevent a double close, assuming the caller 

101 # will close the file descriptor they passed in. 

102 with os.fdopen(os.dup(file), "wb") as fh: 

103 fh.write(mem.getData()) 

104 elif isinstance(file, subprocess.Popen): 

105 file.communicate(input=mem.getData()) 

106 else: 

107 # Try the write() method directly. 

108 file.write(mem.getData())