diff --git a/pymeos/pymeos/collections/time/tstzspanset.py b/pymeos/pymeos/collections/time/tstzspanset.py index 9e938d7c..624e22f2 100644 --- a/pymeos/pymeos/collections/time/tstzspanset.py +++ b/pymeos/pymeos/collections/time/tstzspanset.py @@ -157,7 +157,7 @@ def timestamp_n(self, n: int) -> datetime: raise IndexError(f"Index {n} out of bounds") return timestamptz_to_datetime(tstzspanset_timestamptz_n(self._inner, n + 1)) - def timestamps(self) -> List[datetime]: + def timestamps(self) -> TsTzSet: """ Returns the list of distinct timestamps in ``self``. Returns: @@ -166,8 +166,9 @@ def timestamps(self) -> List[datetime]: MEOS Functions: spanset_timestamptzs """ - ts, count = tstzspanset_timestamps(self._inner) - return [timestamptz_to_datetime(ts[i]) for i in range(count)] + from .tstzset import TsTzSet + + return TsTzSet(_inner=tstzspanset_timestamps(self._inner)) def start_span(self) -> TsTzSpan: """ diff --git a/pymeos/tests/collections/time/dateset_test.py b/pymeos/tests/collections/time/dateset_test.py index fa5ce36b..21caa07e 100644 --- a/pymeos/tests/collections/time/dateset_test.py +++ b/pymeos/tests/collections/time/dateset_test.py @@ -68,7 +68,7 @@ def test_to_spanset(self): ) spanset = self.date_set.to_spanset() - + assert spanset == expected diff --git a/pymeos/tests/collections/time/tstzspanset_test.py b/pymeos/tests/collections/time/tstzspanset_test.py index f7dd9458..777ce227 100644 --- a/pymeos/tests/collections/time/tstzspanset_test.py +++ b/pymeos/tests/collections/time/tstzspanset_test.py @@ -6,6 +6,7 @@ from pymeos import ( TsTzSpan, + TsTzSet, TsTzSpanSet, TFloatInst, TFloatSeq, @@ -123,12 +124,14 @@ def test_timestamp_n(self): ) def test_timestamps(self): - assert self.tstzspanset.timestamps() == [ - datetime(2019, 9, 1, tzinfo=timezone.utc), - datetime(2019, 9, 2, tzinfo=timezone.utc), - datetime(2019, 9, 3, tzinfo=timezone.utc), - datetime(2019, 9, 4, tzinfo=timezone.utc), - ] + assert self.tstzspanset.timestamps() == TsTzSet( + elements=[ + datetime(2019, 9, 1, tzinfo=timezone.utc), + datetime(2019, 9, 2, tzinfo=timezone.utc), + datetime(2019, 9, 3, tzinfo=timezone.utc), + datetime(2019, 9, 4, tzinfo=timezone.utc), + ] + ) def test_num_tstzspans(self): assert self.tstzspanset.num_spans() == 2 diff --git a/pymeos/tests/main/tfloat_test.py b/pymeos/tests/main/tfloat_test.py index 2545632d..bf3bc41f 100644 --- a/pymeos/tests/main/tfloat_test.py +++ b/pymeos/tests/main/tfloat_test.py @@ -2052,10 +2052,7 @@ def test_derivative(self, temporal, expected): @pytest.mark.parametrize( "temporal", - [ - tfi, - tfds - ], + [tfi, tfds], ids=["Instant", "Discrete Sequence"], ) def test_derivative_without_linear_interpolation_raises(self, temporal):