++

simple72/Tranformer/ace_lib.py

@@ -1,4 +1,3 @@
-import getpass
 import json
 import logging
 import os
@@ -6,23 +5,85 @@ import threading
 import time
 from functools import partial
 from multiprocessing.pool import ThreadPool
-from pathlib import Path
 from typing import Literal, Optional, Union
 from urllib.parse import urljoin
 
 import pandas as pd
 import requests
 import tqdm
-from helpful_functions import (
-    expand_dict_columns,
-    save_pnl,
-    save_simulation_result,
-    save_yearly_stats,
-)
-
 DEV = False
 
 
+def save_simulation_result(result: dict) -> None:
+    """
+    Save the simulation result to a JSON file in the 'simulation_results' folder.
+
+    Args:
+        result (dict): A dictionary containing the simulation result for an alpha.
+    """
+    alpha_id = result["id"]
+    region = result["settings"]["region"]
+    folder_path = "simulation_results/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}")
+
+    os.makedirs(folder_path, exist_ok=True)
+
+    with open(file_path, "w", encoding="utf-8") as file:
+        json.dump(result, file)
+
+
+def save_pnl(pnl_df: pd.DataFrame, alpha_id: str, region: str) -> None:
+    """
+    Save the PnL data for an alpha to a CSV file in the 'alphas_pnl' folder.
+
+    Args:
+        pnl_df (pandas.DataFrame): The DataFrame containing PnL data.
+        alpha_id (str): The ID of the alpha.
+        region (str): The region for which the PnL data was generated.
+    """
+    folder_path = "alphas_pnl/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv")
+    os.makedirs(folder_path, exist_ok=True)
+
+    pnl_df.to_csv(file_path)
+
+
+def save_yearly_stats(yearly_stats: pd.DataFrame, alpha_id: str, region: str):
+    """
+    Save the yearly statistics for an alpha to a CSV file in the 'yearly_stats' folder.
+
+    Args:
+        yearly_stats (pandas.DataFrame): The DataFrame containing yearly statistics.
+        alpha_id (str): The ID of the alpha.
+        region (str): The region for which the statistics were generated.
+    """
+    folder_path = "yearly_stats/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv")
+    os.makedirs(folder_path, exist_ok=True)
+
+    yearly_stats.to_csv(file_path, index=False)
+
+
+def expand_dict_columns(data: pd.DataFrame) -> pd.DataFrame:
+    """
+    Expand dictionary columns in a DataFrame into separate columns.
+
+    Args:
+        data (pandas.DataFrame): The input DataFrame with dictionary columns.
+
+    Returns:
+        pandas.DataFrame: A new DataFrame with expanded columns.
+    """
+    dict_columns = list(filter(lambda x: isinstance(data[x].iloc[0], dict), data.columns))
+    new_columns = pd.concat(
+        [data[col].apply(pd.Series).rename(columns=lambda x: f"{col}_{x}") for col in dict_columns],
+        axis=1,
+    )
+
+    data = pd.concat([data, new_columns], axis=1)
+    return data
+
+
 class SingleSession(requests.Session):
     _instance = None
     _lock = threading.Lock()
@@ -132,7 +193,7 @@ def start_session() -> SingleSession:
     logger.debug(f"New session created (ID: {id(s)}) with authentication response: {r.status_code}, {r.json()}")
     if r.status_code == requests.status_codes.codes.unauthorized:
         if r.headers["WWW-Authenticate"] == "persona":
-            print(
+            logger.info(
                 "Complete biometrics authentication and press any key to continue: \n"
                 + urljoin(r.url, r.headers["Location"])
                 + "\n"
@@ -1192,6 +1253,7 @@ def get_datafields(
     dataset_id: str = "",
     data_type: str = "MATRIX",
     search: str = "",
+    category: str = "",
 ) -> pd.DataFrame:
     """
     Retrieve available datafields based on specified parameters.
@@ -1219,7 +1281,14 @@ def get_datafields(
             + f"&region={region}&delay={str(delay)}&universe={universe}{type_param}&dataset.id={dataset_id}&limit=50"
             + "&offset={x}"
         )
-        count = s.get(url_template.format(x=0)).json()["count"]
+        try:
+            response = s.get(url_template.format(x=0))
+            response_json = response.json()
+            count = response_json.get("count", 0)
+        except Exception as e:
+            logger.warning(f"获取数据字段数量失败: {e}")
+            count = 0
+        
         if count == 0:
             logger.warning(
                 f"No fields found: region={region}, delay={str(delay)}, universe={universe}, "
@@ -1238,13 +1307,15 @@ def get_datafields(
         )
         count = 100
 
+    logger.info(url_template.format(x=0))
+    
     max_try = 5
     datafields_list = []
     for x in range(0, count, 50):
         for _ in range(max_try):
             datafields = s.get(url_template.format(x=x))
             while datafields.status_code == 429:
-                print("status_code 429, sleep 3 seconds")
+                logger.info("status_code 429, sleep 3 seconds")
                 time.sleep(3)
                 datafields = s.get(url_template.format(x=x))
             if "results" in datafields.json():
@@ -1510,5 +1581,4 @@ def main():
     simulate_alpha_list(s, alpha_list)
 
 
-if __name__ == "__main__":
-    main()
+