Merge pull request #201 from dstansby/no-untyped-astro-epochs

No untyped astro epochs

cdflib/epochs_astropy.py

@@ -7,9 +7,10 @@
 """
 import datetime
 from datetime import timezone
-from typing import List
+from typing import List, Union, Optional
 
 import numpy as np
+import numpy.typing as npt
 from astropy.time import Time
 from astropy.time.formats import TimeFromEpoch, erfa
 
@@ -53,7 +54,7 @@ class CDFAstropy:
     increment = 0
 
     @staticmethod
-    def convert_to_astropy(epochs, format=None) -> Time:
+    def convert_to_astropy(epochs: Union[Time, npt.ArrayLike], format: Optional[str]=None) -> Time:
         """
         Convert CDF epochs to astropy time objects.
 
@@ -82,7 +83,7 @@ def convert_to_astropy(epochs, format=None) -> Time:
             raise TypeError("Not sure how to handle type {}".format(type(epochs)))
 
     @staticmethod
-    def encode(epochs, iso_8601: bool = True):  # @NoSelf
+    def encode(epochs: npt.ArrayLike, iso_8601: bool = True) -> npt.NDArray[np.str_]:
         epochs = CDFAstropy.convert_to_astropy(epochs)
         if iso_8601:
             return epochs.iso
@@ -102,7 +103,7 @@ def breakdown(epochs):
         raise TypeError("Not sure how to handle type {}".format(type(epochs)))
 
     @staticmethod
-    def to_datetime(cdf_time) -> Time:
+    def to_datetime(cdf_time: npt.ArrayLike) -> Time:
         cdf_time = CDFAstropy.convert_to_astropy(cdf_time)
         return cdf_time.datetime
 
@@ -116,9 +117,8 @@ def unixtime(cdf_time):  # @NoSelf
         return epochs.unix
 
     @staticmethod
-    def compute(datetimes):
-        if not isinstance(datetimes[0], list):
-            datetimes = [datetimes]
+    def compute(datetimes: npt.ArrayLike) -> npt.NDArray:
+        datetimes = np.atleast_2d(datetimes)
         cdf_time = []
         for d in datetimes:
             unix_seconds = datetime.datetime(d[0], d[1], d[2], d[3], d[4], d[5]).replace(tzinfo=timezone.utc).timestamp()
@@ -134,7 +134,7 @@ def compute(datetimes):
                 remainder_seconds = (d[6] / 1000.0) + (d[7] / 1000000.0) + (d[8] / 1000000000.0) + (d[9] / 1000000000000.0)
                 astrotime = Time(unix_seconds, remainder_seconds, format="unix", precision=9)
                 cdf_time.append(astrotime.cdf_epoch16)
-        return np.array(cdf_time)
+        return np.squeeze(cdf_time)
 
     @staticmethod
     def findepochrange(epochs, starttime=None, endtime=None):  # @NoSelf
@@ -150,10 +150,8 @@ def findepochrange(epochs, starttime=None, endtime=None):  # @NoSelf
         return min(indices[0]), max(indices[0])
 
     @staticmethod
-    def breakdown_tt2000(tt2000):
-        tt2000strings = tt2000.iso
-        if not isinstance(tt2000strings, (list, np.ndarray)):
-            tt2000strings = [tt2000strings]
+    def breakdown_tt2000(tt2000: Time) -> npt.NDArray:
+        tt2000strings = np.atleast_1d(tt2000.iso)
         times = []
         for t in tt2000strings:
             date, time = t.split(" ")
@@ -177,13 +175,11 @@ def breakdown_tt2000(tt2000):
             time_as_list.append(int(nanoseconds))  # microseconds
             times.append(time_as_list)
 
-        return times
+        return np.squeeze(times)
 
     @staticmethod
-    def breakdown_epoch16(epochs):
-        epoch16strings = epochs.iso
-        if not isinstance(epoch16strings, (list, np.ndarray)):
-            epoch16strings = [epoch16strings]
+    def breakdown_epoch16(epochs: Time) -> npt.NDArray:
+        epoch16strings = np.atleast_1d(epochs.iso)
         times = []
         for t in epoch16strings:
             time_as_list: List[int] = []
