Repository: solegalli/feature_engine

Issue:
modify _find_or_check_categorical_variables to support dtype='category'
**Is your feature request related to a problem? Please describe.**
Currently, _find_or_check_categorical_variables in variable_manipulation.py only checks for type object to determine categorical variables.

**Describe the solution you'd like**
The method should be enhanced to support dtype="category" to support cases where user has used category datatype to optimise memory footprint of Pandas dataframe. 
(https://pandas.pydata.org/pandas-docs/stable/user_guide/categorical.html)


Discussion:
solegalli: @bitsbuffer would you like to review this PR?
bitsbuffer: @solegalli sure, will do it asap.
solegalli: thank you! @bitsbuffer 

Patch:
diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py
index 3bcf0eb..3aa4f08 100644
--- a/feature_engine/encoding/base_encoder.py
+++ b/feature_engine/encoding/base_encoder.py
@@ -144,6 +144,14 @@ class BaseCategoricalTransformer(BaseEstimator, TransformerMixin):
         for feature in self.encoder_dict_.keys():
             X[feature] = X[feature].map(self.encoder_dict_[feature])
 
+            # if original variables are cast as categorical, they will remain
+            # categorical after the encoding, and this is probably not desired
+            if pd.api.types.is_categorical_dtype(X[feature]):
+                if all(isinstance(x, int) for x in X[feature]):
+                    X[feature] = X[feature].astype("int")
+                else:
+                    X[feature] = X[feature].astype("float")
+
         # check if NaN values were introduced by the encoding
         if X[self.encoder_dict_.keys()].isnull().sum().sum() > 0:
             warnings.warn(
diff --git a/feature_engine/imputation/categorical.py b/feature_engine/imputation/categorical.py
index 112fd0a..2002330 100644
--- a/feature_engine/imputation/categorical.py
+++ b/feature_engine/imputation/categorical.py
@@ -167,8 +167,24 @@ class CategoricalImputer(BaseImputer):
         return self
 
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:
-        # bring functionality from the BaseImputer
-        X = super().transform(X)
+
+        X = self._check_transform_input_and_state(X)
+
+        # replaces missing data with the learned parameters
+        if self.imputation_method == "frequent":
+            for variable in self.imputer_dict_:
+                X[variable].fillna(self.imputer_dict_[variable], inplace=True)
+
+        else:
+            for variable in self.imputer_dict_:
+                if pd.api.types.is_categorical_dtype(X[variable]):
+                    # if variable is of type category, we first need to add the new
+                    # category, and then fill in the nan
+                    X[variable].cat.add_categories(
+                        self.imputer_dict_[variable], inplace=True
+                    )
+
+                X[variable].fillna(self.imputer_dict_[variable], inplace=True)
 
         # add additional step to return variables cast as object
         if self.return_object:
diff --git a/feature_engine/variable_manipulation.py b/feature_engine/variable_manipulation.py
index 3ac4394..a8430cb 100644
--- a/feature_engine/variable_manipulation.py
+++ b/feature_engine/variable_manipulation.py
@@ -86,8 +86,8 @@ def _find_or_check_categorical_variables(
     X: pd.DataFrame, variables: Variables = None
 ) -> List[Union[str, int]]:
     """
-    Checks that variables provided by the user are of type object. If None, finds all
-    the categorical (object type) variables in the DataFrame.
+    Checks that variables provided by the user are of type object or categorical.
+    If None, finds all the categorical and object type variables in the DataFrame.
 
     Parameters
     ----------
@@ -103,25 +103,25 @@ def _find_or_check_categorical_variables(
 
     Returns
     -------
-    variables : List of numerical variables
+    variables : List of categorical variables
     """
 
     if isinstance(variables, (str, int)):
         variables = [variables]
 
     elif not variables:
-        variables = list(X.select_dtypes(include="O").columns)
+        variables = list(X.select_dtypes(include=["O", "category"]).columns)
         if len(variables) == 0:
             raise ValueError(
-                "No categorical variables in this dataframe. Please check variable "
-                "format with pandas dtypes"
+                "No categorical variables in this dataframe. Please check the "
+                "variables format with pandas dtypes"
             )
 
     else:
-        if any(X[variables].select_dtypes(exclude="O").columns):
+        if any(X[variables].select_dtypes(exclude=["O", "category"]).columns):
             raise TypeError(
                 "Some of the variables are not categorical. Please cast them as object "
-                "before calling this transformer"
+                "or category before calling this transformer"
             )
 
     return variables


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 49f2971..6b56dfe 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -79,6 +79,19 @@ def df_enc():
     return df
 
 
+@pytest.fixture(scope="module")
+def df_enc_category_dtypes():
+    df = {
+        "var_A": ["A"] * 6 + ["B"] * 10 + ["C"] * 4,
+        "var_B": ["A"] * 10 + ["B"] * 6 + ["C"] * 4,
+        "target": [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0],
+    }
+    df = pd.DataFrame(df)
+    df[["var_A", "var_B"]] = df[["var_A", "var_B"]].astype("category")
+
+    return df
+
+
 @pytest.fixture(scope="module")
 def df_enc_numeric():
     df = {
diff --git a/tests/test_encoding/test_count_frequency_encoder.py b/tests/test_encoding/test_count_frequency_encoder.py
index e405b75..e8ef049 100644
--- a/tests/test_encoding/test_count_frequency_encoder.py
+++ b/tests/test_encoding/test_count_frequency_encoder.py
@@ -201,3 +201,40 @@ def test_column_names_are_numbers(df_numeric_columns):
     assert encoder.n_features_in_ == 5
     # transform params
     pd.testing.assert_frame_equal(X, transf_df)
+
+
+def test_variables_cast_as_category(df_enc_category_dtypes):
+    encoder = CountFrequencyEncoder(encoding_method="count", variables=["var_A"])
+    X = encoder.fit_transform(df_enc_category_dtypes)
+
+    # expected result
+    transf_df = df_enc_category_dtypes.copy()
+    transf_df["var_A"] = [
+        6,
+        6,
+        6,
+        6,
+        6,
+        6,
+        10,
+        10,
+        10,
+        10,
+        10,
+        10,
+        10,
+        10,
+        10,
+        10,
+        4,
+        4,
+        4,
+        4,
+    ]
+    # transform params
+    pd.testing.assert_frame_equal(X, transf_df, check_dtype=False)
+    assert X["var_A"].dtypes == int
+
+    encoder = CountFrequencyEncoder(encoding_method="frequency", variables=["var_A"])
+    X = encoder.fit_transform(df_enc_category_dtypes)
+    assert X["var_A"].dtypes == float
diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py
index 0201b8a..f2286f5 100644
--- a/tests/test_encoding/test_decision_tree_encoder.py
+++ b/tests/test_encoding/test_decision_tree_encoder.py
@@ -96,3 +96,16 @@ def test_regression_ignore_format(df_enc_numeric):
     )  # Tree: var_A <= 1.5 -> 0.25 else 0.5
     transf_df["var_B"] = [0.044806] * 10 + [-0.079066] * 10
     pd.testing.assert_frame_equal(X.round(6), transf_df[["var_A", "var_B"]])
+
+
+def test_variables_cast_as_category(df_enc_category_dtypes):
+    df = df_enc_category_dtypes.copy()
+    encoder = DecisionTreeEncoder(regression=False)
+    encoder.fit(df[["var_A", "var_B"]], df["target"])
+    X = encoder.transform(df[["var_A", "var_B"]])
+
+    transf_df = df.copy()
+    transf_df["var_A"] = [0.25] * 16 + [0.5] * 4  # Tree: var_A <= 1.5 -> 0.25 else 0.5
+    transf_df["var_B"] = [0.2] * 10 + [0.4] * 10  # Tree: var_B <= 0.5 -> 0.2 else 0.4
+    pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]], check_dtype=False)
+    assert X["var_A"].dtypes == float
diff --git a/tests/test_encoding/test_mean_encoder.py b/tests/test_encoding/test_mean_encoder.py
index 0c17b56..706997a 100644
--- a/tests/test_encoding/test_mean_encoder.py
+++ b/tests/test_encoding/test_mean_encoder.py
@@ -258,3 +258,38 @@ def test_automatically_find_variables_ignore_format(df_enc_numeric):
     assert encoder.n_features_in_ == 2
     # test transform output
     pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]])
