From 44f766012e6ce5c7575162f21745f8abbf6f2b51 Mon Sep 17 00:00:00 2001 From: Jeremy Goh Date: Thu, 24 Feb 2022 21:30:19 +0800 Subject: [PATCH] Simplify encode_categorical logic slightly --- janitor/functions/encode_categorical.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/janitor/functions/encode_categorical.py b/janitor/functions/encode_categorical.py index 348a9e78a..3e1df0b9e 100644 --- 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)