| 51 | |
| 52 | |
| 53 | class SubtitleMerger(SubsMixin, TransformerMixin): |
| 54 | def __init__(self, reference_subs, first="reference"): |
| 55 | assert first in ("reference", "output") |
| 56 | super(SubsMixin, self).__init__() |
| 57 | self.reference_subs = reference_subs |
| 58 | self.first = first |
| 59 | |
| 60 | def fit(self, output_subs: GenericSubtitlesFile, *_): |
| 61 | def _merger_gen(a, b): |
| 62 | ita, itb = iter(a), iter(b) |
| 63 | cur_a = next(ita, None) |
| 64 | cur_b = next(itb, None) |
| 65 | while True: |
| 66 | if cur_a is None and cur_b is None: |
| 67 | return |
| 68 | elif cur_a is None: |
| 69 | while cur_b is not None: |
| 70 | yield cur_b |
| 71 | cur_b = next(itb, None) |
| 72 | return |
| 73 | elif cur_b is None: |
| 74 | while cur_a is not None: |
| 75 | yield cur_a |
| 76 | cur_a = next(ita, None) |
| 77 | return |
| 78 | # else: neither are None |
| 79 | if cur_a.start < cur_b.start: |
| 80 | swapped = False |
| 81 | else: |
| 82 | swapped = True |
| 83 | cur_a, cur_b = cur_b, cur_a |
| 84 | ita, itb = itb, ita |
| 85 | prev_a = cur_a |
| 86 | while prev_a is not None and cur_a.start < cur_b.start: |
| 87 | cur_a = next(ita, None) |
| 88 | if cur_a is None or cur_a.start < cur_b.start: |
| 89 | yield prev_a |
| 90 | prev_a = cur_a |
| 91 | if prev_a is None: |
| 92 | while cur_b is not None: |
| 93 | yield cur_b |
| 94 | cur_b = next(itb, None) |
| 95 | return |
| 96 | if cur_b.start - prev_a.start < cur_a.start - cur_b.start: |
| 97 | if swapped: |
| 98 | yield cur_b.merge_with(prev_a) |
| 99 | ita, itb = itb, ita |
| 100 | cur_a, cur_b = cur_b, cur_a |
| 101 | cur_a = next(ita, None) |
| 102 | else: |
| 103 | yield prev_a.merge_with(cur_b) |
| 104 | cur_b = next(itb, None) |
| 105 | else: |
| 106 | if swapped: |
| 107 | yield cur_b.merge_with(cur_a) |
| 108 | ita, itb = itb, ita |
| 109 | else: |
| 110 | yield cur_a.merge_with(cur_b) |