++
This commit is contained in:
File diff suppressed because it is too large
Load Diff
+157
@@ -0,0 +1,157 @@
|
||||
import getpass
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# Ensure we can import ace_lib from the project root
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
ROOT_DIR = os.path.dirname(SCRIPT_DIR)
|
||||
if ROOT_DIR not in sys.path:
|
||||
sys.path.append(ROOT_DIR)
|
||||
|
||||
import ace_lib # noqa: E402
|
||||
|
||||
|
||||
def prompt_credentials() -> tuple[str, str]:
|
||||
"""Prompt user for platform credentials."""
|
||||
email = input("Enter BRAIN Email: ").strip()
|
||||
while not email:
|
||||
email = input("Email is required. Enter BRAIN Email: ").strip()
|
||||
|
||||
password = getpass.getpass("Enter BRAIN Password: ").strip()
|
||||
while not password:
|
||||
password = getpass.getpass("Password is required. Enter BRAIN Password: ").strip()
|
||||
|
||||
return email, password
|
||||
|
||||
|
||||
def fetch_all_combinations(session: ace_lib.SingleSession) -> pd.DataFrame:
|
||||
"""Return all valid instrument/region/delay/universe combos from platform settings."""
|
||||
options_df = ace_lib.get_instrument_type_region_delay(session)
|
||||
if options_df is None or options_df.empty:
|
||||
raise RuntimeError("No simulation options fetched; cannot enumerate datasets.")
|
||||
return options_df
|
||||
|
||||
|
||||
def fetch_datasets_for_combo(
|
||||
session: ace_lib.SingleSession,
|
||||
instrument_type: str,
|
||||
region: str,
|
||||
delay: int,
|
||||
universe: str,
|
||||
) -> pd.DataFrame:
|
||||
"""Fetch datasets for one combination (theme ALL to include both theme true/false)."""
|
||||
df = ace_lib.get_datasets(
|
||||
session,
|
||||
instrument_type=instrument_type,
|
||||
region=region,
|
||||
delay=delay,
|
||||
universe=universe,
|
||||
theme="ALL",
|
||||
)
|
||||
if df is None:
|
||||
return pd.DataFrame()
|
||||
|
||||
df = df.copy()
|
||||
df["param_instrument_type"] = instrument_type
|
||||
df["param_region"] = region
|
||||
df["param_delay"] = delay
|
||||
df["param_universe"] = universe
|
||||
df["combo_key"] = df.apply(
|
||||
lambda row: f"{instrument_type}-{region}-D{delay}-{universe}",
|
||||
axis=1,
|
||||
)
|
||||
return df
|
||||
|
||||
|
||||
def merge_and_deduplicate(datasets: List[pd.DataFrame]) -> pd.DataFrame:
|
||||
"""Merge fetched datasets and deduplicate by dataset id, keeping all combo metadata."""
|
||||
combined = pd.concat([df for df in datasets if not df.empty], ignore_index=True)
|
||||
if combined.empty:
|
||||
return combined
|
||||
|
||||
# Aggregate availability combos per dataset id
|
||||
availability = (
|
||||
combined.groupby("id")["combo_key"]
|
||||
.agg(lambda x: " | ".join(sorted(set(x))))
|
||||
.rename("available_in")
|
||||
.reset_index()
|
||||
)
|
||||
|
||||
# Drop duplicate rows by dataset id, keep first occurrence of other columns
|
||||
unique_df = combined.drop_duplicates(subset=["id"]).copy()
|
||||
unique_df = unique_df.merge(availability, on="id", how="left")
|
||||
|
||||
# Sort for readability
|
||||
sort_cols = [col for col in ["category", "subcategory", "id"] if col in unique_df.columns]
|
||||
if sort_cols:
|
||||
# Ensure sort keys are hashable/strings to avoid unhashable dict errors
|
||||
for col in sort_cols:
|
||||
unique_df[col] = unique_df[col].apply(
|
||||
lambda v: v
|
||||
if pd.isna(v) or isinstance(v, (int, float, str, bool))
|
||||
else json.dumps(v, ensure_ascii=False, sort_keys=True)
|
||||
)
|
||||
unique_df = unique_df.sort_values(sort_cols).reset_index(drop=True)
|
||||
|
||||
return unique_df
|
||||
|
||||
|
||||
def main():
|
||||
print("=== Fetch All BRAIN Datasets (all regions/universes/delays) ===")
|
||||
|
||||
email, password = prompt_credentials()
|
||||
|
||||
# Monkey-patch ace_lib credential retrieval so start_session uses provided credentials
|
||||
ace_lib.get_credentials = lambda: (email, password)
|
||||
|
||||
print("Logging in...")
|
||||
try:
|
||||
session = ace_lib.start_session()
|
||||
print("Login successful.")
|
||||
except Exception as exc:
|
||||
print(f"Login failed: {exc}")
|
||||
return
|
||||
|
||||
print("Fetching valid instrument/region/delay/universe combinations from platform settings...")
|
||||
try:
|
||||
options_df = fetch_all_combinations(session)
|
||||
except Exception as exc:
|
||||
print(f"Failed to fetch simulation options: {exc}")
|
||||
return
|
||||
|
||||
all_datasets: List[pd.DataFrame] = []
|
||||
total_combos = 0
|
||||
|
||||
for _, row in options_df.iterrows():
|
||||
instrument_type = row.get("InstrumentType")
|
||||
region = row.get("Region")
|
||||
delay = row.get("Delay")
|
||||
universes = row.get("Universe") or []
|
||||
|
||||
for universe in universes:
|
||||
total_combos += 1
|
||||
print(f"[{total_combos}] Fetching datasets for {instrument_type} / {region} / D{delay} / {universe}...")
|
||||
try:
|
||||
df = fetch_datasets_for_combo(session, instrument_type, region, delay, universe)
|
||||
print(f" -> Retrieved {len(df)} rows")
|
||||
all_datasets.append(df)
|
||||
except Exception as exc:
|
||||
print(f" -> Failed for {instrument_type}-{region}-D{delay}-{universe}: {exc}")
|
||||
|
||||
result_df = merge_and_deduplicate(all_datasets)
|
||||
|
||||
if result_df.empty:
|
||||
print("No datasets fetched; nothing to save.")
|
||||
return
|
||||
|
||||
output_path = os.path.join(SCRIPT_DIR, "all_datasets_full.csv")
|
||||
result_df.to_csv(output_path, index=False)
|
||||
print(f"Saved {len(result_df)} unique datasets to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+132
@@ -0,0 +1,132 @@
|
||||
"""Minimal helper script to log in, list tutorials, and dump each tutorial page.
|
||||
|
||||
- Prompts for email/password at runtime (keeps credentials out of source).
|
||||
- Uses only the requests standard stack; no project-internal imports.
|
||||
- Saves each page JSON to tutorial_dump/<index>_<page_id>.json for inspection.
|
||||
|
||||
Run with: python doc_fetch_sample.py
|
||||
"""
|
||||
import base64
|
||||
import getpass
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import requests
|
||||
|
||||
BASE_URL = "https://api.worldquantbrain.com"
|
||||
|
||||
|
||||
def _basic_auth_header(email: str, password: str) -> Dict[str, str]:
|
||||
token = base64.b64encode(f"{email}:{password}".encode()).decode()
|
||||
return {"Authorization": f"Basic {token}"}
|
||||
|
||||
|
||||
def authenticate(email: str, password: str) -> requests.Session:
|
||||
"""Authenticate and return a session carrying the JWT cookie."""
|
||||
session = requests.Session()
|
||||
resp = session.post(f"{BASE_URL}/authentication", headers=_basic_auth_header(email, password), timeout=30)
|
||||
if resp.status_code != 201:
|
||||
raise RuntimeError(f"Authentication failed (status {resp.status_code}): {resp.text}")
|
||||
return session
|
||||
|
||||
|
||||
def fetch_tutorials(session: requests.Session) -> List[Dict[str, Any]]:
|
||||
"""Fetch tutorials list; handle a few common response shapes."""
|
||||
resp = session.get(f"{BASE_URL}/tutorials", timeout=30)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
if isinstance(data, list):
|
||||
return data
|
||||
if isinstance(data, dict):
|
||||
for key in ("items", "results", "data", "tutorials"):
|
||||
maybe = data.get(key)
|
||||
if isinstance(maybe, list):
|
||||
return maybe
|
||||
return []
|
||||
|
||||
|
||||
def fetch_tutorial_pages(session: requests.Session, tutorial_id: str) -> List[Dict[str, Any]]:
|
||||
"""Fetch pages for a tutorial when the list entry only gives a tutorial id/slug."""
|
||||
resp = session.get(f"{BASE_URL}/tutorials/{tutorial_id}/pages", timeout=30)
|
||||
if resp.status_code == 404:
|
||||
return [] # graceful fallback
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
if isinstance(data, list):
|
||||
return data
|
||||
if isinstance(data, dict):
|
||||
for key in ("items", "results", "pages", "data"):
|
||||
maybe = data.get(key)
|
||||
if isinstance(maybe, list):
|
||||
return maybe
|
||||
return []
|
||||
|
||||
|
||||
def _extract_page_id(entry: Dict[str, Any]) -> Optional[str]:
|
||||
for key in ("page_id", "pageId", "id", "pageID", "slug", "code"):
|
||||
if key in entry and entry[key] is not None:
|
||||
return str(entry[key])
|
||||
return None
|
||||
|
||||
|
||||
def fetch_page(session: requests.Session, page_id: str) -> Dict[str, Any]:
|
||||
resp = session.get(f"{BASE_URL}/tutorial-pages/{page_id}", timeout=30)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def dump_pages(session: requests.Session, tutorials: List[Dict[str, Any]], out_dir: str = "tutorial_dump") -> None:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
# Save raw tutorials list for inspection
|
||||
with open(os.path.join(out_dir, "tutorials_raw.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(tutorials, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def _iter_page_candidates(item: Dict[str, Any]) -> Iterable[Dict[str, Any]]:
|
||||
# If the tutorial entry already has pages array, yield them
|
||||
if isinstance(item.get("pages"), list):
|
||||
for p in item["pages"]:
|
||||
yield p
|
||||
# Else, try fetching pages via tutorial id/slug
|
||||
tutorial_id = _extract_page_id(item)
|
||||
if tutorial_id:
|
||||
pages = fetch_tutorial_pages(session, tutorial_id)
|
||||
for p in pages:
|
||||
yield p
|
||||
# Lastly, treat the tutorial itself as a single page if it has an id/slug
|
||||
if tutorial_id:
|
||||
yield {"id": tutorial_id, "title": item.get("title")}
|
||||
|
||||
seen = 0
|
||||
for idx, item in enumerate(tutorials, start=1):
|
||||
for page_entry in _iter_page_candidates(item):
|
||||
page_id = _extract_page_id(page_entry)
|
||||
if not page_id:
|
||||
print(f"[{idx:03d}] skipped page (no id): {page_entry}")
|
||||
continue
|
||||
try:
|
||||
page = fetch_page(session, page_id)
|
||||
except requests.HTTPError as e:
|
||||
print(f"[{idx:03d}] page {page_id} -> HTTP {e.response.status_code} ({page_entry})")
|
||||
continue
|
||||
seen += 1
|
||||
title = page.get("title") or page_entry.get("title") or item.get("title") or f"page_{page_id}"
|
||||
out_path = os.path.join(out_dir, f"{idx:03d}_{seen:02d}_{page_id}.json")
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
json.dump(page, f, ensure_ascii=False, indent=2)
|
||||
snippet = page.get("code") or page.get("content") or str(page)[:120]
|
||||
print(f"[{idx:03d}] saved {title} -> {out_path}; sample: {str(snippet)[:80]}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
email = input("BRAIN email: ").strip()
|
||||
password = getpass.getpass("BRAIN password: ")
|
||||
session = authenticate(email, password)
|
||||
tutorials = fetch_tutorials(session)
|
||||
print(f"Fetched {len(tutorials)} tutorials")
|
||||
dump_pages(session, tutorials)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
import getpass
|
||||
import os
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# Make ace_lib importable
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
ROOT_DIR = os.path.dirname(SCRIPT_DIR)
|
||||
if ROOT_DIR not in sys.path:
|
||||
sys.path.append(ROOT_DIR)
|
||||
|
||||
import ace_lib # noqa: E402
|
||||
|
||||
|
||||
def prompt_credentials() -> tuple[str, str]:
|
||||
email = input("Enter BRAIN Email: ").strip()
|
||||
while not email:
|
||||
email = input("Email is required. Enter BRAIN Email: ").strip()
|
||||
|
||||
password = getpass.getpass("Enter BRAIN Password: ").strip()
|
||||
while not password:
|
||||
password = getpass.getpass("Password is required. Enter BRAIN Password: ").strip()
|
||||
|
||||
return email, password
|
||||
|
||||
|
||||
def fetch_operators(session: ace_lib.SingleSession) -> pd.DataFrame:
|
||||
df = ace_lib.get_operators(session)
|
||||
if df is None or df.empty:
|
||||
return pd.DataFrame()
|
||||
|
||||
df = df.copy()
|
||||
|
||||
# Choose an identifier column robustly
|
||||
id_col = "id" if "id" in df.columns else None
|
||||
if id_col is None:
|
||||
if "name" in df.columns:
|
||||
id_col = "name"
|
||||
else:
|
||||
id_col = "_row_id"
|
||||
df[id_col] = df.index
|
||||
|
||||
# Re-aggregate scopes so each operator id is unique
|
||||
if "scope" in df.columns:
|
||||
scope_map = (
|
||||
df.groupby(id_col)["scope"]
|
||||
.agg(lambda x: sorted(set([item for item in x if pd.notna(item)])))
|
||||
.rename("scopes")
|
||||
.reset_index()
|
||||
)
|
||||
else:
|
||||
scope_map = pd.DataFrame({id_col: df[id_col].unique(), "scopes": [[] for _ in range(df[id_col].nunique())]})
|
||||
|
||||
unique_df = df.drop(columns=["scope"], errors="ignore").drop_duplicates(subset=[id_col]).merge(
|
||||
scope_map, on=id_col, how="left"
|
||||
)
|
||||
|
||||
# Sort for readability
|
||||
sort_cols: List[str] = [col for col in ["category", "subcategory", "name", id_col] if col in unique_df.columns]
|
||||
if sort_cols:
|
||||
unique_df = unique_df.sort_values(sort_cols).reset_index(drop=True)
|
||||
|
||||
return unique_df
|
||||
|
||||
|
||||
def main():
|
||||
print("=== Fetch All BRAIN Operators ===")
|
||||
|
||||
email, password = prompt_credentials()
|
||||
ace_lib.get_credentials = lambda: (email, password)
|
||||
|
||||
print("Logging in...")
|
||||
try:
|
||||
session = ace_lib.start_session()
|
||||
print("Login successful.")
|
||||
except Exception as exc:
|
||||
print(f"Login failed: {exc}")
|
||||
return
|
||||
|
||||
print("Fetching operators...")
|
||||
try:
|
||||
operators_df = fetch_operators(session)
|
||||
except Exception as exc:
|
||||
print(f"Failed to fetch operators: {exc}")
|
||||
return
|
||||
|
||||
if operators_df.empty:
|
||||
print("No operators returned; nothing to save.")
|
||||
return
|
||||
|
||||
output_path = os.path.join(SCRIPT_DIR, "all_operators.csv")
|
||||
operators_df.to_csv(output_path, index=False)
|
||||
print(f"Saved {len(operators_df)} operators to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,180 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import pandas as pd
|
||||
from pandas.io.formats.style import Styler
|
||||
|
||||
brain_api_url = os.environ.get("BRAIN_API_URL", "https://api.worldquantbrain.com")
|
||||
brain_url = os.environ.get("BRAIN_URL", "https://platform.worldquantbrain.com")
|
||||
|
||||
|
||||
def make_clickable_alpha_id(alpha_id: str) -> str:
|
||||
"""
|
||||
Create a clickable HTML link for an alpha ID.
|
||||
|
||||
Args:
|
||||
alpha_id (str): The ID of the alpha.
|
||||
|
||||
Returns:
|
||||
str: An HTML string containing a clickable link to the alpha's page on the platform.
|
||||
"""
|
||||
|
||||
url = brain_url + "/alpha/"
|
||||
return f'<a href="{url}{alpha_id}">{alpha_id}</a>'
|
||||
|
||||
|
||||
def prettify_result(
|
||||
result: list, detailed_tests_view: bool = False, clickable_alpha_id: bool = False
|
||||
) -> Union[pd.DataFrame, Styler]:
|
||||
"""
|
||||
Combine and format simulation results into a single DataFrame for analysis.
|
||||
|
||||
Args:
|
||||
result (list): A list of dictionaries containing simulation results.
|
||||
detailed_tests_view (bool, optional): If True, include detailed test results. Defaults to False.
|
||||
clickable_alpha_id (bool, optional): If True, make alpha IDs clickable. Defaults to False.
|
||||
|
||||
Returns:
|
||||
pandas.DataFrame or pandas.io.formats.style.Styler: A DataFrame containing formatted results,
|
||||
optionally with clickable alpha IDs.
|
||||
"""
|
||||
list_of_is_stats = [result[x]["is_stats"] for x in range(len(result)) if result[x]["is_stats"] is not None]
|
||||
is_stats_df = pd.concat(list_of_is_stats).reset_index(drop=True)
|
||||
is_stats_df = is_stats_df.sort_values("fitness", ascending=False)
|
||||
|
||||
expressions = {
|
||||
result[x]["alpha_id"]: (
|
||||
{
|
||||
"selection": result[x]["simulate_data"]["selection"],
|
||||
"combo": result[x]["simulate_data"]["combo"],
|
||||
}
|
||||
if result[x]["simulate_data"]["type"] == "SUPER"
|
||||
else result[x]["simulate_data"]["regular"]
|
||||
)
|
||||
for x in range(len(result))
|
||||
if result[x]["is_stats"] is not None
|
||||
}
|
||||
expression_df = pd.DataFrame(list(expressions.items()), columns=["alpha_id", "expression"])
|
||||
|
||||
list_of_is_tests = [result[x]["is_tests"] for x in range(len(result)) if result[x]["is_tests"] is not None]
|
||||
is_tests_df = pd.concat(list_of_is_tests, sort=True).reset_index(drop=True)
|
||||
is_tests_df = is_tests_df[is_tests_df["result"] != "WARNING"]
|
||||
if detailed_tests_view:
|
||||
cols = ["limit", "result", "value"]
|
||||
is_tests_df["details"] = is_tests_df[cols].to_dict(orient="records")
|
||||
is_tests_df = is_tests_df.pivot(index="alpha_id", columns="name", values="details").reset_index()
|
||||
else:
|
||||
is_tests_df = is_tests_df.pivot(index="alpha_id", columns="name", values="result").reset_index()
|
||||
|
||||
alpha_stats = pd.merge(is_stats_df, expression_df, on="alpha_id")
|
||||
alpha_stats = pd.merge(alpha_stats, is_tests_df, on="alpha_id")
|
||||
alpha_stats = alpha_stats.drop(columns=alpha_stats.columns[(alpha_stats == "PENDING").any()])
|
||||
alpha_stats.columns = alpha_stats.columns.str.replace("(?<=[a-z])(?=[A-Z])", "_", regex=True).str.lower()
|
||||
if clickable_alpha_id:
|
||||
return alpha_stats.style.format({"alpha_id": lambda x: make_clickable_alpha_id(str(x))})
|
||||
return alpha_stats
|
||||
|
||||
|
||||
def concat_pnl(result: list) -> pd.DataFrame:
|
||||
"""
|
||||
Combine PnL results from multiple alphas into a single DataFrame.
|
||||
|
||||
Args:
|
||||
result (list): A list of dictionaries containing simulation results with PnL data.
|
||||
|
||||
Returns:
|
||||
pandas.DataFrame: A DataFrame containing combined PnL data for all alphas.
|
||||
"""
|
||||
list_of_pnls = [result[x]["pnl"] for x in range(len(result)) if result[x]["pnl"] is not None]
|
||||
pnls_df = pd.concat(list_of_pnls).reset_index()
|
||||
|
||||
return pnls_df
|
||||
|
||||
|
||||
def concat_is_tests(result: list) -> pd.DataFrame:
|
||||
"""
|
||||
Combine in-sample test results from multiple alphas into a single DataFrame.
|
||||
|
||||
Args:
|
||||
result (list): A list of dictionaries containing simulation results with in-sample test data.
|
||||
|
||||
Returns:
|
||||
pandas.DataFrame: A DataFrame containing combined in-sample test results for all alphas.
|
||||
"""
|
||||
is_tests_list = [result[x]["is_tests"] for x in range(len(result)) if result[x]["is_tests"] is not None]
|
||||
is_tests_df = pd.concat(is_tests_list, sort=True).reset_index(drop=True)
|
||||
return is_tests_df
|
||||
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user