+
+
+def test_variables_cast_as_category(df_enc_category_dtypes):
+    df = df_enc_category_dtypes.copy()
+    encoder = MeanEncoder(variables=["var_A"])
+    encoder.fit(df[["var_A", "var_B"]], df["target"])
+    X = encoder.transform(df[["var_A", "var_B"]])
+
+    # expected output
+    transf_df = df.copy()
+    transf_df["var_A"] = [
+        0.3333333333333333,
+        0.3333333333333333,
+        0.3333333333333333,
+        0.3333333333333333,
+        0.3333333333333333,
+        0.3333333333333333,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.2,
+        0.5,
+        0.5,
+        0.5,
+        0.5,
+    ]
+
+    pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]], check_dtype=False)
+    assert X["var_A"].dtypes == float
diff --git a/tests/test_encoding/test_onehot_encoder.py b/tests/test_encoding/test_onehot_encoder.py
index f2f54c3..4f6b50f 100644
--- a/tests/test_encoding/test_onehot_encoder.py
+++ b/tests/test_encoding/test_onehot_encoder.py
@@ -168,3 +168,36 @@ def test_encode_numerical_variables(df_enc_numeric):
     assert encoder.n_features_in_ == 2
     # test transform output
     pd.testing.assert_frame_equal(X, transf)
