Merge branch 'frame-support'

picamera/streams.py

@@ -713,16 +713,25 @@ def _find_seconds(self, seconds, first_frame):
                     break
         return pos
 
+    def _find_frames(self, frames, first_frame):
+        pos = None
+        for count, frame in enumerate(reversed(self.frames)):
+            if first_frame in (None, frame.frame_type):
+                pos = frame.position
+            if frames < count:
+                break
+        return pos
+
     def _find_all(self, first_frame):
         for frame in self.frames:
             if first_frame in (None, frame.frame_type):
                 return frame.position
 
     def copy_to(
-            self, output, size=None, seconds=None,
+            self, output, size=None, seconds=None, frames=None,
             first_frame=PiVideoFrameType.sps_header):
         """
-        copy_to(output, size=None, seconds=None, first_frame=PiVideoFrameType.sps_header)
+        copy_to(output, size=None, seconds=None, frames=None, first_frame=PiVideoFrameType.sps_header)
 
         Copies content from the stream to *output*.
 
@@ -732,8 +741,9 @@ def copy_to(
         If *size* is specified then the copy will be limited to the whole
         number of frames that fit within the specified number of bytes. If
         *seconds* if specified, then the copy will be limited to that number of
-        seconds worth of frames. Only one of *size* or *seconds* can be
-        specified.  If neither is specified, all frames are copied.
+        seconds worth of frames. If *frames* is specified then the copy will
+        be limited to that number of frames. Only one of *size*, *seconds*, or
+        *frames* can be specified.  If none is specified, all frames are copied.
 
         If *first_frame* is specified, it defines the frame type of the first
         frame to be copied. By default this is
@@ -744,13 +754,14 @@ def copy_to(
         .. warning::
 
             Note that if a frame of the specified type (e.g. SPS header) cannot
-            be found within the specified number of seconds or bytes then this
-            method will simply copy nothing (but no error will be raised).
+            be found within the specified number of seconds, bytes, or frames,
+            then this method will simply copy nothing (but no error will be
+            raised).
 
         The stream's position is not affected by this method.
         """
-        if size is not None and seconds is not None:
-            raise PiCameraValueError('You cannot specify both size and seconds')
+        if (size, seconds, frames).count(None) < 2:
+            raise PiCameraValueError('You can only specify one of size, seconds, or frames')
         if isinstance(output, bytes):
             output = output.decode('utf-8')
         opened = isinstance(output, str)
@@ -764,6 +775,8 @@ def copy_to(
                         pos = self._find_size(size, first_frame)
                     elif seconds is not None:
                         pos = self._find_seconds(seconds, first_frame)
+                    elif frames is not None:
+                        pos = self._find_frames(frames, first_frame)
                     else:
                         pos = self._find_all(first_frame)
                     # Copy chunks efficiently from the position found

tests/test_streams.py

@@ -331,3 +331,16 @@ def test_camera_stream_copy_seconds():
     stream.copy_to(output, seconds=10)
     assert output.getvalue() == b'hkkffkkff'
 
+def test_camera_stream_copy_frames():
+    camera = mock.Mock()
+    encoder = mock.Mock()
+    camera._encoders = {1: encoder}
+    stream = PiCameraCircularIO(camera, size=10)
+    for data, frame in generate_frames('hkffkff'):
+        encoder.frame = frame
+        stream.write(data)
+    output = io.BytesIO()
+    stream.copy_to(output, frames=1)
+    assert output.getvalue() == b''
+    stream.copy_to(output, frames=10)
+    assert output.getvalue() == b'hkkffkkff'