local files stored in custom directory

tkrajina · tkrajina · commit c4254233e048 · 2013-05-28T06:18:52.000+02:00
diff --git a/srtm/__init__.py b/srtm/__init__.py
@@ -14,31 +14,51 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pdb
+
 import cPickle  as mod_pickle
 import os       as mod_os
 import os.path  as mod_path
 
 from . import data      as mod_data
 from . import retriever as mod_retriever
 
-# Local cache path:
-LOCAL_FILES_DIRECTORY = '{0}/.geo_elevation_files'.format(mod_os.environ['HOME'])
+SRTM1_URL = 'http://dds.cr.usgs.gov/srtm/version2_1/SRTM1/'
+SRTM3_URL = 'http://dds.cr.usgs.gov/srtm/version2_1/SRTM3/'
 
-if not mod_path.exists(LOCAL_FILES_DIRECTORY):
-    mod_os.makedirs(LOCAL_FILES_DIRECTORY)
+def get_default_srtm_dir():
+    """ The default path to store files. """
+    # Local cache path:
+    if not mod_os.environ.has_key('HOME'):
+        raise Error('No default HOME directory found, please specify a path where to store files')
 
-FILES_LOCATION = '{0}/list'.format(LOCAL_FILES_DIRECTORY)
+    result = '{0}/.srtm'.format(mod_os.environ['HOME'])
 
-SRTM1_URL = 'http://dds.cr.usgs.gov/srtm/version2_1/SRTM1/'
-SRTM3_URL = 'http://dds.cr.usgs.gov/srtm/version2_1/SRTM3/'
+    if not mod_path.exists(result):
+        mod_os.makedirs(result)
 
-def get_data():
+    return result
+
+def get_data(local_srtm_dir=None):
     """
-    Load the urls of all geo-elevation files. If it can't be found -- it wil 
-    retrieve it and save to LOCAL_FILES_DIRECTORY.
+    Get the utility object for querying elevation data.
+
+    All data files will be stored in local_srtm_dir (note that it may be 
+    gigabytes of data so clean it from time to time).
+
+    On first run -- all files url will be stored and for every next elevation 
+    query if the SRTM file is not found in local_srtm_dir it will be retrieved 
+    and saved.
     """
+    if not local_srtm_dir:
+        local_srtm_dir = get_default_srtm_dir()
+
+    if not local_srtm_dir:
+        raise Error('Please specify a path where to store files')
+
+    files_list_file_name = '{0}/list'.format(local_srtm_dir)
     try:
-        f = open(FILES_LOCATION, 'r')
+        f = open(files_list_file_name, 'r')
         contents = f.read()
         f.close()
 
@@ -50,11 +70,11 @@ def get_data():
         srtm1_files = mod_retriever.retrieve_all_files_urls(SRTM1_URL)
         srtm3_files = mod_retriever.retrieve_all_files_urls(SRTM3_URL)
 
-        f = open(FILES_LOCATION, 'w')
+        f = open(files_list_file_name, 'w')
         f.write(mod_pickle.dumps({'srtm1': srtm1_files, 'srtm3': srtm3_files}))
         f.close()
 
     assert srtm1_files
     assert srtm3_files
 
-    return mod_data.GeoElevationData(srtm1_files, srtm3_files, LOCAL_FILES_DIRECTORY)
+    return mod_data.GeoElevationData(srtm1_files, srtm3_files, local_srtm_dir)
