Merge pull request #5405 from radarhere/fits

Read FITS image mode and size

Tests/test_file_fitsstub.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+
 import pytest
 
 from PIL import FitsStubImagePlugin, Image
@@ -11,10 +13,8 @@ def test_open():
 
         # Assert
         assert im.format == "FITS"
-
+        assert im.size == (128, 128)
-        # Dummy data from the stub
+        assert im.mode == "L"
-        assert im.mode == "F"
-        assert im.size == (1, 1)
 
 
 def test_invalid_file():
@@ -35,6 +35,21 @@ def test_load():
             im.load()
 
 
+def test_truncated_fits():
+    # No END to headers
+    image_data = b"SIMPLE  =                    T" + b" " * 50 + b"TRUNCATE"
+    with pytest.raises(OSError):
+        FitsStubImagePlugin.FITSStubImageFile(BytesIO(image_data))
+
+
+def test_naxis_zero():
+    # This test image has been manually hexedited
+    # to set the number of data axes to zero
+    with pytest.raises(ValueError):
+        with Image.open("Tests/images/hopper_naxis_zero.fits"):
+            pass
+
+
 def test_save():
     # Arrange
     with Image.open(TEST_FILE) as im:

src/PIL/FitsStubImagePlugin.py

@@ -38,21 +38,45 @@ class FITSStubImageFile(ImageFile.StubImageFile):
     format_description = "FITS"
 
     def _open(self):
-
         offset = self.fp.tell()
 
-        if not _accept(self.fp.read(6)):
+        headers = {}
-            raise SyntaxError("Not a FITS file")
+        while True:
+            header = self.fp.read(80)
+            if not header:
+                raise OSError("Truncated FITS file")
+            keyword = header[:8].strip()
+            if keyword == b"END":
+                break
+            value = header[8:].strip()
+            if value.startswith(b"="):
+                value = value[1:].strip()
+            if not headers and (not _accept(keyword) or value != b"T"):
+                raise SyntaxError("Not a FITS file")
+            headers[keyword] = value
 
-        # FIXME: add more sanity checks here; mandatory header items
+        naxis = int(headers[b"NAXIS"])
-        # include SIMPLE, BITPIX, NAXIS, etc.
+        if naxis == 0:
+            raise ValueError("No image data")
+        elif naxis == 1:
+            self._size = 1, int(headers[b"NAXIS1"])
+        else:
+            self._size = int(headers[b"NAXIS1"]), int(headers[b"NAXIS2"])
+
+        number_of_bits = int(headers[b"BITPIX"])
+        if number_of_bits == 8:
+            self.mode = "L"
+        elif number_of_bits == 16:
+            self.mode = "I"
+            # rawmode = "I;16S"
+        elif number_of_bits == 32:
+            self.mode = "I"
+        elif number_of_bits in (-32, -64):
+            self.mode = "F"
+            # rawmode = "F" if number_of_bits == -32 else "F;64F"
 
         self.fp.seek(offset)
 
-        # make something up
-        self.mode = "F"
-        self._size = 1, 1
-
         loader = self._load()
         if loader:
             loader.open(self)