Simplify encode_categorical logic slightly

pyjanitor-devs · Mar 1, 2022 · 44f7660 · 44f7660
1 parent a784a77
commit 44f7660
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/janitor/functions/encode_categorical.py b/janitor/functions/encode_categorical.py
@@ -112,13 +112,11 @@ def encode_categorical(
     # or user supplies specific categories to create the categorical
     if column_names is not None:
         check("column_names", column_names, [list, tuple, Hashable])
-        if isinstance(column_names, (list, tuple)):
-            check_column(df, column_names)
-            dtypes = {col: "category" for col in column_names}
-            return df.astype(dtypes)
         if isinstance(column_names, Hashable):
-            check_column(df, [column_names])
-            return df.astype({column_names: "category"})
+            column_names = [column_names]
+        check_column(df, column_names)
+        dtypes = {col: "category" for col in column_names}
+        return df.astype(dtypes)
 
     return _computations_as_categorical(df, **kwargs)