Manage zipfile for imports

1 files changed, 66 insertions, 22 deletions

diff --git a/chimere/utils.py b/chimere/utils.py
index 97f710f..2459075 100644
--- a/chimere/utils.py
+++ b/chimere/utils.py
@@ -23,6 +23,8 @@ Utilitaries
 
 import urllib2, re
 import unicodedata
+import zipfile
+import StringIO
 from external_utils import OsmApi
 from lxml import etree
 from chimere import get_version
@@ -36,6 +38,27 @@ def unicode_normalize(string):
         (c for c in unicodedata.normalize('NFD', string)
         if unicodedata.category(c) != 'Mn'))
 
+def get_files_inside_zip(zippedfile, suffixes):
+    try:
+        flz = zipfile.ZipFile(zippedfile)
+    except zipfile.BadZipfile:
+        return [], _(u"Bad zip file")
+    namelist = flz.namelist()
+    filenames = []
+    for suffix in suffixes:
+        current_file_name = None
+        for name in namelist:
+            if name.endswith(suffix):
+                current_file_name = name
+        filenames.append(current_file_name)
+    files = []
+    for filename in filenames:
+        if filename:
+            files.append(flz.open(filename))
+        else:
+            files.append(None)
+    return files
+
 class ImportManager:
     u"""
     Generic class for specific importers
@@ -56,35 +79,54 @@ class KMLManager(ImportManager):
     imported
     """
     XPATH = '//kml:Folder/kml:name[text()="%s"]/../kml:Placemark'
+    DEFAULT_XPATH = '//kml:Placemark'
     def __init__(self, importer_instance, ns=''):
         self.importer_instance = importer_instance
         self.ns = ns
 
-    def get(self):
+    def get(self, source=None):
         u"""
         Get data from the source
+        Args:
+         - source (None): input file if not provided get it from the distant
+         source provided in the importer instance.
+
         Return a tuple with:
-        - number of new item ;
-        - number of item updated ;
-        - error detail on error
+         - number of new item ;
+         - number of item updated ;
+         - error detail on error
         """
         from models import Marker
         new_item, updated_item, msg = 0, 0, ''
-        try:
-            source = urllib2.urlopen(self.importer_instance.source)
-        except ValueError:
-            # assume it is a local file
+        if not source:
             try:
-                source = open(self.importer_instance.source)
-            except IOError, msg:
-                return (new_item, updated_item, msg)
-        except urllib2.URLError as error:
-            return (new_item, updated_item, error.message)
+                remotehandle = urllib2.urlopen(self.importer_instance.source)
+                source = StringIO.StringIO(remotehandle.read())
+                remotehandle.close()
+            except ValueError:
+                # assume it is a local file
+                try:
+                    source = open(self.importer_instance.source)
+                except IOError, msg:
+                    return (new_item, updated_item, msg)
+            except urllib2.URLError as error:
+                return (new_item, updated_item, error.message)
+        if self.importer_instance.zipped:
+            try:
+                files = get_files_inside_zip(source, ['.kml'])
+            except zipfile.BadZipfile:
+                return (new_item, updated_item, _(u"Bad zip file"))
+            if not files or not files[0]:
+                return (new_item, updated_item,
+                        _(u"No KML file inside the zip file"))
+            source = files[0]
         tree = etree.parse(source)
         # try to get default namespace
         if not self.ns:
             self.ns = tree.getroot().nsmap[None]
-        for placemark in tree.xpath(self.XPATH % self.importer_instance.filtr,
+        xpath = self.XPATH % self.importer_instance.filtr \
+                  if self.importer_instance.filtr else self.DEFAULT_XPATH
+        for placemark in tree.xpath(xpath,
                                     namespaces={'kml':self.ns}):
             name, point, linestring = None, None, None
             pl_id = placemark.attrib.get('id')
@@ -141,13 +183,17 @@ class OSMManager(ImportManager):
     The filtr argument is XAPI args or empty if it is an OSM file.
     """
 
-    def get(self, get_items=None):
+    def get(self, source=None):
         u"""
         Get data from the source
+        Args:
+         - source (None): input file if not provided get it from the distant
+         source provided in the importer instance.
+
         Return a tuple with:
-        - new items or number of new items if nb_only = True;
-        - updated items or number of updated items if nb_only = True;
-        - error detail on error
+        - new items;
+        - updated items;
+        - error detail on error.
         """
         from models import Marker
         new_item, updated_item = 0 , 0
@@ -208,15 +254,13 @@ class OSMManager(ImportManager):
                 m.categories.clear()
                 for cat in self.importer_instance.categories.all():
                     m.categories.add(cat)
-        if not get_items:
-            return (new_item, updated_item, msg)
-        return (items, new_item, updated_item, msg)
+        return (new_item, updated_item, msg)
 
     def put(self):
         # first of all: reimport in order to verify that no changes has been
         # made since the last import
         from models import Marker
-        items, new_item, updated_item, msg = self.get(get_items=True)
+        new_item, updated_item, msg = self.get()
         # check if import is possible
         if msg:
             return 0, msg