Przeglądaj źródła

Added Flattening support for pd.Series

ogert 5 lat temu
rodzic
commit
6ca8438c32
2 zmienionych plików z 7 dodań i 3 usunięć
  1. 5 1
      cdplib/FlattenData.py
  2. 2 2
      cdplib/unit_tests/TestFlattenData.py

+ 5 - 1
cdplib/FlattenData.py

@@ -25,13 +25,17 @@ class FlattenData():
         :parm data: data given in either dictionary, list or dataframe format.
         '''
 
-        assert(isinstance(data, (list, dict, pd.DataFrame))),\
+        assert(isinstance(data, (list, dict, pd.DataFrame, pd.Series))),\
             "Parameter 'data' either be of List, Dictionary or DataFrame type"
         in_length=0
         start = time.time()
         if type(data) is pd.DataFrame:
             in_length = len(data.columns)
             return_data = self.flatten_dataframe(data)
+        elif type(data) is pd.Series:
+            data = pd.DataFrame(data)
+            in_length = len(data.columns)
+            return_data = self.flatten_dataframe(data)
         elif type(data) is dict:
             in_length = len(data)
             return_data = self.flatten_dict(data)

+ 2 - 2
cdplib/unit_tests/TestFlattenData.py

@@ -35,8 +35,8 @@ class TestMongodbHandler(unittest.TestCase):
 
         flattened_dict = self.flattener.flatten(nested_data_dict)
         flattened_list = self.flattener.flatten(nested_data_list)
-        flattened_df = self.flattener.flatten(nested_data_df) 
-    
+        flattened_df = self.flattener.flatten(nested_data_df['two_levels']) 
+
         self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_dict.loc['two_levels_one_level', 0])
         self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_list.loc['0_two_levels_one_level', 0])
         self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_df.loc[0 , 'two_levels_one_level'])