from sklearn.base import BaseEstimator, TransformerMixin
import pandas as pd

class ProductTypeCategorizer(BaseEstimator, TransformerMixin):

    def __init__(self):
        self.perishables = [
            "Dairy",
            "Meat",
            "Fruits and Vegetables",
            "Breakfast",
            "Breads",
            "Seafood",
        ]

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X = X.copy()
        X = pd.Series(X.squeeze())  # ensure 1D

        return pd.DataFrame(
            X.apply(lambda x: "Perishables" if x in self.perishables else "Non Perishables")
        )