1# -*- coding: utf-8 -*-
2# Copyright 2006 Joe Wreschnig
3#
4# This program is free software; you can redistribute it and/or modify
5# it under the terms of the GNU General Public License as published by
6# the Free Software Foundation; either version 2 of the License, or
7# (at your option) any later version.
8
9"""Read and write Ogg Theora comments.
10
11This module handles Theora files wrapped in an Ogg bitstream. The
12first Theora stream found is used.
13
14Based on the specification at http://theora.org/doc/Theora_I_spec.pdf.
15"""
16
17__all__ = ["OggTheora", "Open", "delete"]
18
19import struct
20
21from mutagen import StreamInfo
22from mutagen._vorbis import VCommentDict
23from mutagen._util import cdata, get_size, loadfile, convert_error
24from mutagen._tags import PaddingInfo
25from mutagen.ogg import OggPage, OggFileType, error as OggError
26
27
28class error(OggError):
29    pass
30
31
32class OggTheoraHeaderError(error):
33    pass
34
35
36class OggTheoraInfo(StreamInfo):
37    """OggTheoraInfo()
38
39    Ogg Theora stream information.
40
41    Attributes:
42        length (`float`): File length in seconds, as a float
43        fps (`float`): Video frames per second, as a float
44        bitrate (`int`): Bitrate in bps (int)
45    """
46
47    length = 0
48    fps = 0
49    bitrate = 0
50
51    def __init__(self, fileobj):
52        page = OggPage(fileobj)
53        while not page.packets[0].startswith(b"\x80theora"):
54            page = OggPage(fileobj)
55        if not page.first:
56            raise OggTheoraHeaderError(
57                "page has ID header, but doesn't start a stream")
58        data = page.packets[0]
59        vmaj, vmin = struct.unpack("2B", data[7:9])
60        if (vmaj, vmin) != (3, 2):
61            raise OggTheoraHeaderError(
62                "found Theora version %d.%d != 3.2" % (vmaj, vmin))
63        fps_num, fps_den = struct.unpack(">2I", data[22:30])
64        self.fps = fps_num / float(fps_den)
65        self.bitrate = cdata.uint_be(b"\x00" + data[37:40])
66        self.granule_shift = (cdata.ushort_be(data[40:42]) >> 5) & 0x1F
67        self.serial = page.serial
68
69    def _post_tags(self, fileobj):
70        page = OggPage.find_last(fileobj, self.serial, finishing=True)
71        if page is None:
72            raise OggTheoraHeaderError
73        position = page.position
74        mask = (1 << self.granule_shift) - 1
75        frames = (position >> self.granule_shift) + (position & mask)
76        self.length = frames / float(self.fps)
77
78    def pprint(self):
79        return u"Ogg Theora, %.2f seconds, %d bps" % (self.length,
80                                                      self.bitrate)
81
82
83class OggTheoraCommentDict(VCommentDict):
84    """Theora comments embedded in an Ogg bitstream."""
85
86    def __init__(self, fileobj, info):
87        pages = []
88        complete = False
89        while not complete:
90            page = OggPage(fileobj)
91            if page.serial == info.serial:
92                pages.append(page)
93                complete = page.complete or (len(page.packets) > 1)
94        data = OggPage.to_packets(pages)[0][7:]
95        super(OggTheoraCommentDict, self).__init__(data, framing=False)
96        self._padding = len(data) - self._size
97
98    def _inject(self, fileobj, padding_func):
99        """Write tag data into the Theora comment packet/page."""
100
101        fileobj.seek(0)
102        page = OggPage(fileobj)
103        while not page.packets[0].startswith(b"\x81theora"):
104            page = OggPage(fileobj)
105
106        old_pages = [page]
107        while not (old_pages[-1].complete or len(old_pages[-1].packets) > 1):
108            page = OggPage(fileobj)
109            if page.serial == old_pages[0].serial:
110                old_pages.append(page)
111
112        packets = OggPage.to_packets(old_pages, strict=False)
113
114        content_size = get_size(fileobj) - len(packets[0])  # approx
115        vcomment_data = b"\x81theora" + self.write(framing=False)
116        padding_left = len(packets[0]) - len(vcomment_data)
117
118        info = PaddingInfo(padding_left, content_size)
119        new_padding = info._get_padding(padding_func)
120
121        packets[0] = vcomment_data + b"\x00" * new_padding
122
123        new_pages = OggPage._from_packets_try_preserve(packets, old_pages)
124        OggPage.replace(fileobj, old_pages, new_pages)
125
126
127class OggTheora(OggFileType):
128    """OggTheora(filething)
129
130    An Ogg Theora file.
131
132    Arguments:
133        filething (filething)
134
135    Attributes:
136        info (`OggTheoraInfo`)
137        tags (`mutagen._vorbis.VCommentDict`)
138    """
139
140    _Info = OggTheoraInfo
141    _Tags = OggTheoraCommentDict
142    _Error = OggTheoraHeaderError
143    _mimes = ["video/x-theora"]
144
145    info = None
146    tags = None
147
148    @staticmethod
149    def score(filename, fileobj, header):
150        return (header.startswith(b"OggS") *
151                ((b"\x80theora" in header) + (b"\x81theora" in header)) * 2)
152
153
154Open = OggTheora
155
156
157@convert_error(IOError, error)
158@loadfile(method=False, writable=True)
159def delete(filething):
160    """ delete(filething)
161
162    Arguments:
163        filething (filething)
164    Raises:
165        mutagen.MutagenError
166
167    Remove tags from a file.
168    """
169
170    t = OggTheora(filething)
171    filething.fileobj.seek(0)
172    t.delete(filething)
173