0.2.4 allow user to specify the database file dir to download to

Author: MacHu-GWU

uszipcode/_version.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 
 if __name__ == "__main__":  # pragma: no cover
     print(__version__)

uszipcode/db.py

@@ -23,45 +23,51 @@
     from uszipcode.pkg.atomicwrites import atomic_write
     from uszipcode.pkg.sqlalchemy_mate import engine_creator
 
-db_file_dir = Path("/tmp")
-db_file_dir.mkdir(exist_ok=True)
 
-simple_db_file_path = db_file_dir.append_parts("uszipcode_simple_db.sqlite")
-db_file_path = db_file_dir.append_parts("uszipcode_db.sqlite")
+def get_simple_db_file_path(db_file_dir):
+    return Path(db_file_dir, "simple_db.sqlite")
 
 
-def is_simple_db_file_exists():
+def get_db_file_path(db_file_dir):
+    return Path(db_file_dir, "db.sqlite")
+
+
+def is_simple_db_file_exists(db_file_dir):
+    simple_db_file_path = get_simple_db_file_path(db_file_dir)
     if simple_db_file_path.exists():
         if simple_db_file_path.size >= 5 * 1000 * 1000:
             return True
     return False
 
 
-def is_db_file_exists():
+def is_db_file_exists(db_file_dir):
+    db_file_path = get_db_file_path(db_file_dir)
     if db_file_path.exists():
         if db_file_path.size >= 100 * 1000 * 1000:
             return True
     return False
 
 
-def connect_to_simple_zipcode_db():
-    return engine_creator.create_sqlite(path=simple_db_file_path.abspath)
+def connect_to_simple_zipcode_db(db_file_dir):
+    return engine_creator.create_sqlite(
+        path=get_simple_db_file_path(db_file_dir).abspath)
 
 
-def connect_to_zipcode_db():
-    return engine_creator.create_sqlite(path=db_file_path.abspath)
+def connect_to_zipcode_db(db_file_dir):
+    return engine_creator.create_sqlite(
+        path=get_db_file_path(db_file_dir).abspath)
 
 
-def download_simple_db_file():
+def download_simple_db_file(db_file_dir):
     simple_db_file_download_url = "https://datahub.io/machu-gwu/uszipcode-0.2.0-simple_db/r/simple_db.sqlite"
 
-    if not is_simple_db_file_exists():
+    if not is_simple_db_file_exists(db_file_dir):
         print("Start downloading data for simple zipcode database, total size 9MB ...")
         response = requests.get(simple_db_file_download_url, stream=True)
         chunk_size = 1 * 1024 ** 2
 
         counter = 0
-        with atomic_write(simple_db_file_path.abspath, mode="wb", overwrite=True) as f:
+        with atomic_write(get_simple_db_file_path(db_file_dir).abspath, mode="wb", overwrite=True) as f:
             for chunk in response.iter_content(chunk_size):
                 if not chunk:
                     break
@@ -71,16 +77,16 @@ def download_simple_db_file():
         print("  Complete!")
 
 
-def download_db_file():
+def download_db_file(db_file_dir):
     db_file_download_url = "https://datahub.io/machu-gwu/uszipcode-0.2.0-db/r/db.sqlite"
-    if not is_db_file_exists():
+    if not is_db_file_exists(db_file_dir):
         print(
             "Start downloading data for rich info zipcode database, total size 450+MB ...")
         response = requests.get(db_file_download_url, stream=True)
         chunk_size = 10 * 1024 ** 2
 
         counter = 0
-        with atomic_write(db_file_path.abspath, mode="wb", overwrite=True) as f:
+        with atomic_write(get_db_file_path(db_file_dir).abspath, mode="wb", overwrite=True) as f:
             for chunk in response.iter_content(chunk_size):
                 if not chunk:
                     break

uszipcode/search.py

@@ -12,6 +12,7 @@
 from collections import OrderedDict
 from sqlalchemy.orm import sessionmaker
 from six import integer_types, string_types
+from pathlib_mate import Path
 
 from .db import (
     is_simple_db_file_exists, is_db_file_exists,
@@ -33,14 +34,23 @@
 """
 
 
+HOME = Path.home().abspath
+HOME_USZIPCODE = Path(HOME, ".uszipcode").abspath
+TMP = "/tmp"
+
+
 class SearchEngine(object):
     """
     Zipcode Search Engine.
 
-    :param simple_zipcode: bool, default True, if True, use the simple zipcode
+    :type simple_zipcode: bool
+    :param simple_zipcode: default True, if True, use the simple zipcode
         db. Rich Demographics, Real Estate, Employment, Education info is not
         available. If False, use the rich info database.
 
+    :type db_file_dir: str
+    :param db_file_dir: where you want to download the sqlite database to.
+
     Usage::
 
         >>> search = SearchEngine()
@@ -75,16 +85,18 @@ class SearchEngine(object):
     _state_to_city_mapper = None
     _city_to_state_mapper = None
 
-    def __init__(self, simple_zipcode=True):
+    def __init__(self, simple_zipcode=True, db_file_dir=HOME_USZIPCODE):
+        Path(db_file_dir).mkdir(exist_ok=True)
+
         if simple_zipcode:
-            if not is_simple_db_file_exists():
-                download_simple_db_file()
-            engine = connect_to_simple_zipcode_db()
+            if not is_simple_db_file_exists(db_file_dir):
+                download_simple_db_file(db_file_dir)
+            engine = connect_to_simple_zipcode_db(db_file_dir)
             self.zip_klass = SimpleZipcode
         else:  # pragma: no cover
-            if not is_db_file_exists():
-                download_db_file()
-            engine = connect_to_zipcode_db()
+            if not is_db_file_exists(db_file_dir):
+                download_db_file(db_file_dir)
+            engine = connect_to_zipcode_db(db_file_dir)
             self.zip_klass = Zipcode
         self.engine = engine
         self.ses = sessionmaker(bind=engine)()