Move dastcom to sbpy.data.utils

Author: shreyasbapat

sbpy/data/orbit.py

@@ -17,7 +17,7 @@
 
 from .. import bib
 from . import conf, DataClass
-from .. import utils
+from . import utils
 
 __all__ = ['Orbit']
 
@@ -135,14 +135,14 @@ def from_horizons(cls, targetids, id_type='smallbody',
         return cls.from_table(all_elem)
 
     @classmethod
-    def from_dastcom5(cls, name):
+    def from_dastcom5(cls, identifier, is_record_number=False):
         """Load orbital elements from the DASTCOM5 Database
         (ftp://ssd.jpl.nasa.gov/pub/ssd/dastcom5.zip).
 
         Parameters
         ----------
-        name: str, mandatory
-            Name of NEO
+        identifier: mandatory
+            Names, numbers, or designations of objects to be queried
 
         Returns
         -------
@@ -158,7 +158,10 @@ def from_dastcom5(cls, name):
         if not os.path.isdir(dastcom5_dir):
             utils.dastcom5.download_dastcom5()
 
-        tb = utils.dastcom5.orbit_from_name(name=name)
+        if is_record_number:
+            tb = utils.dastcom5.orbit_from_record(record=identifier)
+        else:
+            tb = utils.dastcom5.orbit_from_name(name=identifier)
         return cls.from_table(tb)
 
     @classmethod

sbpy/data/utils/__init__.py

@@ -0,0 +1,6 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+# This sub-module is destined for common non-package specific utility
+# functions.
+
+from sbpy.data.utils import dastcom5

sbpy/utils/dastcom5.py renamed to sbpy/data/utils/dastcom5.py

@@ -1,5 +1,6 @@
 import os
 import re
+import shutil
 import urllib.request
 import zipfile
 
@@ -509,7 +510,7 @@ def read_record(record):
     return body_data
 
 
-def download_dastcom5():
+def download_dastcom5(update=False):
     """Downloads DASTCOM5 database.
 
     Downloads and unzip DASTCOM5 file in default sbpy path (~/.sbpy).
@@ -519,10 +520,15 @@ def download_dastcom5():
     dastcom5_dir = os.path.join(SBPY_LOCAL_PATH, "dastcom5")
     dastcom5_zip_path = os.path.join(SBPY_LOCAL_PATH, "dastcom5.zip")
 
-    if os.path.isdir(dastcom5_dir):
+    if os.path.isdir(dastcom5_dir) and not update:
         raise FileExistsError(
             "dastcom5 is already created in " + os.path.abspath(dastcom5_dir)
         )
+
+    if os.path.isdir(dastcom5_dir) and update:
+        shutil.rmtree(dastcom5_dir)
+
+    print("Downloading the latest DASTCOM5 Database")
     if not zipfile.is_zipfile(dastcom5_zip_path):
         if not os.path.isdir(SBPY_LOCAL_PATH):
             os.makedirs(SBPY_LOCAL_PATH)
@@ -534,6 +540,7 @@ def download_dastcom5():
         )
     with zipfile.ZipFile(dastcom5_zip_path) as myzip:
         myzip.extractall(SBPY_LOCAL_PATH)
+    os.remove(dastcom5_zip_path)
 
 
 def show_download_progress(transferred, block, totalsize):

sbpy/data/utils/tests/__init__.py

@@ -4,25 +4,25 @@
 import numpy as np
 import pytest
 
-from sbpy.utils import dastcom5
+from sbpy.data.utils import dastcom5
 
 
-@mock.patch("sbpy.utils.dastcom5.np.fromfile")
-@mock.patch("sbpy.utils.dastcom5.open")
+@mock.patch("sbpy.data.utils.dastcom5.np.fromfile")
+@mock.patch("sbpy.data.utils.dastcom5.open")
 def test_asteroid_db_is_called_with_right_path(mock_open, mock_np_fromfile):
     dastcom5.asteroid_db()
     mock_open.assert_called_with(dastcom5.AST_DB_PATH, "rb")
 
 
