MP4 tools

src/pymp4/cli.py

@@ -3,10 +3,9 @@
 import io
 import logging
 import argparse
-import codecs
 
 from pymp4.parser import Box
-from pymp4.util import BoxUtil
+from pymp4.util import BoxUtil, escape_decode
 from construct import setGlobalPrintFullStrings
 
 log = logging.getLogger(__name__)
@@ -23,20 +22,14 @@ def dump(args=None):
 
     setGlobalPrintFullStrings(args.full)
 
-    boxes_to_dump = {codecs.escape_decode(b)[0] for b in args.box or [] if b}
+    boxes_to_dump = {escape_decode(b) for b in args.box or [] if b}
 
     fd = args.input_file
     fd.seek(0, io.SEEK_END)
     eof = fd.tell()
     fd.seek(0)
 
-    def dump_box(box):
-        for t in boxes_to_dump:
-            for b in BoxUtil.find(box, t):
-                print(b)
-
-    print_fn = dump_box if boxes_to_dump else print
-
     while fd.tell() < eof:
         box = Box.parse_stream(fd)
-        print_fn(box)
+        for b in BoxUtil.search(box, *boxes_to_dump, decode=False):
+            print(b)

src/pymp4/util.py

@@ -15,13 +15,18 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
+import codecs
 import logging
 
 from pymp4.exceptions import BoxNotFound
 
 log = logging.getLogger(__name__)
 
 
+def escape_decode(s):
+    return codecs.escape_decode(s)[0]
+
+
 class BoxUtil(object):
     @classmethod
     def child(cls, box, type_):
@@ -59,6 +64,21 @@ def find(cls, box, type_):
                 for fbox in cls.find(sbox, type_):
                     yield fbox
 
+    @classmethod
+    def search(cls, box, *types, decode=True, key='type', children='children'):
+        if decode:
+            types = {escape_decode(t) for t in types if t}
+        if not types:
+            yield box
+        else:
+            if getattr(box, key, None) in types:
+                yield box
+            for sbox in (getattr(box, children, None) or []):
+                for fbox in cls.search(sbox, *types, decode=False,
+                        key=key, children=children):
+                    yield fbox
+
+
     @classmethod
     def find_extended(cls, box, extended_type_):
         if hasattr(box, "extended_type"):

tests/test_box.py

