From: Guillaume Pellerin <guillaume.pellerin@ircam.fr>
Date: Thu, 3 Dec 2015 08:44:01 +0000 (+0100)
Subject: fix CSV stream output
X-Git-Tag: 1.6b~7^2~46^2~1
X-Git-Url: https://git.parisson.com/?a=commitdiff_plain;h=03d5d057a1eaee21b4bf79232fe7dbb0122e979d;p=telemeta.git

fix CSV stream output
---

diff --git a/telemeta/models/collection.py b/telemeta/models/collection.py
index 2b76bda4..58d9ecdf 100644
--- a/telemeta/models/collection.py
+++ b/telemeta/models/collection.py
@@ -273,6 +273,16 @@ class MediaCollection(MediaResource):
         return json.dumps(self.to_dict_with_more())
 
 
+    def to_row(self):
+        row = []
+        _dict = self.to_dict_with_more()
+        for tag in _dict:
+            if tag in _dict.keys():
+                row.append(_dict[tag])
+            else:
+                row.append('')
+        return row
+
 
 class MediaCollectionRelated(MediaRelated):
     "Collection related media"
@@ -295,4 +305,3 @@ class MediaCollectionIdentifier(Identifier):
         verbose_name = _('collection identifier')
         verbose_name_plural = _('collection identifiers')
         unique_together = ('identifier', 'collection')
-
diff --git a/telemeta/models/core.py b/telemeta/models/core.py
index 5e647bee..98853f16 100644
--- a/telemeta/models/core.py
+++ b/telemeta/models/core.py
@@ -264,5 +264,3 @@ class CoreManager(EnhancedManager):
                 return super(CoreManager, self).get(**args)
 
         return super(CoreManager, self).get(**kwargs)
-
-
diff --git a/telemeta/models/item.py b/telemeta/models/item.py
index fc7c78cf..53c5669a 100644
--- a/telemeta/models/item.py
+++ b/telemeta/models/item.py
@@ -339,6 +339,15 @@ class MediaItem(MediaResource):
 
         return metadata
 
+    def to_row(self):
+        row = []
+        _dict = self.to_dict_with_more()
+        for tag in _dict:
+            if tag in _dict.keys():
+                row.append(_dict[tag])
+            else:
+                row.append('')
+        return row
 
 class MediaItemRelated(MediaRelated):
     "Item related media"
diff --git a/telemeta/util/unicode.py b/telemeta/util/unicode.py
index aa81e8cb..0e9cc8ba 100644
--- a/telemeta/util/unicode.py
+++ b/telemeta/util/unicode.py
@@ -50,19 +50,24 @@ def _stringify_list(l, encoding):
     return [_stringify(s, encoding) for s in l]
 
 
+class StreamCSVException(Exception):
+    pass
+
+
 class UnicodeCSVWriter(object):
-    def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
+    def __init__(self, f, elements, dialect=csv.excel, encoding="utf-8", **kwds):
         self.dialect = dialect
         self.encoding = encoding
         self.writer = csv.writer(f, dialect=dialect, **kwds)
         self.line = 0
+        self.elements = elements
         self.tags = []
-
+        self.get_tags(self.elements[0])
 
     def get_tags(self, element):
-        element_dict = element.to_dict_with_more()
+        _dict = element.to_dict_with_more()
 
-        for key in element_dict.keys():
+        for key in _dict.keys():
             if not key in self.tags:
                 self.tags.append(key)
 
@@ -72,23 +77,11 @@ class UnicodeCSVWriter(object):
         self.tags.sort()
         self.tags.insert(0, 'title')
         self.tags.insert(0, 'code')
-        self.init_tags = True
-
-    def get_row(self, element):
-        row = []
-        element_dict = element.to_dict_with_more()
-        for tag in self.tags:
-            if tag in element_dict.keys():
-                    row.append(element_dict[tag])
-            else:
-                row.append('')
-        return row
-
-    def write_element(self, element):
-        if not self.tags:
-            yield self.writer.writerow(self.get_tags())
-        row = self.get_row(element)
-        yield self.writer.writerow(_stringify_list(row, self.encoding))
+
+    def output(self):
+        yield self.writer.writerow(self.tags)
+        for element in self.elements:
+            yield self.writer.writerow(_stringify_list(element.to_row(), self.encoding))
 
 
 class Echo(object):
diff --git a/telemeta/views/playlist.py b/telemeta/views/playlist.py
index a47a66fd..11a55d2a 100644
--- a/telemeta/views/playlist.py
+++ b/telemeta/views/playlist.py
@@ -112,7 +112,7 @@ class PlaylistView(object):
         playlist = Playlist.objects.get(public_id=public_id)
         elements = self.get_elements(playlist, resource_type)
         pseudo_buffer = Echo()
-        writer = UnicodeCSVWriter(pseudo_buffer)
-        response = StreamingHttpResponse((writer.write_element(element) for element in elements), content_type='text/csv')
+        writer = UnicodeCSVWriter(pseudo_buffer, elements)
+        response = StreamingHttpResponse(writer.output(), content_type='text/csv')
         response['Content-Disposition'] = 'attachment; filename='+playlist.title+'_'+resource_type+'.csv'
         return response