Test out of range

DHI · Aug 23, 2023 · ca869a3 · ca869a3
1 parent e894920
commit ca869a3
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 19 deletions.
diff --git a/mikeio/_time.py b/mikeio/_time.py
@@ -7,15 +7,20 @@
 
 
 @dataclass
-class  DateTimeSelector:
+class DateTimeSelector:
     """Helper class for selecting time steps from a pandas DatetimeIndex"""
 
     index: pd.DatetimeIndex
 
     def __len__(self):
         return len(self.index)
 
-    def isel(self, x: Optional[int | Iterable[int] | str | datetime | pd.DatetimeIndex | slice]) -> List[int]:
+    def isel(
+        self,
+        x: Optional[
+            int | Iterable[int] | str | datetime | pd.DatetimeIndex | slice
+        ] = None,
+    ) -> List[int]:
         """Select time steps from a pandas DatetimeIndex
 
         Parameters
@@ -38,15 +43,15 @@ def isel(self, x: Optional[int | Iterable[int] | str | datetime | pd.DatetimeInd
         >>> dts.isel(-1)
         [3]
         """
-        
+
         indices = list(range(len(self.index)))
 
         if x is None:
             return indices
-        
+
         if isinstance(x, int):
             return [indices[x]]
-        
+
         if isinstance(x, (datetime, str)):
             loc = self.index.get_loc(x)
             if isinstance(loc, int):
@@ -60,28 +65,26 @@ def isel(self, x: Optional[int | Iterable[int] | str | datetime | pd.DatetimeInd
             else:
                 s = self.index.slice_indexer(x.start, x.stop)
                 return list(range(s.start, s.stop))
-        
+
         if isinstance(x, Iterable):
             return [self.isel(t)[0] for t in x]
-        
+
         return indices
 
 
 if __name__ == "__main__":
-        
+
     idx = pd.date_range("2000-01-01", periods=4, freq="D")
     assert len(idx) == 4
-    
+
     dts = DateTimeSelector(idx)
 
-    assert dts.isel(None) == [0,1,2,3]
+    assert dts.isel(None) == [0, 1, 2, 3]
     assert dts.isel(0) == [0]
     assert dts.isel(-1) == [3]
-    assert dts.isel([0,1]) == [0,1]
+    assert dts.isel([0, 1]) == [0, 1]
     assert dts.isel("2000-01-02") == [1]
-    assert dts.isel(["2000-01-02", "2000-01-03"]) == [1,2]
-    assert dts.isel(idx) == [0,1,2,3]
-    assert dts.isel(slice(1,4)) == [1,2,3]
-    assert dts.isel(slice("2000-01-02", "2000-01-04")) == [1,2,3]
-
-
+    assert dts.isel(["2000-01-02", "2000-01-03"]) == [1, 2]
+    assert dts.isel(idx) == [0, 1, 2, 3]
+    assert dts.isel(slice(1, 4)) == [1, 2, 3]
+    assert dts.isel(slice("2000-01-02", "2000-01-04")) == [1, 2, 3]
diff --git a/tests/test_time.py b/tests/test_time.py
@@ -1,14 +1,17 @@
 import pandas as pd
+import pytest
+from datetime import datetime
 
 from mikeio._time import DateTimeSelector
 
-def test_date_time_selector():
+def test_date_time_selector_valid():
 
     idx = pd.date_range("2000-01-01", periods=4, freq="D")
     assert len(idx) == 4
 
     dts = DateTimeSelector(idx)
 
+    assert dts.isel() == [0,1,2,3]
     assert dts.isel(None) == [0,1,2,3]
     assert dts.isel(0) == [0]
     assert dts.isel(-1) == [3]
@@ -17,4 +20,23 @@ def test_date_time_selector():
     assert dts.isel(["2000-01-02", "2000-01-03"]) == [1,2]
     assert dts.isel(idx) == [0,1,2,3]
     assert dts.isel(slice(1,4)) == [1,2,3]
-    assert dts.isel(slice("2000-01-02", "2000-01-04")) == [1,2,3]
+    assert dts.isel(slice("2000-01-02", "2000-01-04")) == [1,2,3]
+    assert dts.isel(datetime(2000,1,2)) == [1]
+
+def test_out_of_range_int():
+    idx = pd.date_range("2000-01-01", periods=4, freq="D")
+    dts = DateTimeSelector(idx)
+    with pytest.raises(IndexError):
+        dts.isel(4)
+
+def test_out_of_range_str():
+    idx = pd.date_range("2000-01-01", periods=4, freq="D")
+    dts = DateTimeSelector(idx)
+    with pytest.raises(KeyError):
+        dts.isel("2000-01-05")
+
+def test_out_of_range_datetime():
+    idx = pd.date_range("2000-01-01", periods=4, freq="D")
+    dts = DateTimeSelector(idx)
+    with pytest.raises(KeyError):
+        dts.isel(datetime(2000,1,5))