Merge pull request #200 from dstansby/no-untyped-epochs

No untyped epochs
MAVENSDC · May 26, 2023 · ae35bf3 · ae35bf3
2 parents e279197 + 9d4dfa2
commit ae35bf3
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 29 deletions.
diff --git a/cdflib/cdfread.py b/cdflib/cdfread.py
@@ -1,4 +1,3 @@
-import datetime
 import gzip
 import hashlib
 import io
@@ -375,8 +374,8 @@ def varget(
         self,
         variable: Optional[str] = None,
         epoch: Optional[str] = None,
-        starttime: Optional[datetime.datetime] = None,
-        endtime: Optional[datetime.datetime] = None,
+        starttime: Optional[epoch.epoch_types] = None,
+        endtime: Optional[epoch.epoch_types] = None,
         startrec: int = 0,
         endrec: Optional[int] = None,
     ) -> Union[str, np.ndarray]:
@@ -1792,8 +1791,8 @@ def _read_vardata(
         self,
         vdr_info: VDR,
         epoch: Optional[str] = None,
-        starttime: Optional[datetime.datetime] = None,
-        endtime: Optional[datetime.datetime] = None,
+        starttime: Optional[epoch.epoch_types] = None,
+        endtime: Optional[epoch.epoch_types] = None,
         startrec: int = 0,
         endrec: Optional[int] = None,
     ) -> Optional[Union[str, np.ndarray]]:
@@ -1838,7 +1837,7 @@ def _read_vardata(
             return data[0]
 
     def _findtimerecords(
-        self, var_name: str, starttime: datetime.datetime, endtime: datetime.datetime, epoch: Optional[str] = None
+        self, var_name: str, starttime: epoch.epoch_types, endtime: epoch.epoch_types, epoch: Optional[str] = None
     ) -> np.ndarray:
         if epoch is not None:
             vdr_info = self.varinq(epoch)
@@ -1875,7 +1874,7 @@ def _findtimerecords(
         return self._findrangerecords(vdr_info.Data_Type, epochtimes, starttime, endtime)
 
     def _findrangerecords(
-        self, data_type: int, epochtimes: epoch.epochs_type, starttime: datetime.datetime, endtime: datetime.datetime
+        self, data_type: int, epochtimes: epoch.epochs_type, starttime: epoch.epoch_types, endtime: epoch.epoch_types
     ) -> np.ndarray:
         if data_type == 31 or data_type == 32 or data_type == 33:
             # CDF_EPOCH or CDF_EPOCH16 or CDF_TIME_TT2000

diff --git a/cdflib/epochs.py b/cdflib/epochs.py
@@ -1,7 +1,6 @@
 import csv
 import datetime
 import math
-import numbers
 import os
 import re
 from typing import List, Optional, Tuple, Union
@@ -162,7 +161,6 @@ def _compose_date(
         milliseconds: Optional[npt.NDArray] = None,
         microseconds: Optional[npt.NDArray] = None,
         nanoseconds: Optional[npt.NDArray] = None,
-        *extras,
     ) -> npt.NDArray[np.datetime64]:
         """
         Take date components and return a numpy datetime array.
@@ -185,7 +183,7 @@ def to_datetime(cls, cdf_time: epoch_types) -> npt.NDArray[np.datetime64]:
         """
         times = cls.breakdown(cdf_time)
         times = np.atleast_2d(times)
-        return cls._compose_date(*times.T).astype("datetime64[us]")
+        return cls._compose_date(*times.T[:9]).astype("datetime64[us]")
 
     @staticmethod
     def unixtime(cdf_time: npt.ArrayLike) -> npt.NDArray:
@@ -266,7 +264,9 @@ def compute(datetimes: npt.ArrayLike) -> npt.NDArray:
         return np.squeeze(ret)
 
     @staticmethod
-    def findepochrange(epochs: epochs_type, starttime=None, endtime=None) -> np.ndarray:
+    def findepochrange(
+        epochs: epochs_type, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None
+    ) -> np.ndarray:
         """
         Finds the record range within the start and end time from values
         of a CDF epoch data type. It returns a list of record numbers.
@@ -490,15 +490,12 @@ def breakdown_tt2000(tt2000: cdf_tt2000_type) -> np.ndarray:
         return np.squeeze(toutcs.T)
 
     @staticmethod
-    def compute_tt2000(datetimes) -> npt.NDArray[np.int64]:
+    def compute_tt2000(datetimes: npt.ArrayLike) -> npt.NDArray[np.int64]:
         if not isinstance(datetimes, (list, tuple, np.ndarray)):
             raise TypeError("datetime must be in list form")
 
-        if isinstance(datetimes[0], numbers.Number):
-            datetimes = [datetimes]
-
-        new_datetimes = np.array(datetimes)
-        count = len(datetimes)
+        new_datetimes = np.atleast_2d(datetimes)
+        count = len(new_datetimes)
         nanoSecSinceJ2000s = []
         for x in range(count):
             datetime = new_datetimes[x]
@@ -723,7 +720,9 @@ def _EPOCHbreakdownTT2000(epoch: npt.ArrayLike) -> npt.NDArray:
         return date
 
     @staticmethod
-    def epochrange_tt2000(epochs: cdf_tt2000_type, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None):
+    def epochrange_tt2000(
+        epochs: cdf_tt2000_type, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None
+    ) -> npt.NDArray:
         if isinstance(epochs, int) or isinstance(epochs, np.int64):
             pass
         elif isinstance(epochs, list) or isinstance(epochs, tuple) or isinstance(epochs, np.ndarray):
@@ -830,7 +829,7 @@ def _JulianDay(y: int, m: int, d: int) -> int:
         return 367 * y - a1 - a2 + a3 + d + 1721029
 
     @staticmethod
-    def compute_epoch16(datetimes) -> npt.NDArray[np.complex128]:
+    def compute_epoch16(datetimes: npt.ArrayLike) -> npt.NDArray[np.complex128]:
         new_dates = np.atleast_2d(datetimes)
         count = len(new_dates)
         epochs = []
@@ -1061,7 +1060,7 @@ def _calc_from_julian(epoch0: npt.ArrayLike, epoch1: npt.ArrayLike) -> npt.NDArr
         return out
 
     @staticmethod
-    def breakdown_epoch16(epochs: cdf_epoch16_type):
+    def breakdown_epoch16(epochs: cdf_epoch16_type) -> npt.NDArray:
         """
         Calculate date and time from epochs
 
@@ -1072,7 +1071,7 @@ def breakdown_epoch16(epochs: cdf_epoch16_type):
 
         Returns
         -------
-        components : list
+        components : ndarray
             List or array of date and time values.  The last axis contains
             (in order): year, month, day, hour, minute, second, millisecond,
             microsecond, nanosecond, and picosecond
@@ -1150,7 +1149,7 @@ def _computeEpoch16(y: int, m: int, d: int, h: int, mn: int, s: int, ms: int, ms
 
     @staticmethod
     def epochrange_epoch16(
-        epochs: cdf_epoch16_type, starttime: Union[complex, np.complex128] = None, endtime=None
+        epochs: cdf_epoch16_type, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None
     ) -> Optional[np.ndarray]:
         new_epochs = np.atleast_1d(epochs)
 
@@ -1242,7 +1241,7 @@ def encode_epoch(epochs: cdf_epoch_type, iso_8601: bool = True) -> encoded_type:
         return encodeds
 
     @staticmethod
-    def _encodex_epoch(epoch: cdf_epoch_type, iso_8601: bool = True):
+    def _encodex_epoch(epoch: cdf_epoch_type, iso_8601: bool = True) -> str:
         components = CDFepoch.breakdown_epoch(epoch)
         if iso_8601:
             # year-mm-ddThh:mm:ss.mmm
@@ -1370,7 +1369,7 @@ def compute_epoch(dates: npt.ArrayLike) -> np.ndarray:
         return np.squeeze(epochs)
 
     @staticmethod
-    def _computeEpoch(y: int, m, d, h, mn, s, ms):
+    def _computeEpoch(y: int, m: int, d: int, h: int, mn: int, s: int, ms: int) -> float:
         if m == 0:
             daysSince0AD = CDFepoch._JulianDay(y, 1, 1) + (d - 1) - 1721060
         else:
@@ -1446,7 +1445,9 @@ def breakdown_epoch(epochs: cdf_epoch_type) -> np.ndarray:
         return np.squeeze(components)
 
     @staticmethod
-    def epochrange_epoch(epochs, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None) -> np.ndarray:
+    def epochrange_epoch(
+        epochs: epoch_types, starttime: Optional[epoch_types] = None, endtime: Optional[epoch_types] = None
+    ) -> np.ndarray:
         if isinstance(epochs, (float, np.float64)):
             pass
         elif isinstance(epochs, (list, tuple, np.ndarray)):
@@ -1480,10 +1481,7 @@ def epochrange_epoch(epochs, starttime: Optional[epoch_types] = None, endtime: O
         if stime > etime:
             raise ValueError("Invalid start/end time")
 
-        if isinstance(epochs, (list, tuple)):
-            new_epochs = np.array(epochs)
-        else:
-            new_epochs = epochs
+        new_epochs = np.array(epochs)
         return np.where(np.logical_and(new_epochs >= stime, new_epochs <= etime))[0]
 
     @staticmethod

diff --git a/mypy.ini b/mypy.ini
@@ -27,3 +27,5 @@ disallow_incomplete_defs = True
 
 [mypy-cdflib.epochs.*]
 disallow_untyped_calls = True
+disallow_incomplete_defs = True
+disallow_untyped_defs = True