+
+
+def test_variables_cast_as_category(df_enc_numeric):
+
+    encoder = OneHotEncoder(
+        top_categories=None,
+        variables=None,
+        drop_last=False,
+        ignore_format=True,
+    )
+
+    df = df_enc_numeric.copy()
+    df[["var_A", "var_B"]] = df[["var_A", "var_B"]].astype("category")
+
+    X = encoder.fit_transform(df[["var_A", "var_B"]])
+
+    # test fit attr
+    transf = {
+        "var_A_1": [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        "var_A_2": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+        "var_A_3": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
+        "var_B_1": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        "var_B_2": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+        "var_B_3": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1],
+    }
+
+    transf = pd.DataFrame(transf).astype("int32")
+    X = pd.DataFrame(X).astype("int32")
+
+    assert encoder.variables_ == ["var_A", "var_B"]
+    assert encoder.n_features_in_ == 2
+    # test transform output
+    pd.testing.assert_frame_equal(X, transf)
diff --git a/tests/test_encoding/test_ordinal_encoder.py b/tests/test_encoding/test_ordinal_encoder.py
index 339dead..c1ab2e7 100644
--- a/tests/test_encoding/test_ordinal_encoder.py
+++ b/tests/test_encoding/test_ordinal_encoder.py
@@ -141,3 +141,18 @@ def test_arbitrary_encoding_automatically_find_variables_ignore_format(df_enc_nu
     assert encoder.n_features_in_ == 2
     # test transform output
     pd.testing.assert_frame_equal(X, transf_df)
+
+
+def test_variables_cast_as_category(df_enc_category_dtypes):
+    df = df_enc_category_dtypes.copy()
+    encoder = OrdinalEncoder(encoding_method="ordered", variables=["var_A"])
+    encoder.fit(df[["var_A", "var_B"]], df["target"])
+    X = encoder.transform(df[["var_A", "var_B"]])
+
+    # expected output
+    transf_df = df.copy()
+    transf_df["var_A"] = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2]
+
+    # test transform output
+    pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]], check_dtype=False)
+    assert X["var_A"].dtypes == int
diff --git a/tests/test_encoding/test_rare_label_encoder.py b/tests/test_encoding/test_rare_label_encoder.py
index ecab2b1..96bae1d 100644
--- a/tests/test_encoding/test_rare_label_encoder.py
+++ b/tests/test_encoding/test_rare_label_encoder.py
@@ -171,3 +171,44 @@ def test_max_n_categories_with_numeric_var(df_enc_numeric):
     for i in range(len(df)):
         assert str(list(X["var_A"])[i]) == str(list(df["var_A"])[i])
         assert str(list(X["var_B"])[i]) == str(list(df["var_B"])[i])