@@ -15,15 +15,23 @@
    limitations under the License.
 """
 import logging
+import pytest
 import unittest
 
-from construct import Container, ListContainer
+from construct import Container, ListContainer, StreamError
 from pymp4.parser import Box
 
 log = logging.getLogger(__name__)
 
 
 class BoxTests(unittest.TestCase):
+    def test_build_non_fourcc_type(self):
+        with pytest.raises(StreamError) as err_info:
+            Box.build(Container(
+                type=b"MP4",
+            ))
+        assert 'bytes object of wrong length, expected 4, found 3' in err_info.value.args[0]
+
     def test_ftyp_parse(self):
         self.assertEqual(
             Box.parse(b'\x00\x00\x00\x18ftypiso5\x00\x00\x00\x01iso5avc1'),

tests/test_cli.py

@@ -14,8 +14,7 @@ def test_missing_arg(capsys):
     captured = capsys.readouterr()
     assert "the following arguments are required: FILE" in captured.err
 
-def test_dump(capsys):
-    cli.dump(args=[TEST_FILE])
+def __assert_full_dump(capsys):
     captured = capsys.readouterr()
     assert "" == captured.err
     all_types = re.findall(r'type = [^\s]+', captured.out)
@@ -24,15 +23,13 @@ def test_dump(capsys):
     assert captured.out.rstrip().endswith('end = 60')
     assert 'truncated,' not in captured.out
 
+def test_dump(capsys):
+    cli.dump(args=[TEST_FILE])
+    __assert_full_dump(capsys)
+
 def test_full_dump(capsys):
     cli.dump(args=['--full', TEST_FILE])
-    captured = capsys.readouterr()
-    assert "" == captured.err
-    all_types = re.findall(r'type = [^\s]+', captured.out)
-    assert all_types == ["type = b'ftyp'", "type = b'free'"]
-    assert "major_brand = b'M4A ' " in captured.out
-    assert captured.out.rstrip().endswith('end = 60')
-    assert 'truncated,' not in captured.out
+    __assert_full_dump(capsys)
 
 def test_truncated_dump(capsys):
     cli.dump(args=['--truncated', TEST_FILE])
@@ -42,8 +39,7 @@ def test_truncated_dump(capsys):
     assert captured.out.rstrip().endswith('end = 60')
     assert '\\x16\'... (truncated, total 24)' in captured.out
 
-def test_dumping_specific_box(capsys):
-    cli.dump(args=['-b', 'free', TEST_FILE])
+def __assert_free_box_only(capsys):
     captured = capsys.readouterr()
     assert "" == captured.err
 
@@ -55,15 +51,14 @@ def test_dumping_specific_box(capsys):
     assert captured.out.rstrip().endswith('end = 60')
     assert "!\"#$' (total 24)" in captured.out
 
+def test_dumping_specific_box(capsys):
+    cli.dump(args=['-b', 'free', TEST_FILE])
+    __assert_free_box_only(capsys)
+
 def test_dumping_specific_box_with_escaped_input(capsys):
     cli.dump(args=['-b', '\\x66ree', TEST_FILE])
-    captured = capsys.readouterr()
-    assert "" == captured.err
+    __assert_free_box_only(capsys)
 
-    all_types = re.findall(r'type = [^\s]+', captured.out)
-    assert all_types == ["type = b'free'"]
-
-    assert "major_brand = b'M4A ' " not in captured.out
-    assert "type = b'free'" in captured.out
-    assert captured.out.rstrip().endswith('end = 60')
-    assert "!\"#$' (total 24)" in captured.out
+def test_dumping_specific_box_with_empty_string(capsys):
+    cli.dump(args=['-b', '', TEST_FILE])
+    __assert_full_dump(capsys)

tests/test_util.py

@@ -76,6 +76,114 @@ class BoxTests(unittest.TestCase):
         ])
     )
 
+    box_search_data = Container(
+        type=b"demo",
+        children=ListContainer([
+            Container(type=b"a   ", id=1),
+            Container(type=b"b   ", id=2),
+            Container(
+                type=b"c   ",
+                children=ListContainer([
+                    Container(type="a   ", id=3),
+                    Container(type=b"b   ", id=4),
+                    Container(type=b"\xa9too", id=99),
+                    Container(type=b"c   ", id=98),
+                ])
+            ),
+            Container(type=b"d   ", id=5),
+            Container(type="zzzz", id=100,
+                entries=ListContainer([
+                    Container(type=b"zzzz", id=101),
+                    Container(type="xyzz", id=102)
+                ])
+            ),
+        ])
+    )
+
+    def test_search_convert_to_bytes_false(self):
+        actual = list(BoxUtil.search(self.box_data, "a   ", decode=False))
+        expect = [Container(type="a   ", id=1), Container(type="a   ", id=3)]
+        assert actual == expect
+
+    def test_search_not_found(self):
+        # convert_to_bytes true so "a   " converts to b"a   "
+        actual = list(BoxUtil.search(self.box_data, "a   "))
+        expect = []
+        assert actual == expect
+
+    def test_search(self):
+        # convert_to_bytes true so "a   " converts to b"a   "
+        actual = list(BoxUtil.search(self.box_search_data, "a   "))
+        expect = [Container(type=b"a   ", id=1)]
+        assert actual == expect
+
+    def test_search_converts_bytes2bytes(self):
+        # convert_to_bytes true so b"\xa9too" converts to b"\xa9to"
+        actual = list(BoxUtil.search(self.box_search_data, b"\xa9too"))
+        expect = [Container(type=b"\xa9too", id=99)]
+        assert actual == expect
+
+    def test_search_converts_str2bytes(self):
+        # convert_to_bytes true so "\\xa9too" converts to b"\xa9too"
+        actual = list(BoxUtil.search(self.box_search_data, "\\xa9too"))
+        expect = [Container(type=b"\xa9too", id=99)]
+        assert actual == expect
+
+    def test_search_converts_utf8_str2bytes(self):
+        # convert_to_bytes true so "\xa9too" converts to b"\xc2\xa9too"
+        actual = list(BoxUtil.search(self.box_search_data, "\xa9too"))
+        expect = []
+        assert actual == expect
+
+    def test_search_includes_childrens(self):
+        # convert_to_bytes true so "b   " converts to b"b   "
+        actual = list(BoxUtil.search(self.box_search_data, "b   ", decode=True))
+        expect = [Container(type=b"b   ", id=2), Container(type=b"b   ", id=4)]
+        assert actual == expect
+
+    def test_search_includes_childrens_of_matched(self):
+        # convert_to_bytes true so "c   " converts to b"c   "
+        actual = list(BoxUtil.search(self.box_search_data, "c   ", decode=True))
+        expect = [self.box_search_data.children[2], Container(type=b"c   ", id=98)]
+        assert actual == expect
+
+    def test_search_key_is_id(self):
+        actual = list(BoxUtil.search(self.box_search_data, 2, 3, key='id', decode=False))
+        expect = [Container(type=b"b   ", id=2), Container(type="a   ", id=3)]
+        assert actual == expect
+
+    def test_search_children_key_is_entries(self):
+        actual = list(BoxUtil.search(self.box_search_data.children[-1], 'zzzz', children='entries'))
+        expect = [Container(type=b"zzzz", id=101)]
+        assert actual == expect
+
+    def test_search_key_does_not_exist(self):
+        # convert_to_bytes true so "a   " converts to b"a   "
+        actual = list(BoxUtil.search(self.box_search_data, "a   ", key='NOTE'))
+        expect = []
+        assert actual == expect
+
+    def test_search_childen_key_does_not_exist_in_root_box(self):
+        # convert_to_bytes true so "a   " converts to b"a   "
+        actual = list(BoxUtil.search(self.box_search_data, "a   ", children='entries'))
+        expect = []
+        assert actual == expect
+
+    def test_search_empty_key_returns_self(self):
+        actual = list(BoxUtil.search(self.box_search_data.children[0]))
+        expect = [self.box_search_data.children[0]]
+        assert actual == expect
+
+    def test_search_none_key_returns_self(self):
+        actual = list(BoxUtil.search(self.box_search_data.children[0], None))
+        expect = [self.box_search_data.children[0]]
+        assert actual == expect
+
+    def test_search_empty_str_key_returns_self(self):
+        actual = list(BoxUtil.search(self.box_search_data.children[0], ""))
+        expect = [self.box_search_data.children[0]]
+        assert actual == expect
+
     def test_child(self):
         self.assertListEqual(
             list(BoxUtil.child(self.box_data, "a   ")),