-@mock.patch("sbpy.utils.dastcom5.np.fromfile")
-@mock.patch("sbpy.utils.dastcom5.open")
+@mock.patch("sbpy.data.utils.dastcom5.np.fromfile")
+@mock.patch("sbpy.data.utils.dastcom5.open")
 def test_comet_db_is_called_with_right_path(mock_open, mock_np_fromfile):
     dastcom5.comet_db()
     mock_open.assert_called_with(dastcom5.COM_DB_PATH, "rb")
 
 
-@mock.patch("sbpy.utils.dastcom5.np.fromfile")
-@mock.patch("sbpy.utils.dastcom5.open")
+@mock.patch("sbpy.data.utils.dastcom5.np.fromfile")
+@mock.patch("sbpy.data.utils.dastcom5.open")
 def test_read_headers(mock_open, mock_np_fromfile):
     dastcom5.read_headers()
     mock_open.assert_any_call(
@@ -33,9 +33,9 @@ def test_read_headers(mock_open, mock_np_fromfile):
     )
 
 
-@mock.patch("sbpy.utils.dastcom5.read_headers")
-@mock.patch("sbpy.utils.dastcom5.np.fromfile")
-@mock.patch("sbpy.utils.dastcom5.open")
+@mock.patch("sbpy.data.utils.dastcom5.read_headers")
+@mock.patch("sbpy.data.utils.dastcom5.np.fromfile")
+@mock.patch("sbpy.data.utils.dastcom5.open")
 def test_read_record(mock_open, mock_np_fromfile, mock_read_headers):
     mocked_ast_headers = np.array(
         [(3184, -1, b"00740473", b"00496815")],
@@ -59,10 +59,10 @@ def test_read_record(mock_open, mock_np_fromfile, mock_read_headers):
     )
 
 
-@mock.patch("sbpy.utils.dastcom5.os.makedirs")
-@mock.patch("sbpy.utils.dastcom5.zipfile")
-@mock.patch("sbpy.utils.dastcom5.os.path.isdir")
-@mock.patch("sbpy.utils.dastcom5.urllib.request")
+@mock.patch("sbpy.data.utils.dastcom5.os.makedirs")
+@mock.patch("sbpy.data.utils.dastcom5.zipfile")
+@mock.patch("sbpy.data.utils.dastcom5.os.path.isdir")
+@mock.patch("sbpy.data.utils.dastcom5.urllib.request")
 def test_download_dastcom5_raises_error_when_folder_exists(
     mock_request, mock_isdir, mock_zipfile, mock_makedirs
 ):
@@ -76,10 +76,10 @@ def test_download_dastcom5_raises_error_when_folder_exists(
     )
 
 
-@mock.patch("sbpy.utils.dastcom5.urllib.request")
-@mock.patch("sbpy.utils.dastcom5.os.makedirs")
-@mock.patch("sbpy.utils.dastcom5.zipfile")
-@mock.patch("sbpy.utils.dastcom5.os.path.isdir")
+@mock.patch("sbpy.data.utils.dastcom5.urllib.request")
+@mock.patch("sbpy.data.utils.dastcom5.os.makedirs")
+@mock.patch("sbpy.data.utils.dastcom5.zipfile")
+@mock.patch("sbpy.data.utils.dastcom5.os.path.isdir")
 def test_download_dastcom5_creates_folder(
     mock_isdir, mock_zipfile, mock_makedirs, mock_request
 ):
@@ -89,9 +89,9 @@ def test_download_dastcom5_creates_folder(
     mock_makedirs.assert_called_once_with(dastcom5.SBPY_LOCAL_PATH)
 
 
-@mock.patch("sbpy.utils.dastcom5.zipfile")
-@mock.patch("sbpy.utils.dastcom5.os.path.isdir")
-@mock.patch("sbpy.utils.dastcom5.urllib.request.urlretrieve")
+@mock.patch("sbpy.data.utils.dastcom5.zipfile")
+@mock.patch("sbpy.data.utils.dastcom5.os.path.isdir")
+@mock.patch("sbpy.data.utils.dastcom5.urllib.request.urlretrieve")
 def test_download_dastcom5_downloads_file(
     mock_request, mock_isdir, mock_zipfile
 ):

sbpy/utils/tests/test_dastcom5.py renamed to sbpy/data/utils/tests/test_dastcom5.py

@@ -3,4 +3,4 @@
 # This sub-module is destined for common non-package specific utility
 # functions.
 
-from sbpy.utils import dastcom5
+from .core import *