feat: Add get_data_parts() and get_file_parts() helper methods (#312)

Fixes #311 🦕

Author: holtskinner

src/a2a/utils/__init__.py

@@ -17,6 +17,8 @@
     create_task_obj,
 )
 from a2a.utils.message import (
+    get_data_parts,
+    get_file_parts,
     get_message_text,
     get_text_parts,
     new_agent_parts_message,
@@ -37,6 +39,8 @@
     'build_text_artifact',
     'completed_task',
     'create_task_obj',
+    'get_data_parts',
+    'get_file_parts',
     'get_message_text',
     'get_text_parts',
     'new_agent_parts_message',

src/a2a/utils/message.py

@@ -2,7 +2,13 @@
 
 import uuid
 
+from typing import Any
+
 from a2a.types import (
+    DataPart,
+    FilePart,
+    FileWithBytes,
+    FileWithUri,
     Message,
     Part,
     Role,
@@ -70,6 +76,30 @@ def get_text_parts(parts: list[Part]) -> list[str]:
     return [part.root.text for part in parts if isinstance(part.root, TextPart)]
 
 
+def get_data_parts(parts: list[Part]) -> list[dict[str, Any]]:
+    """Extracts dictionary data from all DataPart objects in a list of Parts.
+
+    Args:
+        parts: A list of `Part` objects.
+
+    Returns:
+        A list of dictionaries containing the data from any `DataPart` objects found.
+    """
+    return [part.root.data for part in parts if isinstance(part.root, DataPart)]
+
+
+def get_file_parts(parts: list[Part]) -> list[FileWithBytes | FileWithUri]:
+    """Extracts file data from all FilePart objects in a list of Parts.
+
+    Args:
+        parts: A list of `Part` objects.
+
+    Returns:
+        A list of `FileWithBytes` or `FileWithUri` objects containing the file data from any `FilePart` objects found.
+    """
+    return [part.root.file for part in parts if isinstance(part.root, FilePart)]
+
+
 def get_message_text(message: Message, delimiter: str = '\n') -> str:
     """Extracts and joins all text content from a Message's parts.
 

tests/utils/test_message.py

@@ -4,12 +4,17 @@
 
 from a2a.types import (
     DataPart,
+    FilePart,
+    FileWithBytes,
+    FileWithUri,
     Message,
     Part,
     Role,
     TextPart,
 )
 from a2a.utils.message import (
+    get_data_parts,
+    get_file_parts,
     get_message_text,
     get_text_parts,
     new_agent_parts_message,
@@ -178,6 +183,141 @@ def test_get_text_parts_empty_list(self):
         assert result == []
 
 
+class TestGetDataParts:
+    def test_get_data_parts_single_data_part(self):
+        # Setup
+        parts = [Part(root=DataPart(data={'key': 'value'}))]
+
+        # Exercise
+        result = get_data_parts(parts)
+
+        # Verify
+        assert result == [{'key': 'value'}]
+
+    def test_get_data_parts_multiple_data_parts(self):
+        # Setup
+        parts = [
+            Part(root=DataPart(data={'key1': 'value1'})),
+            Part(root=DataPart(data={'key2': 'value2'})),
+        ]
+
+        # Exercise
+        result = get_data_parts(parts)
+
+        # Verify
+        assert result == [{'key1': 'value1'}, {'key2': 'value2'}]
+
+    def test_get_data_parts_mixed_parts(self):
+        # Setup
+        parts = [
+            Part(root=TextPart(text='some text')),
+            Part(root=DataPart(data={'key1': 'value1'})),
+            Part(root=DataPart(data={'key2': 'value2'})),
+        ]
+
+        # Exercise
+        result = get_data_parts(parts)
+
+        # Verify
+        assert result == [{'key1': 'value1'}, {'key2': 'value2'}]
+
+    def test_get_data_parts_no_data_parts(self):
+        # Setup
+        parts = [
+            Part(root=TextPart(text='some text')),
+        ]
+
+        # Exercise
+        result = get_data_parts(parts)
+
+        # Verify
+        assert result == []
+
+    def test_get_data_parts_empty_list(self):
+        # Setup
+        parts = []
+
+        # Exercise
+        result = get_data_parts(parts)
+
+        # Verify
+        assert result == []
+
+
+class TestGetFileParts:
+    def test_get_file_parts_single_file_part(self):
+        # Setup
+        file_with_uri = FileWithUri(
+            uri='file://path/to/file', mimeType='text/plain'
+        )
+        parts = [Part(root=FilePart(file=file_with_uri))]
+
+        # Exercise
+        result = get_file_parts(parts)
+
+        # Verify
+        assert result == [file_with_uri]
+
+    def test_get_file_parts_multiple_file_parts(self):
+        # Setup
+        file_with_uri1 = FileWithUri(
+            uri='file://path/to/file1', mimeType='text/plain'
+        )
+        file_with_bytes = FileWithBytes(
+            bytes='ZmlsZSBjb250ZW50',
+            mimeType='application/octet-stream',  # 'file content'
+        )
+        parts = [
+            Part(root=FilePart(file=file_with_uri1)),
+            Part(root=FilePart(file=file_with_bytes)),
+        ]
+
+        # Exercise
+        result = get_file_parts(parts)
+
+        # Verify
+        assert result == [file_with_uri1, file_with_bytes]
+
+    def test_get_file_parts_mixed_parts(self):
+        # Setup
+        file_with_uri = FileWithUri(
+            uri='file://path/to/file', mimeType='text/plain'
+        )
+        parts = [
+            Part(root=TextPart(text='some text')),
+            Part(root=FilePart(file=file_with_uri)),
+        ]
+
+        # Exercise
+        result = get_file_parts(parts)
+
+        # Verify
+        assert result == [file_with_uri]
+
+    def test_get_file_parts_no_file_parts(self):
+        # Setup
+        parts = [
+            Part(root=TextPart(text='some text')),
+            Part(root=DataPart(data={'key': 'value'})),
+        ]
+
+        # Exercise
+        result = get_file_parts(parts)
+
+        # Verify
+        assert result == []
+
+    def test_get_file_parts_empty_list(self):
+        # Setup
+        parts = []
+
+        # Exercise
+        result = get_file_parts(parts)
+
+        # Verify
+        assert result == []
+
+
 class TestGetMessageText:
     def test_get_message_text_single_part(self):
         # Setup