diff --git a/src/gpt_copy/gpt_copy.py b/src/gpt_copy/gpt_copy.py index 720645d..50c4aa5 100755 --- a/src/gpt_copy/gpt_copy.py +++ b/src/gpt_copy/gpt_copy.py @@ -437,6 +437,7 @@ def write_output( tree_output: str, file_sections: list[str], unrecognized_files: list[str], + tree_only: bool = False, ) -> None: """ Write collected outputs to the specified output (file or stdout). @@ -446,21 +447,28 @@ def write_output( tree_output (str): The generated folder structure tree. file_sections (List[str]): The collected file sections. unrecognized_files (List[str]): The list of unrecognized files. + tree_only (bool): If True, only output the tree structure. """ - output.write("# Folder Structure\n\n```\n") - output.write(tree_output) - output.write("\n```\n\n") + if tree_only: + # Only output the tree structure without markdown headers + output.write(tree_output) + output.write("\n") + else: + # Original behavior: output everything with markdown formatting + output.write("# Folder Structure\n\n```\n") + output.write(tree_output) + output.write("\n```\n\n") - for section in file_sections: - output.write(section) + for section in file_sections: + output.write(section) - if unrecognized_files: - output.write("# Unrecognized Files\n\n") - output.write( - "The following files were not recognized by extension and were skipped:\n\n" - ) - for rel_path in unrecognized_files: - output.write(f"- `{rel_path}`\n") + if unrecognized_files: + output.write("# Unrecognized Files\n\n") + output.write( + "The following files were not recognized by extension and were skipped:\n\n" + ) + for rel_path in unrecognized_files: + output.write(f"- `{rel_path}`\n") @click.command() @@ -507,6 +515,12 @@ def write_output( default=False, help="Disable line numbers for file content.", ) +@click.option( + "--tree-only", + is_flag=True, + default=False, + help="Output only the folder structure tree without file contents.", +) def main( root_path: Path, output_file: str | None, @@ -514,6 +528,7 @@ def main( include_patterns: tuple[str, ...], exclude_patterns: tuple[str, ...], no_line_numbers: bool, + tree_only: bool, ) -> None: """ Main function to start the script. @@ -525,6 +540,7 @@ def main( include_patterns (Tuple[str, ...]): The tuple of include glob patterns. exclude_patterns (Tuple[str, ...]): The tuple of exclude glob patterns. no_line_numbers (bool): If True, disable line numbers. + tree_only (bool): If True, output only the folder structure tree. """ root_path = root_path.resolve() @@ -533,22 +549,29 @@ def main( tree_output = generate_tree( root_path, gitignore_specs, tracked_files, list(exclude_patterns) ) - file_sections, unrecognized_files = collect_files_content( - root_path, - gitignore_specs, - output_file, - tracked_files, - include_patterns=list(include_patterns), - exclude_patterns=list(exclude_patterns), - line_numbers=not no_line_numbers, - ) + + if tree_only: + # Only output the tree structure + file_sections = [] + unrecognized_files = [] + else: + # Collect file contents as usual + file_sections, unrecognized_files = collect_files_content( + root_path, + gitignore_specs, + output_file, + tracked_files, + include_patterns=list(include_patterns), + exclude_patterns=list(exclude_patterns), + line_numbers=not no_line_numbers, + ) if output_file: print(f"Writing output to {output_file}...", file=sys.stderr) with open(output_file, "w", encoding="utf-8") as out: - write_output(out, tree_output, file_sections, unrecognized_files) + write_output(out, tree_output, file_sections, unrecognized_files, tree_only) else: - write_output(sys.stdout, tree_output, file_sections, unrecognized_files) + write_output(sys.stdout, tree_output, file_sections, unrecognized_files, tree_only) print(f"All files merged into {output_file or 'stdout'}", file=sys.stderr) if unrecognized_files: diff --git a/tests/test_tree_only.py b/tests/test_tree_only.py new file mode 100644 index 0000000..a9e91f7 --- /dev/null +++ b/tests/test_tree_only.py @@ -0,0 +1,101 @@ +# File: tests/test_tree_only.py + +import tempfile +import shutil +from pathlib import Path +from click.testing import CliRunner + +from gpt_copy.gpt_copy import main + + +def test_tree_only_flag(): + """Test that --tree-only outputs only the folder structure.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create test files + (temp_path / "file1.py").write_text("print('hello')") + (temp_path / "file2.md").write_text("# Header") + subdir = temp_path / "subdir" + subdir.mkdir() + (subdir / "file3.txt").write_text("content") + + runner = CliRunner() + + # Test tree-only output + result = runner.invoke(main, [str(temp_path), "--tree-only"]) + assert result.exit_code == 0 + + # Should contain tree structure + assert temp_path.name in result.output + assert "├── file1.py" in result.output or "file1.py" in result.output + assert "├── file2.md" in result.output or "file2.md" in result.output + assert "subdir" in result.output + + # Should NOT contain file contents + assert "print('hello')" not in result.output + assert "# Header" not in result.output + assert "content" not in result.output + + # Should NOT contain markdown headers + assert "# Folder Structure" not in result.output + assert "## File:" not in result.output + + +def test_tree_only_vs_normal_output(): + """Test that tree-only output is different from normal output.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create a test file + (temp_path / "test.py").write_text("print('test')") + + runner = CliRunner() + + # Normal output + normal_result = runner.invoke(main, [str(temp_path)]) + assert normal_result.exit_code == 0 + + # Tree-only output + tree_result = runner.invoke(main, [str(temp_path), "--tree-only"]) + assert tree_result.exit_code == 0 + + # Tree-only should be shorter + assert len(tree_result.output) < len(normal_result.output) + + # Normal output should contain file contents and markdown + assert "print('test')" in normal_result.output + assert "# Folder Structure" in normal_result.output + assert "## File:" in normal_result.output + + # Tree-only should not contain file contents or markdown headers + assert "print('test')" not in tree_result.output + assert "# Folder Structure" not in tree_result.output + assert "## File:" not in tree_result.output + + +def test_tree_only_with_output_file(): + """Test that --tree-only works with output file.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + output_file = temp_path / "output.txt" + + # Create test files + (temp_path / "test.py").write_text("print('hello')") + + runner = CliRunner() + result = runner.invoke(main, [str(temp_path), "--tree-only", "-o", str(output_file)]) + assert result.exit_code == 0 + + # Check that output file was created and contains only tree + assert output_file.exists() + content = output_file.read_text() + + # Should contain tree structure + assert temp_path.name in content + assert "test.py" in content + + # Should NOT contain file contents or markdown headers + assert "print('hello')" not in content + assert "# Folder Structure" not in content + assert "## File:" not in content \ No newline at end of file