Merge pull request #23 from omerholz/read_from_memory

let the parser read content either from a file or from a file-like object

Author: glut23

tests/webvtt.py

@@ -5,6 +5,7 @@
 import webvtt
 from webvtt.structures import Caption, Style
 from .generic import GenericParserTestCase
+from webvtt.errors import MalformedFileError
 
 
 BASE_DIR = os.path.dirname(__file__)
@@ -222,6 +223,28 @@ def test_manipulate_lines(self):
             'Caption line #1 updated'
         )
 
+    def test_read_file_buffer(self):
+        with open(self._get_file('sample.vtt'), 'r', encoding='utf-8') as f:
+            vtt = webvtt.read_buffer(f)
+            self.assertIsInstance(vtt.captions, list)
+
+    def test_read_memory_buffer(self):
+        payload = ''
+        with open(self._get_file('sample.vtt'), 'r', encoding='utf-8') as f:
+            payload = f.read()
+
+        buffer = io.StringIO(payload)
+        vtt = webvtt.read_buffer(buffer)
+        self.assertIsInstance(vtt.captions, list)
+
+    def test_read_malformed_buffer(self):
+        malformed_payloads = ['', 'MOCK MELFORMED CONTENT']
+        for payload in malformed_payloads:
+            buffer = io.StringIO(payload)
+            with self.assertRaises(MalformedFileError):
+                webvtt.read_buffer(buffer)
+
+
     def test_captions(self):
         vtt = webvtt.read(self._get_file('sample.vtt'))
         self.assertIsInstance(vtt.captions, list)

webvtt/__init__.py

@@ -8,6 +8,7 @@
 __all__ = webvtt.__all__ + segmenter.__all__ + structures.__all__ + errors.__all__
 
 read = WebVTT.read
+read_buffer = WebVTT.read_buffer
 from_srt = WebVTT.from_srt
 from_sbv = WebVTT.from_sbv
 list_formats = WebVTT.list_formats

webvtt/parsers.py

@@ -21,31 +21,46 @@ def __init__(self, parse_options=None):
 
     def read(self, file):
         """Reads the captions file."""
-        content = self._read_content(file)
+        content = self._get_content_from_file(file_path=file)
         self._validate(content)
         self._parse(content)
 
         return self
 
-    def _read_content(self, file):
+    def read_from_buffer(self, buffer):
+        content = self._read_content_lines(buffer)
+        self._validate(content)
+        self._parse(content)
+
+        return self
 
-        first_bytes = min(32, os.path.getsize(file))
-        with open(file, 'rb') as f:
+    def _get_content_from_file(self, file_path):
+        encoding = self._read_file_encoding(file_path)
+        with open(file_path, encoding=encoding) as f:
+            return self._read_content_lines(f)
+
+    def _read_file_encoding(self, file_path):
+        first_bytes = min(32, os.path.getsize(file_path))
+        with open(file_path, 'rb') as f:
             raw = f.read(first_bytes)
 
         if raw.startswith(codecs.BOM_UTF8):
-            encoding = 'utf-8-sig'
+            return 'utf-8-sig'
         else:
-            encoding = 'utf-8'
+            return 'utf-8'
+
+    def _read_content_lines(self, file_obj):
 
-        with open(file, encoding=encoding) as f:
-            lines = [line.rstrip('\n') for line in f.readlines()]
+        lines = [line.rstrip('\n') for line in file_obj.readlines()]
 
         if not lines:
             raise MalformedFileError('The file is empty.')
 
         return lines
 
+    def _read_content(self, file):
+        return self._get_content_from_file(file_path=file)
+
     def _parse_timeframe_line(self, line):
         """Parse timeframe line and return start and end timestamps."""
         tf = self._validate_timeframe_line(line)

webvtt/webvtt.py

@@ -60,6 +60,14 @@ def read(cls, file):
         parser = WebVTTParser().read(file)
         return cls(file=file, captions=parser.captions, styles=parser.styles)
 
+    @classmethod
+    def read_buffer(cls, buffer):
+        """Reads a WebVTT captions from a file-like object.
+        Such file-like object may be the return of an io.open call,
+        io.StringIO object, tempfile.TemporaryFile object, etc."""
+        parser = WebVTTParser().read_from_buffer(buffer)
+        return cls(captions=parser.captions, styles=parser.styles)
+
     def _get_output_file(self, output, extension='vtt'):
         if not output:
             if not self.file: