# Copyright (C) 2019 Alteryx, Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
"""Definition of the DropNaRowsTransformer."""

import pandas as pd

import sklearn.base


class DropNaRowsTransformer(sklearn.base.TransformerMixin):
    """Custom transformer for dropping rows when a given column has a null value."""

    def __init__(self, colname, inplace=False):
        self.colname = colname
        self.inplace = bool(inplace)

    def fit(self, x, y=None, **kwargs):
        """Fit the transformer."""

        return self

    def transform(self, df):
        """Transform to remove NA rows."""
        inplace = self.inplace

        assert isinstance(df, pd.DataFrame)
        assert self.colname in list(df)

        result = df.dropna(subset=[self.colname], inplace=inplace)
        if not inplace:
            df = result

        return df
