Move shot_select from postprocess to imread/open_file/open_buffer

Co-authored-by: letmaik <530988+letmaik@users.noreply.github.com>

Author: Copilot

rawpy/__init__.py

@@ -5,17 +5,19 @@
 import rawpy._rawpy
 globals().update({k:v for k,v in rawpy._rawpy.__dict__.items() if not k.startswith('_')})
 
-def imread(pathOrFile):
+def imread(pathOrFile, shot_select=0):
     """
     Convenience function that creates a :class:`rawpy.RawPy` instance, opens the given file,
     and returns the :class:`rawpy.RawPy` instance for further processing.
     
     :param str|file pathOrFile: path or file object of RAW image that will be read
+    :param int shot_select: select which image to extract from RAW files that contain multiple images
+                            (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
     :rtype: :class:`rawpy.RawPy`
     """
     d = RawPy()
     if hasattr(pathOrFile, 'read'):
-        d.open_buffer(pathOrFile)
+        d.open_buffer(pathOrFile, shot_select=shot_select)
     else:
-        d.open_file(pathOrFile)
+        d.open_file(pathOrFile, shot_select=shot_select)
     return d

rawpy/_rawpy.pyx

@@ -411,13 +411,15 @@ cdef class RawPy:
         with nogil:
             self.p.recycle()
 
-    def open_file(self, path):
+    def open_file(self, path, shot_select=0):
         """
         Opens the given RAW image file. Should be followed by a call to :meth:`~rawpy.RawPy.unpack`.
         
         .. NOTE:: This is a low-level method, consider using :func:`rawpy.imread` instead.
         
         :param str path: The path to the RAW image.
+        :param int shot_select: select which image to extract from RAW files that contain multiple images
+                                (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
         """
         cdef wchar_t *wchars
         cdef Py_ssize_t wchars_len
@@ -433,14 +435,19 @@ cdef class RawPy:
         ELSE:
             res = self.p.open_file(path.encode('UTF-8'))
         self.handle_error(res)
+        # Set shot_select after opening file
+        cdef libraw_raw_unpack_params_t* rp = &self.p.imgdata.rawparams
+        rp.shot_select = shot_select
 
-    def open_buffer(self, fileobj):
+    def open_buffer(self, fileobj, shot_select=0):
         """
         Opens the given RAW image file-like object. Should be followed by a call to :meth:`~rawpy.RawPy.unpack`.
         
         .. NOTE:: This is a low-level method, consider using :func:`rawpy.imread` instead.
         
         :param file fileobj: The file-like object.
+        :param int shot_select: select which image to extract from RAW files that contain multiple images
+                                (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
         """
         self.unpack_called = False
         self.unpack_thumb_called = False
@@ -451,6 +458,9 @@ cdef class RawPy:
         with nogil:
             e = self.p.open_buffer(buf, buf_len)
         self.handle_error(e)
+        # Set shot_select after opening buffer
+        cdef libraw_raw_unpack_params_t* rp = &self.p.imgdata.rawparams
+        rp.shot_select = shot_select
 
     def unpack(self):
         """
@@ -971,10 +981,6 @@ cdef class RawPy:
         p.gamm[1] = params.gamm[1]
         p.aber[0] = params.aber[0]
         p.aber[2] = params.aber[1]
-        
-        # shot_select is in rawparams, not params
-        cdef libraw_raw_unpack_params_t* rp = &self.p.imgdata.rawparams
-        rp.shot_select = params.shot_select
 
     cdef handle_error(self, int code):
         if code > 0:
@@ -1122,7 +1128,7 @@ class Params(object):
                  bright=1.0, highlight_mode=HighlightMode.Clip,
                  exp_shift=None, exp_preserve_highlights=0.0, no_auto_scale=False,
                  gamma=None, chromatic_aberration=None,
-                 bad_pixels_path=None, shot_select=0):
+                 bad_pixels_path=None):
         """
 
         If use_camera_wb and use_auto_wb are False and user_wb is None, then
@@ -1165,8 +1171,6 @@ class Params(object):
         :param str bad_pixels_path: path to dcraw bad pixels file. Each bad pixel will be corrected using
                                     the mean of the neighbor pixels. See the :mod:`rawpy.enhance` module
                                     for alternative repair algorithms, e.g. using the median.
-        :param int shot_select: select which image to extract from RAW files that contain multiple images
-                                (e.g., Dual Pixel RAW). Default is 0 for the first/main image.
         """
 
         if demosaic_algorithm:
@@ -1227,7 +1231,6 @@ class Params(object):
         else:
             self.aber = (1, 1)
         self.bad_pixels = bad_pixels_path
-        self.shot_select = shot_select
 
 cdef class processed_image_wrapper:
     cdef RawPy raw

test/test_shot_select.py

@@ -9,39 +9,37 @@
 rawTestPath = os.path.join(thisDir, 'iss030e122639.NEF')
 
 
-def test_shot_select_parameter_exists():
-    """Test that shot_select parameter can be passed to postprocess"""
-    with rawpy.imread(rawTestPath) as raw:
-        # Test default shot_select=0
-        rgb = raw.postprocess(shot_select=0, no_auto_bright=True)
+def test_shot_select_parameter_via_imread():
+    """Test that shot_select parameter can be passed to imread"""
+    # Test default shot_select=0
+    with rawpy.imread(rawTestPath, shot_select=0) as raw:
+        rgb = raw.postprocess(no_auto_bright=True)
         assert rgb is not None
         assert rgb.shape[2] == 3  # RGB image
 
-def test_shot_select_different_values():
-    """Test that different shot_select values can be set"""
-    with rawpy.imread(rawTestPath) as raw:
-        # Test shot_select=0
-        rgb0 = raw.postprocess(shot_select=0, no_auto_bright=True)
-        assert rgb0 is not None
-        
-        # Test shot_select=1 (may not have effect on single-shot files)
-        rgb1 = raw.postprocess(shot_select=1, no_auto_bright=True)
-        assert rgb1 is not None
-        
-        # Both should produce valid images (same for single-shot files)
-        assert rgb0.shape == rgb1.shape
-
-
-def test_shot_select_via_params():
-    """Test that shot_select can be passed via Params object"""
-    with rawpy.imread(rawTestPath) as raw:
-        params = rawpy.Params(shot_select=0)
-        rgb = raw.postprocess(params=params)
-        assert rgb is not None
-        assert rgb.shape[2] == 3
+
+def test_shot_select_nonexistent_image():
+    """Test that shot_select=1 raises error for single-image files"""
+    # Test shot_select=1 on a file with only one image should raise an error
+    with pytest.raises(rawpy.LibRawRequestForNonexistentImageError):
+        with rawpy.imread(rawTestPath, shot_select=1) as raw:
+            rgb = raw.postprocess(no_auto_bright=True)
+
+
+def test_shot_select_via_open_file():
+    """Test that shot_select can be passed via open_file"""
+    raw = rawpy.RawPy()
+    raw.open_file(rawTestPath, shot_select=0)
+    raw.unpack()
+    rgb = raw.postprocess()
+    assert rgb is not None
+    assert rgb.shape[2] == 3
+    raw.close()
 
 
 def test_shot_select_default_value():
     """Test that default shot_select value is 0"""
-    params = rawpy.Params()
-    assert params.shot_select == 0
+    with rawpy.imread(rawTestPath) as raw:
+        rgb = raw.postprocess()
+        assert rgb is not None
+        assert rgb.shape[2] == 3