feat: define which columns to convert to numbers

Bisonai · Jul 20, 2024 · f30ed88 · f30ed88
1 parent cb3edb3
commit f30ed88
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/datamaxi/datamaxi/dex_trade.py b/datamaxi/datamaxi/dex_trade.py
@@ -95,7 +95,7 @@ def next_request():
             )
 
         if pandas:
-            df = convert_data_to_data_frame(res["data"])
+            df = convert_data_to_data_frame(res["data"], ["b", "bq", "qq", "p", "usd"])
             return df, next_request
         else:
             return res, next_request

diff --git a/datamaxi/datamaxi/utils.py b/datamaxi/datamaxi/utils.py
@@ -2,9 +2,21 @@
 import pandas as pd
 
 
-def convert_data_to_data_frame(data: List) -> pd.DataFrame:
+def convert_data_to_data_frame(
+    data: List,
+    columns_to_replace: List[str] = [],
+) -> pd.DataFrame:
     df = pd.DataFrame(data)
     df = df.set_index("d")
-    df.replace("NaN", pd.NA, inplace=True)
-    df = df.apply(pd.to_numeric, errors="coerce")
+
+    if len(columns_to_replace) == 0:
+        df.replace("NaN", pd.NA, inplace=True)
+        df = df.apply(pd.to_numeric, errors="coerce")
+        return df
+
+    df[columns_to_replace] = df[columns_to_replace].replace("NaN", pd.NA)
+    df[columns_to_replace] = df[columns_to_replace].apply(
+        pd.to_numeric, errors="coerce"
+    )
+
     return df