+
+
+def test_variables_cast_as_category(df_enc_big):
+    # test case 1: defo params, automatically select variables
+    encoder = RareLabelEncoder(
+        tol=0.06, n_categories=5, variables=None, replace_with="Rare"
+    )
+
+    df_enc_big = df_enc_big.copy()
+    df_enc_big["var_B"] = df_enc_big["var_B"].astype("category")
+
+    X = encoder.fit_transform(df_enc_big)
+
+    # expected output
+    df = {
+        "var_A": ["A"] * 6
+        + ["B"] * 10
+        + ["C"] * 4
+        + ["D"] * 10
+        + ["Rare"] * 4
+        + ["G"] * 6,
+        "var_B": ["A"] * 10
+        + ["B"] * 6
+        + ["C"] * 4
+        + ["D"] * 10
+        + ["Rare"] * 4
+        + ["G"] * 6,
+        "var_C": ["A"] * 4
+        + ["B"] * 6
+        + ["C"] * 10
+        + ["D"] * 10
+        + ["Rare"] * 4
+        + ["G"] * 6,
+    }
+    df = pd.DataFrame(df)
+
+    # test fit attr
+    assert encoder.variables_ == ["var_A", "var_B", "var_C"]
+    assert encoder.n_features_in_ == 3
+    # test transform output
+    pd.testing.assert_frame_equal(X, df)
diff --git a/tests/test_encoding/test_woe_encoder.py b/tests/test_encoding/test_woe_encoder.py
index 826845e..535dfc2 100644
--- a/tests/test_encoding/test_woe_encoder.py
+++ b/tests/test_encoding/test_woe_encoder.py
@@ -235,3 +235,60 @@ def test_on_numerical_variables(df_enc_numeric):
     assert encoder.n_features_in_ == 2
     # transform params
     pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]])
+
+
+def test_variables_cast_as_category(df_enc_category_dtypes):
+    df = df_enc_category_dtypes.copy()
+    encoder = WoEEncoder(variables=None)
+    encoder.fit(df[["var_A", "var_B"]], df["target"])
+    X = encoder.transform(df[["var_A", "var_B"]])
+
+    # transformed dataframe
+    transf_df = df.copy()
+    transf_df["var_A"] = [
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        0.8472978603872037,
+        0.8472978603872037,
+        0.8472978603872037,
+        0.8472978603872037,
+    ]
+    transf_df["var_B"] = [
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        -0.5389965007326869,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.15415067982725836,
+        0.8472978603872037,
+        0.8472978603872037,
+        0.8472978603872037,
+        0.8472978603872037,
+    ]
+
+    pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]], check_dtype=False)
+    assert X["var_A"].dtypes == float
diff --git a/tests/test_imputation/test_categorical_imputer.py b/tests/test_imputation/test_categorical_imputer.py
index 84a5f08..774b28f 100644
--- a/tests/test_imputation/test_categorical_imputer.py
+++ b/tests/test_imputation/test_categorical_imputer.py
@@ -109,6 +109,7 @@ def test_mode_imputation_with_multiple_variables(df_na):
 
 def test_imputation_of_numerical_vars_cast_as_object_and_returned_as_numerical(df_na):
     # test case: imputing of numerical variables cast as object + return numeric
