Added tests for get and more tests for getitem

tests/test_core.py

@@ -2,18 +2,20 @@ import datetime
 import random
 from typing import Literal, Mapping, Sequence
 
+import pytest
 from fincal.core import AllFrequencies, Frequency, Series, TimeSeriesCore
 from fincal.fincal import create_date_series
+from fincal.utils import FincalOptions
 
 
 class TestFrequency:
     def test_creation(self):
-        D = Frequency('daily', 'days', 1, 1, 'D')
+        D = Frequency("daily", "days", 1, 1, "D")
         assert D.days == 1
-        assert D.symbol == 'D'
+        assert D.symbol == "D"
-        assert D.name == 'daily'
+        assert D.name == "daily"
         assert D.value == 1
-        assert D.freq_type == 'days'
+        assert D.freq_type == "days"
 
 
 def create_test_data(
@@ -60,9 +62,9 @@ def create_test_data(
 
 class TestAllFrequencies:
     def test_attributes(self):
-        assert hasattr(AllFrequencies, 'D')
+        assert hasattr(AllFrequencies, "D")
-        assert hasattr(AllFrequencies, 'M')
+        assert hasattr(AllFrequencies, "M")
-        assert hasattr(AllFrequencies, 'Q')
+        assert hasattr(AllFrequencies, "Q")
 
     def test_days(self):
         assert AllFrequencies.D.days == 1
@@ -70,64 +72,135 @@ class TestAllFrequencies:
         assert AllFrequencies.Q.days == 91
 
     def test_symbol(self):
-        assert AllFrequencies.H.symbol == 'H'
+        assert AllFrequencies.H.symbol == "H"
-        assert AllFrequencies.W.symbol == 'W'
+        assert AllFrequencies.W.symbol == "W"
 
     def test_values(self):
         assert AllFrequencies.H.value == 6
         assert AllFrequencies.Y.value == 1
 
     def test_type(self):
-        assert AllFrequencies.Q.freq_type == 'months'
+        assert AllFrequencies.Q.freq_type == "months"
-        assert AllFrequencies.W.freq_type == 'days'
+        assert AllFrequencies.W.freq_type == "days"
 
 
 class TestSeries:
     def test_creation(self):
-        series = Series([1, 2, 3, 4, 5, 6, 7], data_type='number')
+        series = Series([1, 2, 3, 4, 5, 6, 7], data_type="number")
         assert series.dtype == float
         assert series[2] == 3
 
-        dates = create_date_series('2021-01-01', '2021-01-31', frequency='D')
+        dates = create_date_series("2021-01-01", "2021-01-31", frequency="D")
-        series = Series(dates, data_type='date')
+        series = Series(dates, data_type="date")
         assert series.dtype == datetime.datetime
 
 
 class TestTimeSeriesCore:
-    data = [('2021-01-01', 220), ('2021-02-01', 230), ('2021-03-01', 240)]
+    data = [("2021-01-01", 220), ("2021-02-01", 230), ("2021-03-01", 240)]
 
     def test_repr_str(self):
-        ts = TimeSeriesCore(self.data, frequency='M')
+        ts = TimeSeriesCore(self.data, frequency="M")
-        assert str(ts) in repr(ts).replace('\t', ' ')
+        assert str(ts) in repr(ts).replace("\t", " ")
 
         data = create_test_data(frequency="D", eomonth=False, n=50, gaps=0, month_position="start", date_as_str=True)
         ts = TimeSeriesCore(data, frequency="D")
-        assert '...' in str(ts)
+        assert "..." in str(ts)
-        assert '...' in repr(ts)
+        assert "..." in repr(ts)
 
     def test_creation(self):
-        ts = TimeSeriesCore(self.data, frequency='M')
+        ts = TimeSeriesCore(self.data, frequency="M")
         assert isinstance(ts, TimeSeriesCore)
         assert isinstance(ts, Mapping)
 
     def test_getitem(self):
-        ts = TimeSeriesCore(self.data, frequency='M')
+        ts = TimeSeriesCore(self.data, frequency="M")
         assert ts.dates[0] == datetime.datetime(2021, 1, 1, 0, 0)
         assert ts.values[0] == 220
-        assert ts['2021-01-01'][1] == 220
+        assert ts["2021-01-01"][1] == 220
-        assert len(ts[ts.dates > '2021-01-01']) == 2
+        assert len(ts[ts.dates > "2021-01-01"]) == 2
-        assert ts[ts.dates == '2021-02-01'].iloc[0][1] == 230
+        assert ts[ts.dates == "2021-02-01"].iloc[0][1] == 230
         assert ts.iloc[2][0] == datetime.datetime(2021, 3, 1)
         assert len(ts.iloc[:2]) == 2
+        with pytest.raises(KeyError):
+            ts["2021-02-03"]
+        subset_ts = ts[["2021-01-01", "2021-03-01"]]
+        assert len(subset_ts) == 2
+        assert isinstance(subset_ts, TimeSeriesCore)
+        assert subset_ts.iloc[1][1] == 240
+
+    def test_get(self):
+        ts = TimeSeriesCore(self.data, frequency="M")
+        assert ts.dates[0] == datetime.datetime(2021, 1, 1, 0, 0)
+        assert ts.values[0] == 220
+        assert ts.get("2021-01-01")[1] == 220
+        assert ts.get("2021-02-15") is None
+        assert ts.get("2021-02-23", -1) == -1
+        assert ts.get("2021-02-10", closest="previous")[1] == 230
+        assert ts.get("2021-02-10", closest="next")[1] == 240
+        FincalOptions.get_closest = "previous"
+        assert ts.get("2021-02-10")[1] == 230
+        FincalOptions.get_closest = "next"
+        assert ts.get("2021-02-10")[1] == 240
 
     def test_contains(self):
-        ts = TimeSeriesCore(self.data, frequency='M')
+        ts = TimeSeriesCore(self.data, frequency="M")
         assert datetime.datetime(2021, 1, 1) in ts
-        assert '2021-01-01' in ts
+        assert "2021-01-01" in ts
-        assert '2021-01-14' not in ts
+        assert "2021-01-14" not in ts
 
     def test_items(self):
-        ts = TimeSeriesCore(self.data, frequency='M')
+        ts = TimeSeriesCore(self.data, frequency="M")
         for i, j in ts.items():
             assert j == self.data[0][1]
             break
+
+    def test_special_keys(self):
+        ts = TimeSeriesCore(self.data, frequency="M")
+        dates = ts["dates"]
+        values = ts["values"]
+        assert isinstance(dates, Series)
+        assert isinstance(values, Series)
+        assert len(dates) == 3
+        assert len(values) == 3
+        assert dates[0] == datetime.datetime(2021, 1, 1, 0, 0)
+        assert values[0] == 220
+
+
+class TestTimeSeriesCoreHeadTail:
+    data = [
+        ("2021-01-01", 220),
+        ("2021-02-01", 230),
+        ("2021-03-01", 240),
+        ("2021-04-01", 250),
+        ("2021-05-01", 260),
+        ("2021-06-01", 270),
+        ("2021-07-01", 280),
+        ("2021-08-01", 290),
+        ("2021-09-01", 300),
+        ("2021-10-01", 310),
+        ("2021-11-01", 320),
+        ("2021-12-01", 330),
+    ]
+
+    def test_head(self):
+        ts = TimeSeriesCore(self.data, frequency="M")
+        assert len(ts.head()) == 6
+        assert len(ts.head(3)) == 3
+        assert isinstance(ts.head(), TimeSeriesCore)
+        head_ts = ts.head(6)
+        assert head_ts.iloc[-1][1] == 270
+
+    def test_tail(self):
+        ts = TimeSeriesCore(self.data, frequency="M")
+        assert len(ts.tail()) == 6
+        assert len(ts.tail(8)) == 8
+        assert isinstance(ts.tail(), TimeSeriesCore)
+        tail_ts = ts.tail(6)
+        assert tail_ts.iloc[0][1] == 280
+
+    def test_head_tail(self):
+        ts = TimeSeriesCore(self.data, frequency="M")
+        head_tail_ts = ts.head(8).tail(2)
+        assert isinstance(head_tail_ts, TimeSeriesCore)
+        assert "2021-07-01" in head_tail_ts
+        assert head_tail_ts.iloc[1][1] == 290