@@ -210,13 +206,12 @@ def breakdown_epoch16(epochs):
             time_as_list.append(int(picoseconds))  # picoseconds
             times.append(time_as_list)
 
-        return times
+        return np.squeeze(times)
 
     @staticmethod
-    def breakdown_epoch(epochs):
-        epochstrings = epochs.iso
-        if not isinstance(epochstrings, (list, np.ndarray)):
-            epochstrings = [epochstrings]
+    def breakdown_epoch(epochs: Time) -> npt.NDArray:
+        epochstrings = np.atleast_1d(epochs.iso)
+
         times = []
         for t in epochstrings:
             date, time = t.split(" ")
@@ -235,7 +230,7 @@ def breakdown_epoch(epochs):
             milliseconds = decimal_seconds * 1000
             time_as_list.append(int(milliseconds))  # milliseconds
             times.append(time_as_list)
-        return times
+        return np.squeeze(times)
 
     @staticmethod
     def parse(value):

mypy.ini

@@ -29,3 +29,7 @@ disallow_incomplete_defs = True
 disallow_untyped_calls = True
 disallow_incomplete_defs = True
 disallow_untyped_defs = True
+
+[mypy-cdflib.epochs_astropy.*]
+disallow_untyped_calls = True
+disallow_incomplete_defs = True

tests/test_astropy_epochs.py

@@ -67,31 +67,31 @@ def test_breakdown_cdfepoch():
 
 def test_breakdown_cdfepoch16():
     x = cdfepoch.breakdown(np.complex128(63300946758.000000 + 176214648000.00000j))
-    assert x[0][0] == 2005
-    assert x[0][1] == 12
-    assert x[0][2] == 4
-    assert x[0][3] == 20
-    assert x[0][4] == 19
-    assert x[0][5] == 18
-    assert x[0][6] == 176
-    assert x[0][7] == 214
-    assert x[0][8] == 648
-    assert x[0][9] == 0
+    assert x[0] == 2005
+    assert x[1] == 12
+    assert x[2] == 4
+    assert x[3] == 20
+    assert x[4] == 19
+    assert x[5] == 18
+    assert x[6] == 176
+    assert x[7] == 214
+    assert x[8] == 648
+    assert x[9] == 0
 
 
 def test_breakdown_cdftt2000():
     x = cdfepoch.breakdown(123456789101112131)
-    assert x[0][0] == 2003
-    assert x[0][1] == 11
-    assert x[0][2] == 30
-    assert x[0][3] == 9
-    assert x[0][4] == 33
-    assert x[0][5] == 9
-    assert x[0][6] == 101
-    assert x[0][7] == 112
+    assert x[0] == 2003
+    assert x[1] == 11
+    assert x[2] == 30
+    assert x[3] == 9
+    assert x[4] == 33
+    assert x[5] == 9
+    assert x[6] == 101
+    assert x[7] == 112
 
     # Apparently there is a loss of precision at this level
-    # assert x[0][8] == 131
+    # assert x[8] == 131
 
 
 def test_compute_cdfepoch():
@@ -108,7 +108,7 @@ def test_compute_cdfepoch():
     random_time.append(randint(0, 999))  # Millisecond
     x = cdfepoch.breakdown(cdfepoch.compute(random_time))
     i = 0
-    for t in x[0]:
+    for t in x:
         assert t == random_time[i], f"Time {random_time} was not equal to {x}"
         i += 1
 
@@ -128,7 +128,7 @@ def test_compute_cdfepoch16():
     cdftime = cdfepoch.convert_to_astropy(cdfepoch.compute(random_time), format="cdf_epoch16")
     x = cdfepoch.breakdown(cdftime)
     i = 0
-    for t in x[0]:
+    for t in x:
         assert t == random_time[i], f"Time {random_time} was not equal to {x}"
         i += 1
         # Unfortunately, currently there is a pretty big loss of precision that comes with
@@ -170,7 +170,7 @@ def test_parse_cdfepoch():
 
 def test_parse_cdfepoch16():
     input_time = 53467976543.0 + 543218654100j
-    x = cdfepoch.encode(input_time)
+    x = str(cdfepoch.encode(input_time))
     assert x == "1694-05-01 07:42:23.543218654"
     add_precision = x + "000"
     parsed = cdfepoch.parse(add_precision)