Testiranje

tkrajina · tkrajina · commit 20457c8a20b2 · 2013-06-08T22:00:27.000+02:00
diff --git a/srtm/__init__.py b/srtm/__init__.py
@@ -39,7 +39,7 @@ def get_default_srtm_dir():
 
     return result
 
-def get_data(local_srtm_dir=None, reduce_big_files=False):
+def get_data(local_srtm_dir=None, reduce_big_files=False, leave_zipped=False):
     """
     Get the utility object for querying elevation data.
 
@@ -54,6 +54,11 @@ def get_data(local_srtm_dir=None, reduce_big_files=False):
     same size as the rest of the world. It may be used to save disk space 
     (because otherwise the tipical USA file is 25 megabaytes comparet to 2-3 
     megabytes for the rest).
+
+    If leave_zipped is True then files will be stored locally as compressed 
+    zip files. That means less disk space but more computing space for every 
+    file loaded. Note that if you leave this to True and don't reduce big 
+    files -- unzipping them on runtime will be slow.
     """
     if not local_srtm_dir:
         local_srtm_dir = get_default_srtm_dir()
@@ -84,4 +89,5 @@ def get_data(local_srtm_dir=None, reduce_big_files=False):
     assert srtm3_files
 
     return mod_data.GeoElevationData(srtm1_files, srtm3_files, local_srtm_dir,
-                                     reduce_big_files=reduce_big_files)
+                                     reduce_big_files=reduce_big_files,
+                                     leave_zipped=leave_zipped)
diff --git a/srtm/data.py b/srtm/data.py
@@ -25,7 +25,6 @@
 import re as mod_re
 import urllib as mod_urllib
 import os.path as mod_path
-import zipfile as mod_zipfile
 import cStringIO as mod_cstringio
 
 from . import utils as mod_utils
@@ -45,13 +44,14 @@ class GeoElevationData:
     files = None
 
     def __init__(self, srtm1_files, srtm3_files, files_directory,
-                 reduce_big_files=False):
+                 reduce_big_files=False, leave_zipped=False):
         self.srtm1_files = srtm1_files
         self.srtm3_files = srtm3_files
         # Place where local files are stored:
         self.files_directory = files_directory
 
         self.reduce_big_files = reduce_big_files
+        self.leave_zipped = leave_zipped
 
         self.files = {}
 
@@ -116,23 +116,24 @@ def retrieve_or_load_file_data(self, file_name):
         mod_logging.info('Retrieved {0}'.format(url))
         f.close()
 
+        if not data:
+            return None
+
         # data is zipped:
         mod_logging.info('Unzipping {0}'.format(url))
-        zip_file = mod_zipfile.ZipFile(mod_cstringio.StringIO(data))
-        zip_info_list = zip_file.infolist()
-        zip_info = zip_info_list[0]
-        data = zip_file.open(zip_info).read()
         mod_logging.info('Unzipped {0}'.format(url))
 
-        if not data:
-            return None
-
         if self.reduce_big_files:
+            data = mod_utils.unzip(data)
             data = self._reduce_file(data, file_name)
-
-        f = open(data_file_name, 'w')
-        f.write(data)
-        f.close()
+            data = mod_utils.zip(data)
+
+        with open(data_file_name, 'w') as f:
+            if self.leave_zipped:
+                f = open(data_file_name + '.zip', 'w')
+            else:
+                data = mod_utils.unzip(data)
+                f = open(data_file_name, 'w')
 
         return data
 
diff --git a/srtm/utils.py b/srtm/utils.py
@@ -14,7 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pdb
+
 import math as mod_math
+import zipfile as mod_zipfile
+import cStringIO as mod_cstringio
 
 ONE_DEGREE = 1000. * 10000.8 / 90.
 
@@ -38,3 +42,17 @@ def get_color_between(color1, color2, i):
     return (int(color1[0] + (color2[0] - color1[0]) * i),
             int(color1[1] + (color2[1] - color1[1]) * i),
             int(color1[2] + (color2[2] - color1[2]) * i))
+
+def zip(contents):
+    pdb.set_trace()
+    result = mod_cstringio.StringIO()
+    zip_file = mod_zipfile.ZipFile(result, 'w', mod_zipfile.ZIP_DEFLATED, False)
+    zip_file.writestr('archive.zip', contents)
+    result.seek(0)
+    return result.read()
+
+def unzip(contents):
+    zip_file = mod_zipfile.ZipFile(mod_cstringio.StringIO(contents))
+    zip_info_list = zip_file.infolist()
+    zip_info = zip_info_list[0]
+    return zip_file.open(zip_info).read()
