using dtypekind enum values

Abhay-Lejith · Abhay-Lejith · commit 898b839c4bce · 2024-05-29T11:04:28.000+05:30
diff --git a/sktime/datatypes/_dtypekind.py b/sktime/datatypes/_dtypekind.py
@@ -0,0 +1,32 @@
+from enum import IntEnum
+
+
+class DtypeKind(IntEnum):
+    """
+    Integer enum for data types.
+
+    Attributes
+    ----------
+    INT : int
+        Matches to signed integer data type.
+    UINT : int
+        Matches to unsigned integer data type.
+    FLOAT : int
+        Matches to floating point data type.
+    BOOL : int
+        Matches to boolean data type.
+    STRING : int
+        Matches to string data type (UTF-8 encoded).
+    DATETIME : int
+        Matches to datetime data type.
+    CATEGORICAL : int
+        Matches to categorical data type.
+    """
+
+    INT = 0
+    UINT = 1
+    FLOAT = 2
+    BOOL = 20
+    STRING = 21  # UTF-8
+    DATETIME = 22
+    CATEGORICAL = 23
diff --git a/sktime/datatypes/_series/_check.py b/sktime/datatypes/_series/_check.py
@@ -44,6 +44,7 @@
 
 from sktime.datatypes._common import _req
 from sktime.datatypes._common import _ret as ret
+from sktime.datatypes._dtypekind import DtypeKind
 from sktime.utils.validation._dependencies import _check_soft_dependencies
 from sktime.utils.validation.series import is_in_valid_index_types
 
@@ -73,8 +74,23 @@ def check_pddataframe_series(obj, return_metadata=False, var_name="obj"):
         metadata["n_features"] = len(obj.columns)
     if _req("feature_names", return_metadata):
         metadata["feature_names"] = obj.columns.to_list()
+
     if _req("feature_kind", return_metadata):
-        metadata["feature_kind"] = obj.dtypes.to_list()
+        col_dtypes = obj.dtypes.to_list()
+        for i, dtype in enumerate(col_dtypes):
+            if dtype == float:
+                col_dtypes[i] = DtypeKind.FLOAT
+            elif dtype == int:
+                col_dtypes[i] = DtypeKind.INT
+            elif dtype == np.uint:
+                col_dtypes[i] = DtypeKind.UINT
+            elif dtype == object:
+                col_dtypes[i] = DtypeKind.CATEGORICAL
+            elif dtype == bool:
+                col_dtypes[i] = DtypeKind.BOOL
+            elif dtype == pd.DatetimeIndex:
+                col_dtypes[i] = DtypeKind.DATETIME
+        metadata["feature_kind"] = col_dtypes
 
     # check that columns are unique
     if not obj.columns.is_unique:
