diff --git a/rdagent/scenarios/qlib/experiment/utils.py b/rdagent/scenarios/qlib/experiment/utils.py index f2ff6375..5b81505b 100644 --- a/rdagent/scenarios/qlib/experiment/utils.py +++ b/rdagent/scenarios/qlib/experiment/utils.py @@ -50,6 +50,56 @@ def generate_data_folder_from_qlib(): ) +def get_file_desc(p: Path) -> str: + """ + Get the description of a file based on its type. + + Parameters + ---------- + p : Path + The path of the file. + + Returns + ------- + str + The description of the file. + """ + p = Path(p) + + JJ_TPL = Environment(undefined=StrictUndefined).from_string( + """ +{{file_name}} +```{{type_desc}} +{{content}} +``` +""" + ) + + if p.name.endswith(".h5"): + df = pd.read_hdf(p) + # get df.head() as string with full width + pd.set_option("display.max_columns", None) # or 1000 + pd.set_option("display.max_rows", None) # or 1000 + pd.set_option("display.max_colwidth", None) # or 199 + return JJ_TPL.render( + file_name=p.name, + type_desc="generated by `pd.read_hdf(filename).head()`", + content=df.head().to_string(), + ) + elif p.name.endswith(".md"): + with open(p) as f: + content = f.read() + return JJ_TPL.render( + file_name=p.name, + type_desc="markdown", + content=content, + ) + else: + raise NotImplementedError( + f"file type {p.name} is not supported. Please implement its description function.", + ) + + def get_data_folder_intro(fname_reg: str = ".*", flags=0) -> str: """ Directly get the info of the data folder. @@ -76,41 +126,8 @@ def get_data_folder_intro(fname_reg: str = ".*", flags=0) -> str: # FIXME: (xiao) I think this is writing in a hard-coded way. # get data folder intro does not imply that we are generating the data folder. generate_data_folder_from_qlib() - - JJ_TPL = Environment(undefined=StrictUndefined).from_string( - """ -{{file_name}} -```{{type_desc}} -{{content}} -``` -""" - ) content_l = [] for p in Path(FACTOR_IMPLEMENT_SETTINGS.data_folder_debug).iterdir(): if re.match(fname_reg, p.name, flags) is not None: - if p.name.endswith(".h5"): - df = pd.read_hdf(p) - # get df.head() as string with full width - pd.set_option("display.max_columns", None) # or 1000 - pd.set_option("display.max_rows", None) # or 1000 - pd.set_option("display.max_colwidth", None) # or 199 - rendered = JJ_TPL.render( - file_name=p.name, - type_desc="generated by `pd.read_hdf(filename).head()`", - content=df.head().to_string(), - ) - content_l.append(rendered) - elif p.name.endswith(".md"): - with open(p) as f: - content = f.read() - rendered = JJ_TPL.render( - file_name=p.name, - type_desc="markdown", - content=content, - ) - content_l.append(rendered) - else: - raise NotImplementedError( - f"file type {p.name} is not supported. Please implement its description function.", - ) + content_l.append(get_file_desc(p)) return "\n----------------- file splitter -------------\n".join(content_l)