PyFacts/tests/test_core.py

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")
        assert D.days == 1
        assert D.symbol == "D"
        assert D.name == "daily"
        assert D.value == 1
        assert D.freq_type == "days"


def create_test_data(
    frequency: str,
    eomonth: bool,
    n: int,
    gaps: float,
    month_position: Literal["start", "middle", "end"],
    date_as_str: bool,
    as_outer_type: Literal["dict", "list"] = "list",
    as_inner_type: Literal["dict", "list", "tuple"] = "tuple",
) -> Sequence[tuple]:
    start_dates = {
        "start": datetime.datetime(2016, 1, 1),
        "middle": datetime.datetime(2016, 1, 15),
        "end": datetime.datetime(2016, 1, 31),
    }
    end_date = datetime.datetime(2021, 12, 31)
    dates = create_date_series(start_dates[month_position], end_date, frequency=frequency, eomonth=eomonth)
    dates = dates[:n]
    if gaps:
        num_gaps = int(len(dates) * gaps)
        to_remove = random.sample(dates, num_gaps)
        for i in to_remove:
            dates.remove(i)
    if date_as_str:
        dates = [i.strftime("%Y-%m-%d") for i in dates]

    values = [random.randint(8000, 90000) / 100 for _ in dates]

    data = list(zip(dates, values))
    if as_outer_type == "list":
        if as_inner_type == "list":
            data = [list(i) for i in data]
        elif as_inner_type == "dict[1]":
            data = [dict((i,)) for i in data]
        elif as_inner_type == "dict[2]":
            data = [dict(date=i, value=j) for i, j in data]
    elif as_outer_type == "dict":
        data = dict(data)

    return data


class TestAllFrequencies:
    def test_attributes(self):
        assert hasattr(AllFrequencies, "D")
        assert hasattr(AllFrequencies, "M")
        assert hasattr(AllFrequencies, "Q")

    def test_days(self):
        assert AllFrequencies.D.days == 1
        assert AllFrequencies.M.days == 30
        assert AllFrequencies.Q.days == 91

    def test_symbol(self):
        assert AllFrequencies.H.symbol == "H"
        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.W.freq_type == "days"


class TestSeries:
    def test_creation(self):
        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")
        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)]

    def test_repr_str(self):
        ts = TimeSeriesCore(self.data, frequency="M")
        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 repr(ts)

    def test_creation(self):
        ts = TimeSeriesCore(self.data, frequency="M")
        assert isinstance(ts, TimeSeriesCore)
        assert isinstance(ts, Mapping)

    def test_getitem(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["2021-01-01"][1] == 220
        assert len(ts[ts.dates > "2021-01-01"]) == 2
        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")
        assert datetime.datetime(2021, 1, 1) in ts
        assert "2021-01-01" in ts
        assert "2021-01-14" not in ts

    def test_items(self):
        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