From 05d5134c1049eec22d638b8b1bc91d6fc98266dd Mon Sep 17 00:00:00 2001
From: Finn Hughes <finn.hughes@arconics.com>
Date: Fri, 16 Nov 2018 00:13:51 +0000
Subject: [PATCH] Extend avc1 to support multiple video sample description
 extensions

avc1 can support multiple extensions including avcC and pasp (found in QT documentation)

avcc has optional fields at the end which are now ignored with GreedyBytes if present
---
 src/pymp4/parser.py | 38 +++++++++++++++++++++++++-------------
 tests/test_box.py   | 25 ++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/src/pymp4/parser.py b/src/pymp4/parser.py
index ddaf25a..9709e4e 100644
--- a/src/pymp4/parser.py
+++ b/src/pymp4/parser.py
@@ -332,6 +332,30 @@ def _encode(self, obj, context):
         return obj & 0x1F
 
 
+VideoSampleEntryExtensionBox = PrefixedIncludingSize(Int32ub, Struct(
+    "type" / String(4, padchar=b" ", paddir="right"),
+    Embedded(Switch(this.type, {
+        b"avcC": Struct(
+            "version" / Const(Int8ub, 1),
+            "profile" / Int8ub,
+            "compatibility" / Int8ub,
+            "level" / Int8ub,
+            EmbeddedBitStruct(
+                Padding(6, pattern=b'\x01'),
+                "nal_unit_length_field" / Default(BitsInteger(2), 3),
+            ),
+            "sps" / Default(PrefixedArray(MaskedInteger(Int8ub), PascalString(Int16ub)), []),
+            "pps" / Default(PrefixedArray(Int8ub, PascalString(Int16ub)), []),
+            # if profile_idc takes specific values there can be additional information - ignoring
+            GreedyBytes
+        ),
+        b"pasp": Struct(
+            "h_spacing" / Int32ub,
+            "v_spacing" /Int32ub
+        )
+    }, default=Struct(RawBox))),
+))
+
 AVC1SampleEntryBox = Struct(
     "version" / Default(Int16ub, 0),
     "revision" / Const(Int16ub, 0),
@@ -349,19 +373,7 @@ def _encode(self, obj, context):
     "compressor_name" / Default(String(32, padchar=b" "), ""),
     "depth" / Default(Int16ub, 24),
     "color_table_id" / Default(Int16sb, -1),
-    "avc_data" / PrefixedIncludingSize(Int32ub, Struct(
-        "type" / Const(b"avcC"),
-        "version" / Const(Int8ub, 1),
-        "profile" / Int8ub,
-        "compatibility" / Int8ub,
-        "level" / Int8ub,
-        EmbeddedBitStruct(
-            Padding(6, pattern=b'\x01'),
-            "nal_unit_length_field" / Default(BitsInteger(2), 3),
-        ),
-        "sps" / Default(PrefixedArray(MaskedInteger(Int8ub), PascalString(Int16ub)), []),
-        "pps" / Default(PrefixedArray(Int8ub, PascalString(Int16ub)), [])
-    ))
+    "extensions" / GreedyRange(VideoSampleEntryExtensionBox),
 )
 
 
diff --git a/tests/test_box.py b/tests/test_box.py
index c76ba5e..d59e9cc 100644
--- a/tests/test_box.py
+++ b/tests/test_box.py
@@ -16,9 +16,10 @@
 """
 import logging
 import unittest
+import io
 
 from construct import Container
-from pymp4.parser import Box
+from pymp4.parser import Box, SampleEntryBox
 
 log = logging.getLogger(__name__)
 
@@ -103,3 +104,25 @@ def test_moov_build(self):
             b'\x00\x00\x00\x20trex\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
             b'\x00\x00\x00\x20trex\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
         )
+
+    def test_avc1_parse(self):
+        compressor = b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+        sps = b'\x67\x4d\x40\x29\xe8\x80\x28\x02\xdd\xff\x80\x0d\x80\x0a\x08\x00\x00\x1f\x48\x00\x05\xdc\x00\x78\xc1\x88\x90'
+        pps = b'\x68\xeb\x8c\xb2'
+        input_bytes = (
+            b'\x00\x00\x00\x98avc1\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x05\x00\x02\xd0\x00\x48\x00\x00\x00\x48\x00\x00\x00\x00\x00\x00\x00\x01' + compressor + b'\x00\x18\xff\xff'
+            b'\x00\x00\x00\x32avcC\x01\x4d\x40\x29\xff\xe1\x00\x1b' + sps + b'\x01\x00\x04' + pps +
+            b'\x00\x00\x00\x10pasp\x00\x00\x00\x1b\x00\x00\x00\x14'
+        )
+        expected = (
+            Container(format=b'avc1')(data_reference_index=1)(version=0)(revision=0)(vendor=b'\x00\x00\x00\x00')
+            (temporal_quality=0)(spatial_quality=0)(width=1280)(height=720)(horizontal_resolution=72)
+            (vertical_resolution=72)(data_size=0)(frame_count=1)(compressor_name=compressor)(depth=24)(color_table_id=-1)
+            (extensions=[
+                Container(type=b'avcC')(version=1)(profile=77)(compatibility=64)(level=41)(nal_unit_length_field=3)(sps=[sps])(pps=[pps]),
+                Container(type=b'pasp')(h_spacing=27)(v_spacing=20)
+            ])
+        )
+        input_stream = io.BytesIO(input_bytes + b'padding')
+        self.assertEqual(SampleEntryBox.parse_stream(input_stream), expected)
+        self.assertEqual(input_stream.tell(), len(input_bytes))