+    df_na = df_na.copy()
     df_na["Marks"] = df_na["Marks"].astype("O")
     imputer = CategoricalImputer(
         imputation_method="frequent", variables=["City", "Studies", "Marks"]
@@ -133,6 +134,7 @@ def test_imputation_of_numerical_vars_cast_as_object_and_returned_as_numerical(d
 def test_imputation_of_numerical_vars_cast_as_object_and_returned_as_object(df_na):
     # test case 6: imputing of numerical variables cast as object + return as object
     # after imputation
+    df_na = df_na.copy()
     df_na["Marks"] = df_na["Marks"].astype("O")
     imputer = CategoricalImputer(
         imputation_method="frequent",
@@ -215,3 +217,66 @@ def test_impute_numerical_variables_with_mode(df_na):
 
     # test transform output
     pd.testing.assert_frame_equal(X_transformed, X_reference)
+
+
+def test_variables_cast_as_category_missing(df_na):
+    # string missing
+    df_na = df_na.copy()
+    df_na["City"] = df_na["City"].astype("category")
+
+    imputer = CategoricalImputer(imputation_method="missing", variables=None)
+    X_transformed = imputer.fit_transform(df_na)
+
+    # set up expected output
+    X_reference = df_na.copy()
+    X_reference["Name"] = X_reference["Name"].fillna("Missing")
+    X_reference["Studies"] = X_reference["Studies"].fillna("Missing")
+
+    X_reference["City"].cat.add_categories("Missing", inplace=True)
+    X_reference["City"] = X_reference["City"].fillna("Missing")
+
+    # test fit attributes
+    assert imputer.variables_ == ["Name", "City", "Studies"]
+    assert imputer.imputer_dict_ == {
+        "Name": "Missing",
+        "City": "Missing",
+        "Studies": "Missing",
+    }
+
+    # test transform output
+    # selected columns should have no NA
+    # non selected columns should still have NA
+    assert X_transformed[["Name", "City", "Studies"]].isnull().sum().sum() == 0
+    assert X_transformed[["Age", "Marks"]].isnull().sum().sum() > 0
+    pd.testing.assert_frame_equal(X_transformed, X_reference)
+
+
+def test_variables_cast_as_category_frequent(df_na):
+
+    df_na = df_na.copy()
+    df_na["City"] = df_na["City"].astype("category")
+
+    # this variable does not have a mode, so drop
+    df_na.drop(labels=["Name"], axis=1, inplace=True)
+
+    imputer = CategoricalImputer(imputation_method="frequent", variables=None)
+    X_transformed = imputer.fit_transform(df_na)
+
+    # set up expected output
+    X_reference = df_na.copy()
+    X_reference["Studies"] = X_reference["Studies"].fillna("Bachelor")
+    X_reference["City"] = X_reference["City"].fillna("London")
+
+    # test fit attributes
+    assert imputer.variables_ == ["City", "Studies"]
+    assert imputer.imputer_dict_ == {
+        "City": "London",
+        "Studies": "Bachelor",
+    }
+
+    # test transform output
+    # selected columns should have no NA
+    # non selected columns should still have NA
+    assert X_transformed[["City", "Studies"]].isnull().sum().sum() == 0
+    assert X_transformed[["Age", "Marks"]].isnull().sum().sum() > 0
+    pd.testing.assert_frame_equal(X_transformed, X_reference)
diff --git a/tests/test_imputation/test_random_sample_imputer.py b/tests/test_imputation/test_random_sample_imputer.py
index 1ad29f3..52c62b3 100644
--- a/tests/test_imputation/test_random_sample_imputer.py
+++ b/tests/test_imputation/test_random_sample_imputer.py
@@ -70,11 +70,14 @@ def test_general_seed_plus_automatically_select_variables(df_na):
 def test_seed_per_observation_and_multiple_variables_in_random_state(df_na):
     # test case 2: imputer seed per observation using multiple variables to determine
     # the random_state
-    # Note the variables used as seed should not have missing data
+    # Note the variables used as seed should not have missing data, this I fill
+    df_na = df_na.copy()
+    df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1)
+
     imputer = RandomSampleImputer(
         variables=["City", "Studies"], random_state=["Marks", "Age"], seed="observation"
     )
-    df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1)
+
     X_transformed = imputer.fit_transform(df_na)
 
     # expected output
@@ -120,13 +123,17 @@ def test_seed_per_observation_and_multiple_variables_in_random_state(df_na):
 
 def test_seed_per_observation_plus_product_of_seeding_variables(df_na):
     # test case 3: observation seed, 2 variables as seed, product of seed variables
+    # need to fill variables used as seed
+    df_na = df_na.copy()
+    df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1)
+
     imputer = RandomSampleImputer(
         variables=["City", "Studies"],
         random_state=["Marks", "Age"],
         seed="observation",
         seeding_method="multiply",
     )
-    df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1)
+
     X_transformed = imputer.fit_transform(df_na)
 
     # expected output
