From 28b26724e4077fd8dbf5154691f677d6024899a1 Mon Sep 17 00:00:00 2001
From: Henrik Andersson <jan@dhigroup.com>
Date: Fri, 10 May 2024 11:52:45 +0200
Subject: [PATCH] Consistency dataarray and dataset

---
 mikeio/dataset/_dataarray.py | 13 +++++++++----
 mikeio/dataset/_dataset.py   |  4 ++--
 tests/test_dataset.py        |  1 +
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/mikeio/dataset/_dataarray.py b/mikeio/dataset/_dataarray.py
index 4d98a92c8..e15a9b1bc 100644
--- a/mikeio/dataset/_dataarray.py
+++ b/mikeio/dataset/_dataarray.py
@@ -124,7 +124,12 @@ def __call__(self, tail: bool = True) -> "DataArray":
             geometry = GeometryUndefined()
 
         return DataArray(
-            data=Hm0, time=self.da.time, item=item, dims=dims, geometry=geometry, dt=self.da._dt
+            data=Hm0,
+            time=self.da.time,
+            item=item,
+            dims=dims,
+            geometry=geometry,
+            dt=self.da._dt,
         )
 
 
@@ -423,11 +428,11 @@ def is_equidistant(self) -> bool:
         return len(self.time.to_series().diff().dropna().unique()) == 1
 
     @property
-    def timestep(self) -> float | None:
+    def timestep(self) -> float:
         """Time step in seconds if equidistant (and at
-        least two time instances); otherwise None
+        least two time instances); otherwise original time step is returned.
         """
-        dt = None
+        dt = self._dt
         if len(self.time) > 1 and self.is_equidistant:
             first: pd.Timestamp = self.time[0]
             second: pd.Timestamp = self.time[1]
diff --git a/mikeio/dataset/_dataset.py b/mikeio/dataset/_dataset.py
index 2b84324b8..6da0ff971 100644
--- a/mikeio/dataset/_dataset.py
+++ b/mikeio/dataset/_dataset.py
@@ -333,9 +333,9 @@ def end_time(self) -> datetime:
         return self.time[-1].to_pydatetime()  # type: ignore
 
     @property
-    def timestep(self) -> float | None:
+    def timestep(self) -> float:
         """Time step in seconds if equidistant (and at
-        least two time instances); otherwise None
+        least two time instances); otherwise original time step is returned.
         """
         dt = self._dt
         if len(self.time) > 1 and self.is_equidistant:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 2e7854803..82576fcec 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -1616,6 +1616,7 @@ def test_select_single_timestep_preserves_dt():
     assert ds.timestep == pytest.approx(1800.0)
     ds2 = ds.isel(time=-1)
     assert ds2.timestep == pytest.approx(1800.0)
+    assert ds2[0].timestep == pytest.approx(1800.0)
 
 
 def test_select_multiple_spaced_timesteps_uses_proper_dt(tmp_path):