From 42658ecd92753ffad33eb5bebe2f4ab085f283f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Mon, 2 Dec 2024 14:49:28 -0800
Subject: [PATCH] frames: use mixins for audio and image data

---
 src/pipecat/frames/frames.py | 48 +++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py
index d10a1cd21..b99a7307c 100644
--- a/src/pipecat/frames/frames.py
+++ b/src/pipecat/frames/frames.py
@@ -67,22 +67,42 @@ class ControlFrame(Frame):
     pass
 
 
+#
+# Mixins
+#
+
+
+@dataclass
+class AudioRawFrame:
+    """A chunk of audio."""
+
+    audio: bytes
+    sample_rate: int
+    num_channels: int
+    num_frames: int = field(init=False)
+
+
+@dataclass
+class ImageRawFrame:
+    """A raw image."""
+
+    image: bytes
+    size: Tuple[int, int]
+    format: str | None
+
+
 #
 # Data frames.
 #
 
 
 @dataclass
-class OutputAudioRawFrame(DataFrame):
+class OutputAudioRawFrame(DataFrame, AudioRawFrame):
     """A chunk of audio. Will be played by the output transport if the
     transport's microphone has been enabled.
 
     """
 
-    audio: bytes
-    sample_rate: int
-    num_channels: int
-
     def __post_init__(self):
         super().__post_init__()
         self.num_frames = int(len(self.audio) / (self.num_channels * 2))
@@ -93,16 +113,12 @@ class OutputAudioRawFrame(DataFrame):
 
 
 @dataclass
-class OutputImageRawFrame(DataFrame):
+class OutputImageRawFrame(DataFrame, ImageRawFrame):
     """An image that will be shown by the transport if the transport's camera is
     enabled.
 
     """
 
-    image: bytes
-    size: Tuple[int, int]
-    format: str | None
-
     def __str__(self):
         pts = format_pts(self.pts)
         return f"{self.name}(pts: {pts}, size: {self.size}, format: {self.format})"
@@ -474,13 +490,9 @@ class UserImageRequestFrame(SystemFrame):
 
 
 @dataclass
-class InputAudioRawFrame(SystemFrame):
+class InputAudioRawFrame(SystemFrame, AudioRawFrame):
     """A chunk of audio usually coming from an input transport."""
 
-    audio: bytes
-    sample_rate: int
-    num_channels: int
-
     def __post_init__(self):
         super().__post_init__()
         self.num_frames = int(len(self.audio) / (self.num_channels * 2))
@@ -491,13 +503,9 @@ class InputAudioRawFrame(SystemFrame):
 
 
 @dataclass
-class InputImageRawFrame(SystemFrame):
+class InputImageRawFrame(SystemFrame, ImageRawFrame):
     """An image usually coming from an input transport."""
 
-    image: bytes
-    size: Tuple[int, int]
-    format: str | None
-
     def __str__(self):
         pts = format_pts(self.pts)
         return f"{self.name}(pts: {pts}, size: {self.size}, format: {self.format})"