From 6f4a0ba90a87e8f2cda3be32ad3dcba647b22827 Mon Sep 17 00:00:00 2001
From: j <j@mailb.org>
Date: Fri, 1 Jul 2016 18:10:01 +0200
Subject: [PATCH] support vtt subtitles

---
 pandora/annotation/management/commands/import_srt.py |  7 +++++--
 pandora/archive/models.py                            |  8 ++++++--
 pandora/archive/views.py                             | 11 +++++++----
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/pandora/annotation/management/commands/import_srt.py b/pandora/annotation/management/commands/import_srt.py
index 03daf68cb..f62990dc6 100644
--- a/pandora/annotation/management/commands/import_srt.py
+++ b/pandora/annotation/management/commands/import_srt.py
@@ -19,7 +19,7 @@ class Command(BaseCommand):
     """
     import annotations
     """
-    help = 'import annotations from srt'
+    help = 'import annotations from srt or vtt'
     args = 'username item layername filename.srt'
     option_list = BaseCommand.option_list + (
     )
@@ -33,7 +33,10 @@ class Command(BaseCommand):
         item = Item.objects.get(public_id=public_id)
         layer = filter(lambda l: l['id'] == layer_id, settings.CONFIG['layers'])[0]
 
-        annotations = ox.srt.load(filename)
+        if filename.endswith('.vtt'):
+            annotations = ox.vtt.load(filename)
+        else:
+            annotations = ox.srt.load(filename)
         print('importing %d annotations into %s/%s' % (len(annotations), public_id, layer_id))
         for i in range(len(annotations)-1):
             if annotations[i]['out'] == annotations[i+1]['in']:
diff --git a/pandora/archive/models.py b/pandora/archive/models.py
index 2bfd14a5f..d8187c1e9 100644
--- a/pandora/archive/models.py
+++ b/pandora/archive/models.py
@@ -256,7 +256,7 @@ class File(models.Model):
             self.sort_path = utils.sort_string(self.path)
             self.is_audio = self.type == 'audio'
             self.is_video = self.type == 'video'
-            self.is_subtitle = self.path.endswith('.srt')
+            self.is_subtitle = self.path.split('.')[-1] in ('srt', 'vtt')
 
         if self.type not in ('audio', 'video'):
             self.duration = None
@@ -289,7 +289,11 @@ class File(models.Model):
     def srt(self, offset=0):
         srt = []
         subtitles = []
-        for s in ox.srt.load(self.data.path):
+        if self.data.path.endswith('.vtt'):
+            load = ox.vtt.load
+        else:
+            load = ox.srt.load
+        for s in load(self.data.path):
             if s['in'] <= s['out'] and s['value'].strip():
                 key = '%s --> %s\n%s' % (s['in'], s['out'], s['value'])
                 if key not in subtitles:
diff --git a/pandora/archive/views.py b/pandora/archive/views.py
index 0e349f217..bde0ce880 100644
--- a/pandora/archive/views.py
+++ b/pandora/archive/views.py
@@ -103,10 +103,13 @@ def update(request, data):
                                                                          file__wanted=True)]
 
         if filter(lambda l: l['id'] == 'subtitles', settings.CONFIG['layers']):
-            response['data']['file'] = [f.file.oshash
-                                        for f in files.filter(file__is_subtitle=True,
-                                                              file__available=False,
-                                                              path__endswith='.srt')]
+            qs = files.filter(
+                file__is_subtitle=True,
+                file__available=False
+            ).filter(
+                Q(path__endswith='.srt') | Q(path__endswith='.vtt')
+            )
+            response['data']['file'] = [f.file.oshash for f in qs]
         else:
             response['data']['file'] = []
     return render_to_json_response(response)