diff --git a/pycaption/base.py b/pycaption/base.py index 30994d99..3db46c41 100644 --- a/pycaption/base.py +++ b/pycaption/base.py @@ -242,11 +242,37 @@ def _format_timestamp(self, value, msec_separator=None): class CaptionList(list): """ A list of captions with a layout object attached to it """ - def __init__(self, layout_info=None): + def __init__(self, iterable=None, layout_info=None): """ + :param iterator: An iterator used to populate the caption list :param Layout layout_info: A Layout object with the positioning info """ self.layout_info = layout_info + args = [iterable] if iterable else [] + super(CaptionList, self).__init__(*args) + + def __getslice__(self, i, j): + return CaptionList( + list.__getslice__(self, i, j), layout_info=self.layout_info) + + def __add__(self, other): + add_is_safe = ( + not hasattr(other, 'layout_info') or + not other.layout_info or + self.layout_info == other.layout_info + ) + if add_is_safe: + return CaptionList( + list.__add__(self, other), layout_info=self.layout_info) + else: + raise ValueError( + "Cannot add CaptionList objects with different layout_info") + + def __mul__(self, other): + return CaptionList( + list.__mul__(self, other), layout_info=self.layout_info) + + __rmul__ = __mul__ class CaptionSet(object): diff --git a/pycaption/dfxp/base.py b/pycaption/dfxp/base.py index 88ed3f60..f981cea8 100644 --- a/pycaption/dfxp/base.py +++ b/pycaption/dfxp/base.py @@ -98,10 +98,10 @@ def _get_dfxp_parser_class(): return LayoutAwareDFXPParser def _translate_div(self, div): - captions = CaptionList(div.layout_info) - for p_tag in div.find_all(u'p'): - captions.append(self._translate_p_tag(p_tag)) - return captions + return CaptionList( + [self._translate_p_tag(p_tag) for p_tag in div.find_all(u'p')], + div.layout_info + ) def _translate_p_tag(self, p_tag): start, end = self._find_times(p_tag) @@ -109,9 +109,8 @@ def _translate_p_tag(self, p_tag): self._translate_tag(p_tag) styles = self._translate_style(p_tag) - caption = Caption(start, end, self.nodes, style=styles, - layout_info=p_tag.layout_info) - return caption + return Caption( + start, end, self.nodes, style=styles, layout_info=p_tag.layout_info) def _find_times(self, p_tag): start = self._translate_time(p_tag[u'begin']) diff --git a/pycaption/sami.py b/pycaption/sami.py index 54af8a40..99e53400 100644 --- a/pycaption/sami.py +++ b/pycaption/sami.py @@ -182,7 +182,7 @@ def _translate_lang(self, language, sami_soup, parent_layout): :rtype: list """ - captions = CaptionList(parent_layout) + captions = CaptionList(layout_info=parent_layout) milliseconds = 0 for p in sami_soup.select(u'p[lang|=%s]' % language): diff --git a/tests/test_base.py b/tests/test_base.py new file mode 100644 index 00000000..e1dea95f --- /dev/null +++ b/tests/test_base.py @@ -0,0 +1,42 @@ +import unittest + +from pycaption.base import CaptionList + + +class CaptionListTestCase(unittest.TestCase): + + def setUp(self): + self.layout_info = "My Layout" + self.caps = CaptionList([1, 2, 3], layout_info=self.layout_info) + + def test_splice(self): + newcaps = self.caps[1:] + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + def test_mul(self): + newcaps = self.caps * 2 + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + def test_rmul(self): + newcaps = 2 * self.caps + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + def test_add_list_to_caption_list(self): + newcaps = self.caps + [9, 8, 7] + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + def test_add_two_caption_lists(self): + newcaps = self.caps + CaptionList([4], layout_info=None) + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + newcaps = self.caps + CaptionList([4], layout_info=self.layout_info) + self.assertTrue(isinstance(newcaps, CaptionList)) + self.assertTrue(newcaps.layout_info == self.layout_info) + + with self.assertRaises(ValueError): + newcaps = self.caps + CaptionList([4], layout_info="Other Layout")