From 6c007ee7fbb18793d3e97ce4c47034e8e11bd14c Mon Sep 17 00:00:00 2001 From: amstrongzyf <201840057@smail.nju.edu.cn> Date: Sat, 28 Jun 2025 21:23:40 +0800 Subject: [PATCH] refactor: convert data describe functions to class-based implementation and add output example (#999) * feat: Enhance data folder description for clarity and robustness * fix bug * fix present bugs * delete useless files * add output example and refactor the hole util.py * fix bug for file tree * add corner case example * delete useless file --- .../scenarios/data_science/scen/__init__.py | 5 +- rdagent/scenarios/data_science/scen/utils.py | 920 +++++++++++------- 2 files changed, 572 insertions(+), 353 deletions(-) diff --git a/rdagent/scenarios/data_science/scen/__init__.py b/rdagent/scenarios/data_science/scen/__init__.py index 808297e8..fac00486 100644 --- a/rdagent/scenarios/data_science/scen/__init__.py +++ b/rdagent/scenarios/data_science/scen/__init__.py @@ -10,10 +10,7 @@ from rdagent.core.scenario import Scenario from rdagent.log import rdagent_logger as logger from rdagent.oai.llm_utils import APIBackend from rdagent.scenarios.data_science.debug.data import create_debug_data -from rdagent.scenarios.data_science.scen.utils import ( - describe_data_folder, - describe_data_folder_v2, -) +from rdagent.scenarios.data_science.scen.utils import describe_data_folder_v2 from rdagent.scenarios.kaggle.kaggle_crawler import ( crawl_descriptions, download_data, diff --git a/rdagent/scenarios/data_science/scen/utils.py b/rdagent/scenarios/data_science/scen/utils.py index 6f112016..3517ae8c 100644 --- a/rdagent/scenarios/data_science/scen/utils.py +++ b/rdagent/scenarios/data_science/scen/utils.py @@ -1,227 +1,85 @@ -import os +""" +An example of the generated data folder description: -import pandas as pd -from PIL import Image, TiffTags - -from rdagent.log import rdagent_logger as logger - -""" data folder description version 1 """ +## File tree: +``` +plant-pathology-2020-fgvc7/ +├── images/ +│ ├── Test_0.jpg (182.7 kB) +│ ├── Test_1.jpg (362.4 kB) +│ ├── ... (+1819 more files) +├── train.csv (30.1 kB) +├── description.md (5.3 kB) +├── sample_submission.csv (5.2 kB) +├── test.csv (1.5 kB)``` -def read_csv_head(file_path, indent=0, lines=5, max_col_width=100): - """ - Reads the first few rows of a CSV file and formats them with indentation and optional truncation. +## File details: - Parameters: - file_path (str): Path to the CSV file. - indent (int): Number of spaces to prepend to each line for indentation. - lines (int): Number of rows to read from the CSV file. - max_col_width (int): Maximum width of each column's content. + (Showing details for representative files out of many) - Returns: - str: A formatted string of the first few rows of the CSV file. - """ - try: - # Read the CSV file with specified rows - df = pd.read_csv(file_path, nrows=lines) +### sample_submission.csv: +#### 1.DataFrame preview: +It has 183 rows and 5 columns. +Here is some information about the columns: +healthy (float64) has 1 unique values: [0.25] +image_id (object) has 183 unique values. Some example values: ['Test_0', 'Test_1', 'Test_2', 'Test_3'] +multiple_diseases (float64) has 1 unique values: [0.25] +rust (float64) has 1 unique values: [0.25] +scab (float64) has 1 unique values: [0.25] +#### 2.DataFrame preview:(only show the first 5 rows and 15 columns) + image_id healthy multiple_diseases rust scab +0 Test_0 0.25 0.25 0.25 0.25 +1 Test_1 0.25 0.25 0.25 0.25 +2 Test_2 0.25 0.25 0.25 0.25 +3 Test_3 0.25 0.25 0.25 0.25 +4 Test_4 0.25 0.25 0.25 0.25 - if df.empty: - return " " * indent + "(No data in the file)" +### test.csv: +#### 1.DataFrame preview: +It has 183 rows and 1 columns. +Here is some information about the columns: +image_id (object) has 183 unique values. Some example values: ['Test_0', 'Test_1', 'Test_2', 'Test_3'] +#### 2.DataFrame preview:(only show the first 5 rows and 15 columns) + image_id +0 Test_0 +1 Test_1 +2 Test_2 +3 Test_3 +4 Test_4 - # Truncate column contents to a maximum width - truncated_df = df.copy() - for col in truncated_df.columns: - truncated_df[col] = ( - truncated_df[col] - .astype(str) - .apply(lambda x: (x[:max_col_width] + "...") if len(x) > max_col_width else x) - ) +### train.csv: +#### 1.DataFrame preview: +It has 1638 rows and 5 columns. +Here is some information about the columns: +healthy (int64) has 2 unique values: [0, 1] +image_id (object) has 1638 unique values. Some example values: ['Train_1637', 'Train_0', 'Train_1', 'Train_2'] +multiple_diseases (int64) has 2 unique values: [0, 1] +rust (int64) has 2 unique values: [1, 0] +scab (int64) has 2 unique values: [0, 1] +#### 2.DataFrame preview:(only show the first 5 rows and 15 columns) + image_id healthy multiple_diseases rust scab +0 Train_0 0 0 1 0 +1 Train_1 1 0 0 0 +2 Train_2 0 0 1 0 +3 Train_3 1 0 0 0 +4 Train_4 0 0 1 0 - # Convert DataFrame to a string representation - df_string_lines = truncated_df.to_string(index=False).split("\n") +""" - # Add indentation to each line - indented_lines = [" " * indent + line for line in df_string_lines] - - return "\n".join(indented_lines) - except FileNotFoundError: - return f"Error: File not found at path '{file_path}'." - except pd.errors.EmptyDataError: - return f"Error: The file at '{file_path}' is empty." - except Exception as e: - return f"Error reading CSV: {e}" - - -def get_dir_snapshot(folder_path): - """ - [note] - - Returns a set of file extensions within the subfolder (excluding subfolder names) - - Compares only the types of files contained, not specific file names or quantities - """ - exts = set() - try: - with os.scandir(folder_path) as it: - for entry in it: - if entry.is_file(): - file_ext = os.path.splitext(entry.name)[1] - exts.add(file_ext) - except Exception as e: - logger.error(f"Error scanning directory: {e}") - - return frozenset(exts) - - -def describe_data_folder(folder_path, indent=0, max_files=2, partial_expand_subfolders=2, is_top_level=True): - """ - folder_path : Current directory path - indent : Current indentation - max_files : Maximum number of files of the same type to display - partial_expand_subfolders: When all subfolders have the same internal file types, only expand this many subfolders, the rest are omitted - is_top_level : Indicates if the current folder is the top-level folder - """ - result = [] - files_count = {} - files_details = {} - - for root, dirs, files in os.walk(folder_path): - dirs.sort() - files.sort() - if not dirs: - for file in files: - file_path = os.path.join(root, file) - file_type = os.path.splitext(file)[1][1:] - file_size = os.path.getsize(file_path) - - if file_type not in files_count: - files_count[file_type] = 0 - files_details[file_type] = [] - files_count[file_type] += 1 - - # At top level, collect all CSV and Markdown files without restrictions - # In deeper levels, follow the max_files restriction - if is_top_level and file_type in ["csv", "md"]: - files_details[file_type].append((file, file_size, file_path)) - elif len(files_details[file_type]) < max_files: - files_details[file_type].append((file, file_size, file_path)) - break - - # Collect "type snapshots" of subfolders - snapshots = [] - for d in dirs: - subfolder_path = os.path.join(root, d) - snapshot = get_dir_snapshot(subfolder_path) - snapshots.append(snapshot) - - # Determine if all subfolders have the same file type distribution - first_snapshot = snapshots[0] - all_same_structure = all(s == first_snapshot for s in snapshots) - - if all_same_structure: - for i, d in enumerate(dirs): - if i < partial_expand_subfolders: - result.append(" " * indent + f"- Folder: {d}") - subfolder_path = os.path.join(root, d) - result.append( - describe_data_folder( - folder_path=subfolder_path, - indent=indent + 2, - max_files=max_files, - partial_expand_subfolders=partial_expand_subfolders, - is_top_level=False, - ) - ) - else: - remaining = len(dirs) - i - result.append(" " * indent + f"... ({remaining} more subfolders)") - break - else: - for d in dirs: - result.append(" " * indent + f"- Folder: {d}") - subfolder_path = os.path.join(root, d) - result.append( - describe_data_folder( - folder_path=subfolder_path, - indent=indent + 2, - max_files=max_files, - partial_expand_subfolders=partial_expand_subfolders, - is_top_level=False, - ) - ) - - for file in files: - file_path = os.path.join(root, file) - file_type = os.path.splitext(file)[1][1:] - file_size = os.path.getsize(file_path) - - if file_type not in files_count: - files_count[file_type] = 0 - files_details[file_type] = [] - files_count[file_type] += 1 - - # At top level, collect all CSV and Markdown files without restrictions - # In deeper levels, follow the max_files restriction - if is_top_level and file_type in ["csv", "md"]: - files_details[file_type].append((file, file_size, file_path)) - elif not is_top_level and len(files_details[file_type]) <= max_files: - files_details[file_type].append((file, file_size, file_path)) - - break - - # Print the folder and its contents - for file_type, count in files_count.items(): - if count > max_files and file_type not in ["csv", "md", "txt"]: - result.append(" " * indent + f"{count} {file_type}s:") - for file, size, path in files_details[file_type]: - result.append(" " * (indent + 2) + f"- {file} ({size} bytes)") - result.append(" " * (indent + 2) + "... (file limit reached)") - else: - for file, size, path in files_details[file_type]: - if file_type == "csv": - df = pd.read_csv(path) - result.append( - " " * indent + f"- {file} ({size} bytes, with {df.shape[0]} rows and {df.shape[1]} columns)" - ) - result.append(" " * (indent + 2) + f"- Head of {file}:") - csv_head = read_csv_head(path, indent + 4) - result.append(csv_head) - continue - result.append(" " * indent + f"- {file} ({size} bytes)") - if file_type == "md": - result.append(" " * (indent + 2) + f"- Content of {file}:") - if file == "description.md": - result.append(" " * (indent + 4) + f"Please refer to the background of the scenario context.") - continue - with open(path, "r", encoding="utf-8") as f: - result.append(" " * (indent + 4) + f.read()) - if file_type == "tif": - result.append(" " * (indent + 2) + f"- Metadata of {file}:") - with Image.open(path) as img: - for tag, value in img.tag_v2.items(): - tag_name = TiffTags.TAGS_V2.get(tag, f"Unknown Tag {tag}") - result.append(" " * (indent + 4) + f"{tag_name}: {value}") - if file_type in ["json", "txt"]: - result.append(" " * (indent + 2) + f"- Content of {file}:") - with open(path, "r", encoding="utf-8") as f: - for i, line in enumerate(f): - if i < 2: - result.append( - " " * (indent + 4) + line.strip()[:100] + ("..." if len(line.strip()) > 100 else "") - ) - else: - break - - return "\n".join(result) + "\n" - - -""" data folder description version 2 """ import json +import os +import reprlib +from collections import defaultdict from pathlib import Path +from typing import Dict, List, Optional, Set, Tuple, Union import humanize import pandas as pd -from genson import SchemaBuilder from pandas.api.types import is_numeric_dtype +from rdagent.log import rdagent_logger as logger + # these files are treated as code (e.g. markdown wrapped) code_files = {".py", ".sh", ".yaml", ".yml", ".md", ".html", ".xml", ".log", ".rst"} # we treat these files as text (rather than binary) files @@ -230,7 +88,25 @@ plaintext_files = {".txt", ".csv", ".json", ".tsv"} | code_files system_names = {"__MACOSX", ".DS_Store", "Thumbs.db"} -def get_file_len_size(f: Path) -> tuple[int, str]: +class FileTreeGenerationError(Exception): + """File tree generation related errors""" + + pass + + +class MaxLinesExceededError(FileTreeGenerationError): + """Raised when max lines limit is exceeded""" + + pass + + +class DirectoryPermissionError(FileTreeGenerationError): + """Raised when directory access is denied""" + + pass + + +def get_file_len_size(f: Path) -> Tuple[int, str]: """ Calculate the size of a file (#lines for plaintext files, otherwise #bytes) Also returns a human-readable string representation of the size. @@ -243,83 +119,13 @@ def get_file_len_size(f: Path) -> tuple[int, str]: return s, humanize.naturalsize(s) -def file_tree(path: Path, depth=0) -> str: - """Generate a tree structure of files in a directory""" - result = [] - - files = [p for p in Path(path).iterdir() if not p.is_dir() and p.name not in system_names] - - max_n = 4 if len(files) > 30 else 8 - for p in sorted(files)[:max_n]: - result.append(f"{' '*depth*4}{p.name} ({get_file_len_size(p)[1]})") - if len(files) > max_n: - result.append(f"{' '*depth*4}... and {len(files)-max_n} other files") - - dirs = [ - p - for p in Path(path).iterdir() - if (p.is_dir() or (p.is_symlink() and p.resolve().is_dir())) and p.name not in system_names - ] - - # Calculate base_path (the top-level resolved absolute directory) - base_path = Path(path).resolve() - # Find the top-level base_path when in recursion (depth>0) - if depth > 0: - # The top-level base_path is the ancestor at depth==0 - ancestor = Path(path) - for _ in range(depth): - ancestor = ancestor.parent - base_path = ancestor.resolve() - - for p in sorted(dirs): - if p.is_symlink(): - target = p.resolve() - if str(target).startswith(str(base_path)): - # avoid recursing into symlinks pointing inside base path - result.append( - f"{' ' * depth * 4}{p.name}@ -> {os.path.relpath(target, base_path)} (symlinked dir, not expanded)" - ) - continue - result.append(f"{' ' * depth * 4}{p.name}/") - result.append(file_tree(p, depth + 1)) - - return "\n".join(result) - - -def _walk(path: Path): - """Recursively walk a directory (analogous to os.walk but for pathlib.Path)""" - for p in sorted(Path(path).iterdir()): - # Filter out system-generated directories/files - if p.name in system_names: - continue - - if p.is_dir(): - # If this is a symlinked dir to a parent/ancestor, do not expand it - if p.is_symlink(): - target = p.resolve() - cur_path = p.parent.resolve() - if target == cur_path or str(cur_path).startswith(str(target)): - yield p - continue - yield from _walk(p) - else: - yield p - - def preview_df(df: pd.DataFrame, file_name: str, simple=True, show_nan_columns=False) -> str: - """Generate a textual preview of a dataframe - - Args: - df (pd.DataFrame): the dataframe to preview - file_name (str): the file name to use in the preview - simple (bool, optional): whether to use a simplified version of the preview. Defaults to True. - - Returns: - str: the textual preview - """ + """Generate a textual preview of a dataframe""" out = [] - out.append(f"-> {file_name} has {df.shape[0]} rows and {df.shape[1]} columns.") + out.append(f"### {file_name}: ") + out.append(f"#### 1.DataFrame preview:") + out.append(f"It has {df.shape[0]} rows and {df.shape[1]} columns.") if simple: cols = df.columns.tolist() @@ -353,6 +159,18 @@ def preview_df(df: pd.DataFrame, file_name: str, simple=True, show_nan_columns=F if nan_cols: out.append(f"Columns containing NaN values: {', '.join(nan_cols)}") + # Add: Display DataFrame directly + if len(df) > 0: + out.append("#### 2.DataFrame preview:(only show the first 5 rows and 15 columns)") + # Show first 5 rows, max 15 columns to avoid overly wide output + df_preview = df.head(5) + if df.shape[1] > 15: + df_preview = df_preview.iloc[:, :15] + out.append(str(df_preview)) + out.append(f"... (showing first 15 of {df.shape[1]} columns)") + else: + out.append(str(df_preview)) + return "\n".join(out) @@ -369,78 +187,482 @@ def preview_parquet(p: Path, file_name: str, simple=True, show_nan_columns=False def preview_json(p: Path, file_name: str): - """Generate a textual preview of a json file using a generated json schema""" - builder = SchemaBuilder() - with open(p) as f: - first_line = f.readline().strip() + """Generate a textual preview of a json file using reprlib for compact object display""" + result = [] + result.append(f"### {file_name}:") - try: - first_object = json.loads(first_line) + try: + # First check if this is a JSONL format + is_jsonl = False - if not isinstance(first_object, dict): - raise json.JSONDecodeError("The first line isn't JSON", first_line, 0) - - # if the the next line exists and is not empty, then it is a JSONL file + with open(p, "r", encoding="utf-8") as f: + first_line = f.readline().strip() second_line = f.readline().strip() - if second_line: - f.seek(0) # so reset and read line by line - for line in f: - builder.add_object(json.loads(line.strip())) - # if it is empty, then it's a single JSON object file + + # Correct JSONL detection: both lines must be independent complete JSON objects + if first_line and second_line: + try: + # Try to parse the first two lines, both should be complete JSON objects + json.loads(first_line) # First line is complete JSON + json.loads(second_line) # Second line is also complete JSON + is_jsonl = True + except json.JSONDecodeError: + # If any line fails to parse, it's not JSONL + is_jsonl = False + + if is_jsonl: + # JSONL format: one JSON object per line + result.append("#### 1.Format: JSONL (JSON Lines)") + result.append("#### 2.Content preview (first few objects):") + + # Configure reprlib for JSONL + jsonl_repr = reprlib.Repr() + jsonl_repr.maxother = 300 + + with open(p, "r", encoding="utf-8") as f: + for i, line in enumerate(f): + if i >= 3: # Only show first 3 objects + result.append("... (showing first 3 JSONL objects)") + break + if line.strip(): + try: + obj = json.loads(line.strip()) + result.append(f"Object {i+1}: {jsonl_repr.repr(obj)}") + except: + result.append(f"Object {i+1}: Invalid JSON") + else: + # Single JSON file + with open(p, "r", encoding="utf-8") as f: + data = json.load(f) + + result.append("#### 1.Format: Single JSON object") + result.append("#### 2.Structure overview:") + + # Basic information + if isinstance(data, dict): + result.append(f"Type: Object with {len(data)} keys: {list(data.keys())}") + for key, value in data.items(): + if isinstance(value, list): + result.append(f" - {key}: array[{len(value)}]") + elif isinstance(value, dict): + result.append(f" - {key}: object{{{len(value)} keys}}") + else: + result.append(f" - {key}: {type(value).__name__}") + elif isinstance(data, list): + result.append(f"Type: Array with {len(data)} items") + if len(data) > 0: + sample_item = data[0] + if isinstance(sample_item, dict): + result.append(f"Sample item keys: {list(sample_item.keys())}") + + result.append("#### 3.Content preview (reprlib):") + + # Use reprlib to display content + compact_repr = reprlib.Repr() + compact_repr.maxother = 300 + + result.append(compact_repr.repr(data)) + + except Exception as e: + result.append(f"Error processing JSON: {str(e)[:100]}") + + return "\n".join(result) + + +def _walk(path: Path): + """Recursively walk a directory (analogous to os.walk but for pathlib.Path)""" + for p in sorted(Path(path).iterdir()): + # Filter out system-generated directories/files + if p.name in system_names: + continue + + if p.is_dir(): + # If this is a symlinked dir to a parent/ancestor, do not expand it + if p.is_symlink(): + target = p.resolve() + cur_path = p.parent.resolve() + if target == cur_path or str(cur_path).startswith(str(target)): + yield p + continue + yield from _walk(p) + else: + yield p + + +class FileTreeGenerator: + """ + Smart file tree generator with symlink handling and intelligent truncation. + """ + + def __init__(self, max_lines: int = 200, priority_files: Set[str] = None): + """ + Initialize the file tree generator. + + Args: + max_lines: Maximum output lines to prevent overly long output + priority_files: File extensions to prioritize for display + """ + self.max_lines = max_lines + self.priority_files = priority_files or {".csv", ".json", ".parquet", ".md", ".txt"} + self.lines = [] + self.line_count = 0 + + def generate_tree(self, path: Union[str, Path]) -> str: + """ + Generate a tree structure of files in a directory. + + Args: + path: Target directory path + + Returns: + str: Tree structure representation + + Raises: + FileTreeGenerationError: If tree generation fails + """ + try: + path = Path(path) + base_path = path.resolve() + self.lines = [] + self.line_count = 0 + self._add_line(f"{path.name}/") + self._process_directory(path, 0, "", base_path) + except MaxLinesExceededError: + pass # Expected when hitting line limit + except Exception as e: + raise FileTreeGenerationError(f"Failed to generate tree for {path}: {str(e)}") from e + + # CORNER CASE HANDLING: Always check if we hit the limit and add truncation notice if needed + # + # WHY THIS IS NECESSARY: + # The code uses a "mixed exception handling strategy": + # - Sub-methods (_process_subdirectories, _process_files, _process_single_directory) + # catch MaxLinesExceededError and handle it silently (don't re-raise) + # - This means some MaxLinesExceededError exceptions never propagate to generate_tree + # + # CORNER CASE SCENARIO: + # 1. _add_line() is called and line_count reaches max_lines + # 2. _add_line() throws MaxLinesExceededError + # 3. A sub-method catches the exception but doesn't re-raise it (silent handling) + # 4. The exception never reaches generate_tree's except block above + # 5. OLD VERSION: No truncation notice is added → User doesn't know content was truncated + # 6. NEW VERSION: This check below ensures truncation notice is always added + # + # DEMONSTRATION EXAMPLE (max_lines=5, processing 6 files): + # + # 🔴 OLD VERSION RESULT: + # project/ + # ├── file1.csv + # ├── file2.csv + # ├── file3.csv + # ├── file4.csv + # 🔍 Truncation notice? NO → User doesn't know content was truncated! + # + # 🔵 NEW VERSION RESULT: + # project/ + # ├── file1.csv + # ├── file2.csv + # ├── file3.csv + # ├── file4.csv + # ... (display limited) + # 🔍 Truncation notice? YES → User knows content was truncated! + # + # The key difference: + # - OLD: Relies on exception propagation (fails when sub-methods handle silently) + # - NEW: Active check ensures truncation notice is always present + if self.line_count >= self.max_lines and ( + not self.lines or not self.lines[-1].startswith("... (display limited") + ): + self.lines.append("... (display limited, please increase max_lines parameter)") + + return "\n".join(self.lines) + + def _add_line(self, text: str) -> None: + """ + Add a line to the output. + + Args: + text: Line text to add + + Raises: + MaxLinesExceededError: If max lines limit is exceeded + """ + if self.line_count >= self.max_lines: + raise MaxLinesExceededError(f"Exceeded maximum lines limit of {self.max_lines}") + self.lines.append(text) + self.line_count += 1 + + def _process_directory(self, path: Path, depth: int, prefix: str, base_path: Path) -> None: + """ + Process a single directory. + + Args: + path: Directory path to process + depth: Current depth in the tree + prefix: Prefix for tree formatting + base_path: Base path for symlink detection + + Raises: + DirectoryPermissionError: If directory access is denied + FileTreeGenerationError: If processing fails + MaxLinesExceededError: Propagated when line limit is reached + """ + try: + # Get directory contents, filter out system files + items = [p for p in path.iterdir() if not p.name.startswith(".") and p.name not in system_names] + dirs = sorted([p for p in items if p.is_dir()]) + files = sorted([p for p in items if p.is_file()]) + + # Categorize files + priority_files_list, other_files = self._categorize_files(files) + + # Process subdirectories + self._process_subdirectories(dirs, depth, prefix, base_path) + + # Process files + self._process_files(priority_files_list + other_files, depth, prefix) + + except MaxLinesExceededError: + # Propagate this up so generate_tree can handle it + raise + except PermissionError as e: + raise DirectoryPermissionError(f"Permission denied accessing {path}") from e + except OSError as e: + if e.errno == 13: # Permission denied + raise DirectoryPermissionError(f"Permission denied accessing {path}") from e else: - builder.add_object(first_object) + raise FileTreeGenerationError(f"Error processing directory {path}: {str(e)}") from e - except json.JSONDecodeError: - # if first line isn't JSON, then it's prettified and we can read whole file - f.seek(0) - builder.add_object(json.load(f)) + def _process_subdirectories(self, dirs: List[Path], depth: int, prefix: str, base_path: Path) -> None: + """Process subdirectories with proper truncation logic.""" + try: + if depth == 0 or len(dirs) <= 8: + # First level or ≤8 items: show all + for d in dirs: + self._process_single_directory(d, depth, prefix, base_path) + else: + # Not first level and >8 items: show first 2 + show_count = 2 + for d in dirs[:show_count]: + self._process_single_directory(d, depth, prefix, base_path) - return f"-> {file_name} has auto-generated json schema:\n" + builder.to_json(indent=2) + # Show remaining directory count + remaining = len(dirs) - show_count + self._add_line(f"{prefix}├── ... (+{remaining} more directories)") + except MaxLinesExceededError: + # If we hit the line limit, just stop processing + pass + + def _process_single_directory(self, d: Path, depth: int, prefix: str, base_path: Path) -> None: + """Process a single directory entry.""" + try: + # Handle symlinks + if d.is_symlink(): + target = d.resolve() + if str(target).startswith(str(base_path)): + # avoid recursing into symlinks pointing inside base path + self._add_line( + f"{prefix}├── {d.name}@ -> {os.path.relpath(target, base_path)} (symlinked dir, not expanded)" + ) + return + + self._add_line(f"{prefix}├── {d.name}/") + + # Process subdirectory recursively + child_prefix = prefix + "│ " + self._process_directory(d, depth + 1, child_prefix, base_path) + except MaxLinesExceededError: + # If we hit the line limit, just stop processing this directory + pass + + def _process_files(self, all_files: List[Path], depth: int, prefix: str) -> None: + """Process files with proper truncation logic.""" + try: + if depth == 0 or len(all_files) <= 8: + # First level or ≤8 items: show all + for f in all_files: + self._add_line(f"{prefix}├── {f.name} ({self._get_size_str(f)})") + else: + # Not first level and >8 items: show first 2 + show_count = 2 + for f in all_files[:show_count]: + self._add_line(f"{prefix}├── {f.name} ({self._get_size_str(f)})") + + # Show remaining file count + remaining = len(all_files) - show_count + self._add_line(f"{prefix}├── ... (+{remaining} more files)") + except MaxLinesExceededError: + # If we hit the line limit, just stop processing files + pass + + def _categorize_files(self, files: List[Path]) -> Tuple[List[Path], List[Path]]: + """Categorize files into priority and other groups.""" + priority = [] + other = [] + + for f in files: + if f.suffix.lower() in self.priority_files: + priority.append(f) + else: + other.append(f) + + # Sort priority files by size (larger files first) + priority.sort(key=lambda x: x.stat().st_size if x.exists() else 0, reverse=True) + + return priority, other + + def _get_size_str(self, file_path: Path) -> str: + """Get file size string.""" + try: + size = file_path.stat().st_size + return humanize.naturalsize(size) + except (OSError, FileNotFoundError): + return "? B" + + +class DataFolderDescriptor: + """ + Generate detailed descriptions of data folders including file previews. + """ + + def __init__(self, tree_generator: FileTreeGenerator = None): + """ + Initialize the data folder descriptor. + + Args: + tree_generator: Optional FileTreeGenerator instance + """ + self.tree_generator = tree_generator or FileTreeGenerator() + + def describe_folder( + self, + base_path: Union[str, Path], + include_file_details: bool = True, + simple: bool = False, + show_nan_columns: bool = False, + max_length: int = 10000, + ) -> str: + """ + Generate a textual preview of a directory, including an overview of the directory + structure and previews of individual files. + """ + base_path = Path(base_path) + + tree = f"## File tree:\n```\n{self.tree_generator.generate_tree(base_path)}```" + out = [tree] + + if include_file_details: + out.append("\n## File details:") + + # Intelligently select a subset of files to preview + files_to_preview = self._select_files_for_preview(base_path) + out.append(f" (Showing details for representative files out of many)") + + for fn in files_to_preview: + try: + file_name = str(fn.relative_to(base_path)) + except ValueError: + file_name = str(fn) + + try: + if fn.suffix == ".csv": + out.append(preview_csv(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) + elif fn.suffix == ".json": + out.append(preview_json(fn, file_name)) + elif fn.suffix == ".parquet": + out.append(preview_parquet(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) + elif fn.suffix in plaintext_files: + if get_file_len_size(fn)[0] < 30: + with open(fn) as f: + content = f.read() + if fn.suffix in code_files: + content = f"```\n{content}\n```" + out.append(f"-> {file_name} has content:\n\n{content}") + except Exception as e: + out.append(f"-> {file_name}: Error reading file - {str(e)[:100]}") + + if len("\n\n".join(out)) > max_length: + out.append("\n... (File details truncated due to max_length)") + break + + result = "\n\n".join(out) + + # if the result is very long we generate a simpler version + if len(result) > max_length and not simple: + return self.describe_folder( + base_path, + include_file_details=include_file_details, + simple=True, + show_nan_columns=show_nan_columns, + max_length=max_length, + ) + # if still too long, we truncate + if len(result) > max_length and simple: + return result[:max_length] + "\n... (truncated)" + + return result + + def _select_files_for_preview( + self, base_path: Path, max_files_per_group: int = 1, threshold: int = 10 + ) -> List[Path]: + """ + Intelligently select a representative subset of files for detailed preview. + If a directory has more than `threshold` files of the same type, only `max_files_per_group` are selected. + """ + # Group files by (parent_directory, file_extension) + files_by_group = defaultdict(list) + for p in _walk(base_path): + if p.is_file(): + files_by_group[(p.parent, p.suffix)].append(p) + + selected_files = [] + + # Always include a root README.md if it exists + root_readme = base_path / "README.md" + if root_readme.exists(): + selected_files.append(root_readme) + + for group, files in files_by_group.items(): + # Sort files to be deterministic (e.g., by name) + files.sort() + + if root_readme in files: + try: + files.remove(root_readme) + except ValueError: + pass # was not in list + + if len(files) > threshold: + # If many files, select a few representatives + selected_files.extend(files[:max_files_per_group]) + else: + # If few files, select all of them + selected_files.extend(files) + + # Remove duplicates and maintain order + return list(dict.fromkeys(selected_files)) + + +# Convenience functions for backward compatibility +def file_tree_v2(path: Union[str, Path], max_lines: int = 200, priority_files: Set[str] = None) -> str: + """Generate a file tree using FileTreeGenerator.""" + generator = FileTreeGenerator(max_lines=max_lines, priority_files=priority_files) + return generator.generate_tree(path) def describe_data_folder_v2( - base_path, include_file_details=True, simple=False, show_nan_columns=False, max_length: int = 6_000 -): - """ - Generate a textual preview of a directory, including an overview of the directory - structure and previews of individual files - """ - tree = f"```\n{file_tree(base_path)}```" - out = [tree] - - if include_file_details: - for fn in _walk(base_path): - file_name = str(fn.relative_to(base_path)) - - if fn.suffix == ".csv": - out.append(preview_csv(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) - elif fn.suffix == ".json": - out.append(preview_json(fn, file_name)) - elif fn.suffix == ".parquet": - out.append(preview_parquet(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) - elif fn.suffix in plaintext_files: - if get_file_len_size(fn)[0] < 30: - with open(fn) as f: - content = f.read() - if fn.suffix in code_files: - content = f"```\n{content}\n```" - out.append(f"-> {file_name} has content:\n\n{content}") - if len("\n\n".join(out)) > max_length: - break - - result = "\n\n".join(out) - - # if the result is very long we generate a simpler version - if len(result) > max_length and not simple: - return describe_data_folder_v2( - base_path, - include_file_details=include_file_details, - simple=True, - show_nan_columns=show_nan_columns, - max_length=max_length, - ) - # if still too long, we truncate - if len(result) > max_length and simple: - return result[:max_length] + "\n... (truncated)" - - return result + base_path: Union[str, Path], + include_file_details: bool = True, + simple: bool = False, + show_nan_columns: bool = False, + max_length: int = 10000, +) -> str: + """Generate a data folder description using DataFolderDescriptor.""" + descriptor = DataFolderDescriptor() + return descriptor.describe_folder( + base_path, + include_file_details=include_file_details, + simple=simple, + show_nan_columns=show_nan_columns, + max_length=max_length, + )