diff --git a/src/avb/bin.py b/src/avb/bin.py index 3050c91..f4876cb 100644 --- a/src/avb/bin.py +++ b/src/avb/bin.py @@ -232,6 +232,7 @@ def __init__(self): self.sifted_settings.append(s) self.attributes = self.root.create.Attributes() + self.mob_dict ={} def read(self, f): super(Bin, self).read(f) @@ -255,6 +256,7 @@ def read(self, f): object_count = ctx.read_u32(f) self.items = [] + self.mob_dict ={} for i in range(object_count): bin_obj = BinItem.__new__(BinItem, root=self.root) diff --git a/src/avb/components.py b/src/avb/components.py index 6f614cc..7ede820 100644 --- a/src/avb/components.py +++ b/src/avb/components.py @@ -313,6 +313,9 @@ def write(self, f): @property def mob(self): + if not self.root.content.mob_dict: + self.root.content.build_mob_dict() + if hasattr(self, 'mob_id'): mob_id = self.mob_id if mob_id: diff --git a/src/avb/trackgroups.py b/src/avb/trackgroups.py index f029f7d..c951afb 100644 --- a/src/avb/trackgroups.py +++ b/src/avb/trackgroups.py @@ -7,11 +7,11 @@ import datetime from . import core -from .core import AVBPropertyDef, AVBRefList -from .components import Component +from .core import AVBPropertyDef, walk_references from . import utils from . import mobid from . utils import peek_data +from .components import Component, SourceClip TRACK_LABEL_FLAG = 1 << 0 TRACK_ATTRIBUTES_FLAG = 1 << 1 @@ -225,6 +225,37 @@ def write(self, f): else: ctx.write_s16(f, 0) + def dependant_mobs(self): + """ + Yields all mobs that this mob is dependant on in depth first order. + """ + + visited = set() + stack = [self] + + if not self.root.content.mob_dict: + self.root.content.build_mob_dict() + + while stack: + mob = stack[-1] + children_processed = True + + for obj in walk_references(mob): + if isinstance(obj, SourceClip): + ref_mob = obj.mob + if not ref_mob: + continue + if ref_mob.mob_id not in visited: + stack.append(ref_mob) + children_processed = False + + if children_processed: + stack.pop(-1) + if mob.mob_id not in visited: + visited.add(mob.mob_id) + if mob is not self: + yield mob + @utils.register_class class TrackEffect(TrackGroup): diff --git a/tests/test_copy.py b/tests/test_copy.py index 8df63bc..455f3cf 100644 --- a/tests/test_copy.py +++ b/tests/test_copy.py @@ -37,5 +37,64 @@ def test_copy(self): assert mob_b compare(mob_a, mob_b) + def test_copy_mastermob_depends(self): + result_file = os.path.join(result_dir, 'copy_mastermobs.avb') + + mob_ids = set() + with avb.open(test_file_01) as a: + with avb.open() as b: + mobs = {} + for mob in a.content.mastermobs(): + mobs[mob.mob_id] = mob + for m in mob.dependant_mobs(): + mobs[m.mob_id] = m + + for mob in mobs.values(): + mob_ids.add(mob.mob_id) + new_mob = mob.copy(b) + b.content.add_mob(new_mob) + + + b.write(result_file) + + with avb.open(test_file_01) as a: + with avb.open(result_file) as b: + b.content.build_mob_dict() + a.content.build_mob_dict() + for mob_id in mob_ids: + mob_a = a.content.mob_dict.get(mob_id, None) + mob_b = b.content.mob_dict.get(mob_id, None) + compare(mob_a, mob_b) + + def test_copy_compositionmobs(self): + result_file = os.path.join(result_dir, 'copy_compositionmob.avb') + + mob_ids = set() + with avb.open(test_file_01) as a: + with avb.open() as b: + + mobs = {} + for mob in a.content.compositionmobs(): + mobs[mob.mob_id] = mob + for m in mob.dependant_mobs(): + mobs[m.mob_id] = m + + for mob in mobs.values(): + mob_ids.add(mob.mob_id) + new_mob = mob.copy(b) + b.content.add_mob(new_mob) + + + b.write(result_file) + + with avb.open(test_file_01) as a: + with avb.open(result_file) as b: + b.content.build_mob_dict() + a.content.build_mob_dict() + for mob_id in mob_ids: + mob_a = a.content.mob_dict.get(mob_id, None) + mob_b = b.content.mob_dict.get(mob_id, None) + compare(mob_a, mob_b) + if __name__ == "__main__": unittest.main()