Repository: quantopian/zipline
File: zipline/pipeline/loaders/earnings_estimates.py (Python)

Code before:
                    self.create_overwrite_for_quarter(
                        col_to_overwrites,
                        idx,
                        last_per_qtr,
                        qtrs_with_estimates,
                        requested_qtr_data,
                        sid,
                        sid_to_idx[sid],
                        columns,
                    )

        quarter_shifts.groupby(level=SID_FIELD_NAME).apply(collect_adjustments)
        return col_to_overwrites

    def create_overwrite_for_quarter(self,
                                     col_to_overwrites,
                                     next_qtr_start_idx,
                                     last_per_qtr,
                                     quarters_with_estimates_for_sid,
                                     requested_qtr_data,
                                     sid,
                                     sid_idx,
                                     columns):
        """
        Add entries to the dictionary of columns to adjustments for the given
        sid and the given quarter.

        Parameters
        ----------
        col_to_overwrites : dict [column_name -> list of ArrayAdjustment]
            A dictionary mapping column names to all overwrites for those
            columns.
        next_qtr_start_idx : int
            The index of the first day of the next quarter in the calendar
            dates.
        last_per_qtr : pd.DataFrame
            A DataFrame with a column MultiIndex of [self.estimates.columns,
            normalized_quarters, sid] that allows easily getting the timeline
            of estimates for a particular sid for a particular quarter; this
            is particularly useful for getting adjustments for 'next'
            estimates.
        quarters_with_estimates_for_sid : np.array
            An array of all quarters for which there are estimates for the
            given sid.
        sid : int
            The sid for which to create overwrites.
        sid_idx : int
            The index of the sid in `assets`.
        columns : list of BoundColumn
            The columns for which to create overwrites.
        """

Review comment (bug):
We should probably add back the comment about why we're checking for this and what it means if the `idx` fails these conditions.


Code after:
        """
        col_to_overwrites = defaultdict(dict)
        # We no longer need NORMALIZED_QUARTERS in the index, but we do need it
        # as a column to calculate adjustments.
        zero_qtr_data = zero_qtr_data.reset_index(level=NORMALIZED_QUARTERS)
        zero_qtr_data.sort_index(inplace=True)

        quarter_shifts = zero_qtr_data.loc[
            zero_qtr_data.index[
                zero_qtr_data.groupby(level=SID_FIELD_NAME)[
                    NORMALIZED_QUARTERS
                ].diff().nonzero()
            ]
        ]

        sid_to_idx = dict(zip(assets, range(len(assets))))

        def collect_adjustments(group):
            next_qtr_start_indices = dates.searchsorted(
                group[EVENT_DATE_FIELD_NAME].values,
                side=self.searchsorted_side,
            )
            sid = int(group.name)
            qtrs_with_estimates = group[NORMALIZED_QUARTERS].values
            for idx in next_qtr_start_indices:
                if 0 < idx < len(dates):
                    # Only add adjustments if the next quarter starts somewhere
                    # in our date index for this sid. Our 'next' quarter can
                    # never start at index 0; a starting index of 0 means that
                    # the next quarter's event date was NaT.
                    self.create_overwrite_for_quarter(
                        col_to_overwrites,
                        idx,
                        last_per_qtr,
                        qtrs_with_estimates,
                        requested_qtr_data,
                        sid,
                        sid_to_idx[sid],
                        columns,
                    )

        quarter_shifts.groupby(level=SID_FIELD_NAME).apply(collect_adjustments)
        return col_to_overwrites

    def create_overwrite_for_quarter(self,
                                     col_to_overwrites,
                                     next_qtr_start_idx,
                                     last_per_qtr,
                                     quarters_with_estimates_for_sid,
                                     requested_qtr_data,
                                     sid,