Completed
Pull Request — master (#881)
by
unknown
01:24
created

zipline.pipeline.loaders.previous_earnings_date_frame()   B

Complexity

Conditions 3

Size

Total Lines 32

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 3
dl 0
loc 32
rs 8.8571
1
"""
2
Reference implementation for EarningsCalendar loaders.
3
"""
4
from numpy import full_like
5
import pandas as pd
6
from six import iteritems
7
8
from zipline.utils.memoize import lazyval
9
10
from .base import PipelineLoader
11
from .frame import DataFrameLoader
12
from ..data.earnings import EarningsCalendar
13
14
15
class EarningsCalendarLoader(PipelineLoader):
16
    """
17
    Reference loader for `zipline.pipeline.data.earnings.EarningsCalendar`.
18
19
    Does not currently support adjustments to the dates of known earnings.
20
21
    Parameters
22
    ----------
23
    all_dates : pd.DatetimeIndex
24
        Index of dates for which we can serve queries.
25
    announcement_dates : dict[int -> DatetimeIndex]
26
        Dict mapping column labels to an index of dates on which earnings were
27
        announced.
28
    """
29
    def __init__(self, all_dates, announcement_dates):
30
        self._all_dates = all_dates
31
        self._announcment_dates = announcement_dates
32
33
    def get_loader(self, column):
34
        """
35
        Dispatch to the loader for `column`.
36
        """
37
        if column is EarningsCalendar.next_announcement:
38
            return self.next_announcement_loader
39
        elif column is EarningsCalendar.previous_announcement:
40
            return self.previous_annoucement_loader
41
        else:
42
            raise ValueError("Don't know how to load column %s." % column)
43
44
    @lazyval
45
    def next_announcement_loader(self):
46
        return DataFrameLoader(
47
            EarningsCalendar.next_announcement,
48
            next_earnings_date_frame(
49
                self._all_dates,
50
                self._announcement_dates,
51
            ),
52
            adjustments=None,
53
        )
54
55
    @lazyval
56
    def previous_announcement_loader(self):
57
        return DataFrameLoader(
58
            EarningsCalendar.previous_announcement,
59
            previous_earnings_date_frame(
60
                self._all_dates,
61
                self._announcement_dates,
62
            ),
63
            adjustments=None,
64
        )
65
66
    def load_adjusted_array(self, columns, dates, assets, mask):
67
        return {
68
            column: self.get_loader(column).load_adjusted_array(
69
                [column], dates, assets, mask
70
            )
71
            for column in columns
72
        }
73
74
75
def next_earnings_date_frame(dates, announcement_dates):
76
    """
77
    Make a DataFrame representing simulated next earnings dates.
78
79
    Parameters
80
    ----------
81
    dates : DatetimeIndex.
82
        The index of the returned DataFrame.
83
    announcement_dates : dict[int -> DatetimeIndex]
84
        Dict mapping sids to an index of dates on which earnings were announced
85
        for that sid.
86
87
    Returns
88
    -------
89
    next_earnings: pd.DataFrame
90
        A DataFrame representing, for each (label, date) pair, the first entry
91
        in `earnings_calendars[label]` on or after `date`.  Entries falling
92
        after the last date in a calendar will have `NaT` as the result in the
93
        output.
94
95
    See Also
96
    --------
97
    next_earnings_date_frame
98
    """
99
    cols = {equity: full_like(dates, "NaT") for equity in announcement_dates}
100
    for equity, earnings_dates in iteritems(announcement_dates):
101
        next_dt_indices = earnings_dates.searchsorted(dates)
102
        mask = next_dt_indices < len(earnings_dates)
103
        cols[equity][mask] = earnings_dates[next_dt_indices[mask]]
104
105
    return pd.DataFrame(index=dates, data=cols)
106
107
108
def previous_earnings_date_frame(dates, announcement_dates):
109
    """
110
    Make a DataFrame representing simulated next earnings dates.
111
112
    Parameters
113
    ----------
114
    dates : DatetimeIndex.
115
        The index of the returned DataFrame.
116
    announcement_dates : dict[int -> DatetimeIndex]
117
        Dict mapping sids to an index of dates on which earnings were announced
118
        for that sid.
119
120
    Returns
121
    -------
122
    prev_earnings: pd.DataFrame
123
        A DataFrame representing, for (label, date) pair, the first entry in
124
        `announcement_dates[label]` strictly before `date`.  Entries falling
125
        before the first date in a calendar will have `NaT` as the result in
126
        the output.
127
128
    See Also
129
    --------
130
    next_earnings_date_frame
131
    """
132
    cols = {equity: full_like(dates, "NaT") for equity in announcement_dates}
133
    for equity, earnings_dates in iteritems(announcement_dates):
134
        # Subtract one to roll back to the index of the previous date.
135
        prev_dt_indices = earnings_dates.searchsorted(dates) - 1
136
        mask = prev_dt_indices > 0
137
        cols[equity][mask] = earnings_dates[prev_dt_indices[mask]]
138
139
    return pd.DataFrame(index=dates, data=cols)
140