def apply()

in misc/CCSynth/CC/DataInsights/src/prose/datainsights/_assertion/_assertion_helper.py [0:0]

• 28 lines of code
• 10 McCabe index (conditional complexity)

    def apply(self, df, drop_column=True):
        """ Returns the transformed dataframe after the constraint is applied
        :param df: source dataframe
        :param drop_column: boolean flag. denotes whether to drop the column, on which constraint is applied, or not
        :return: df with the constraint applied
        """
        new_df = df
        if self.relational_op == _RelationalOperators.EQUAL:
            new_df = df[df[self.column_name] == self.column_value]
        elif self.relational_op == _RelationalOperators.LESS_THAN:
            new_df = df[df[self.column_name] < self.column_value]
        elif self.relational_op == _RelationalOperators.LESS_THAN_EQUAL_TO:
            new_df = df[df[self.column_name] <= self.column_value]
        elif self.relational_op == _RelationalOperators.GREATER_THAN:
            new_df = df[df[self.column_name] > self.column_value]
        elif self.relational_op == _RelationalOperators.GREATER_THAN_EQUAL_TO:
            new_df = df[df[self.column_name] >= self.column_value]
        elif self.relational_op == _RelationalOperators.NOT_EQUAL:
            new_df = df[df[self.column_name] != self.column_value]
        elif self.relational_op == _RelationalOperators.INCLUSIVE_RANGE:
            new_df = df[
                (df[self.column_name] >= self.column_value[0])
                & (df[self.column_name] <= self.column_value[1])
            ]
        elif self.relational_op == _RelationalOperators.EXCLUSIVE_RANGE:
            new_df = df[
                (df[self.column_name] > self.column_value[0])
                & (df[self.column_name] < self.column_value[1])
            ]
        if drop_column:
            return new_df.loc[:, new_df.columns != self.column_name]
        else:
            return new_df