@@ -176,10 +183,13 @@ def test_seed_per_observation_plus_product_of_seeding_variables(df_na):
 def test_seed_per_observation_with_only_1_variable_as_seed(df_na):
     # test case 4: observation seed, only variable indicated as seed, method: addition
     # Note the variable used as seed should not have missing data
+    df_na = df_na.copy()
+    df_na["Age"] = df_na["Age"].fillna(1)
+
     imputer = RandomSampleImputer(
         variables=["City", "Studies"], random_state="Age", seed="observation"
     )
-    df_na["Age"] = df_na["Age"].fillna(1)
+
     X_transformed = imputer.fit_transform(df_na)
 
     # expected output
@@ -248,3 +258,52 @@ def test_error_if_random_state_is_string(df_na):
     with pytest.raises(ValueError):
         imputer = RandomSampleImputer(seed="observation", random_state="arbitrary")
         imputer.fit(df_na)
+
+
+def test_variables_cast_as_category(df_na):
+
+    df_na = df_na.copy()
+    df_na["City"] = df_na["City"].astype("category")
+
+    # set up transformer
+    imputer = RandomSampleImputer(variables=None, random_state=5, seed="general")
+    X_transformed = imputer.fit_transform(df_na)
+
+    # expected output:
+    # fillna based on seed used (found experimenting on Jupyter notebook)
+    ref = {
+        "Name": ["tom", "nick", "krish", "peter", "peter", "sam", "fred", "sam"],
+        "City": [
+            "London",
+            "Manchester",
+            "London",
+            "Manchester",
+            "London",
+            "London",
+            "Bristol",
+            "Manchester",
+        ],
+        "Studies": [
+            "Bachelor",
+            "Bachelor",
+            "PhD",
+            "Masters",
+            "Bachelor",
+            "PhD",
+            "None",
+            "Masters",
+        ],
+        "Age": [20, 21, 19, 23, 23, 40, 41, 37],
+        "Marks": [0.9, 0.8, 0.7, 0.3, 0.3, 0.6, 0.8, 0.6],
+        "dob": pd.date_range("2020-02-24", periods=8, freq="T"),
+    }
+    ref = pd.DataFrame(ref)
+    ref["City"] = ref["City"].astype("category")
+
+    # test fit attr
+    assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks", "dob"]
+    assert imputer.n_features_in_ == 6
+    pd.testing.assert_frame_equal(imputer.X_, df_na)
+
+    # test transform output
+    pd.testing.assert_frame_equal(X_transformed, ref, check_dtype=False)
diff --git a/tests/test_variable_manipulation.py b/tests/test_variable_manipulation.py
index 21e5479..647eb70 100644
--- a/tests/test_variable_manipulation.py
+++ b/tests/test_variable_manipulation.py
@@ -57,10 +57,9 @@ def test_find_or_check_numerical_variables(df_vartypes, df_numeric_columns):
 def test_find_or_check_categorical_variables(df_vartypes, df_numeric_columns):
     vars_cat = ["Name", "City"]
     vars_mix = ["Age", "Marks", "Name"]
-    vars_none = None
 
     assert _find_or_check_categorical_variables(df_vartypes, vars_cat) == vars_cat
-    assert _find_or_check_categorical_variables(df_vartypes, vars_none) == vars_cat
+    assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat
 
     with pytest.raises(TypeError):
         assert _find_or_check_categorical_variables(df_vartypes, vars_mix)
@@ -71,6 +70,13 @@ def test_find_or_check_categorical_variables(df_vartypes, df_numeric_columns):
     assert _find_or_check_categorical_variables(df_numeric_columns, [0, 1]) == [0, 1]
     assert _find_or_check_categorical_variables(df_numeric_columns, 1) == [1]
 
+    df_vartypes["Age"] = df_vartypes["Age"].astype("category")
+    assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat + ["Age"]
+    assert _find_or_check_categorical_variables(df_vartypes, ["Name", "Age"]) == [
+        "Name",
+        "Age",
+    ]
+
 
 def test_find_all_variables(df_vartypes):
     all_vars = ["Name", "City", "Age", "Marks", "dob"]