diff --git a/summarizeGPT/summarizeGPT.py b/summarizeGPT/summarizeGPT.py index ed5b890..998fba4 100644 --- a/summarizeGPT/summarizeGPT.py +++ b/summarizeGPT/summarizeGPT.py @@ -25,16 +25,21 @@ def setup_logging(verbose): console_handler.setFormatter(formatter) logger.addHandler(console_handler) -def summarize_directory(directory, gitignore_file=None, include_exts=None, exclude_exts=None, show_docker=False, show_only_docker=False, max_lines=None): +def summarize_directory(directory, gitignore_file=None, include_exts=None, + exclude_exts=None, show_docker=False, show_only_docker=False, + max_lines=None, tree_depth=None, file_depth=None): directory = directory.replace("\\", "/") prompt_md = f"# Summary of directory: {directory}\n\n" - tree_view = get_tree_view(directory, gitignore_file=gitignore_file) + tree_view = get_tree_view(directory, gitignore_file=gitignore_file, max_depth=tree_depth) prompt_md += "```\n" + tree_view + "\n```\n\n" - file_contents = get_file_contents(directory, gitignore_file, include_exts, exclude_exts, show_docker=show_docker, show_only_docker=show_only_docker, max_lines=max_lines) + file_contents = get_file_contents(directory, gitignore_file, include_exts, + exclude_exts, show_docker=show_docker, + show_only_docker=show_only_docker, + max_lines=max_lines, max_depth=file_depth) prompt_md += file_contents return prompt_md -def get_tree_view(directory, gitignore_file=None): +def get_tree_view(directory, gitignore_file=None, max_depth=None): tree_view = "" if gitignore_file: gitignore = gitignore_parser.parse_gitignore(gitignore_file) @@ -44,20 +49,30 @@ def get_tree_view(directory, gitignore_file=None): for root, dirs, files in os.walk(directory): if '.git' in root: continue + + level = root.replace(directory, '').count(os.sep) + 1 # +1 because root is level 1 + + # Skip if we've exceeded the maximum depth + if max_depth is not None and level > max_depth: + dirs[:] = [] # Clear dirs to prevent further recursion + continue + if gitignore: dirs[:] = [d for d in dirs if not gitignore(os.path.join(root, d))] files = [f for f in files if not gitignore(os.path.join(root, f))] - level = root.replace(directory, '').count(os.sep) - indent = ' ' * 4 * (level) + + indent = ' ' * 4 * (level - 1) # Adjust indent because level starts at 1 tree_view += f"{indent}{os.path.basename(root)}/\n" - sub_indent = ' ' * 4 * (level + 1) + sub_indent = ' ' * 4 * level for file in files: if file == output_file: continue tree_view += f"{sub_indent}{file}\n" return tree_view -def get_file_contents(directory, gitignore_file=None, include_exts=None, exclude_exts=None, show_docker=False, show_only_docker=False, max_lines=None): +def get_file_contents(directory, gitignore_file=None, include_exts=None, + exclude_exts=None, show_docker=False, show_only_docker=False, + max_lines=None, max_depth=None): file_contents = "" excluded_files = ['docker', 'Dockerfile'] if gitignore_file: @@ -68,6 +83,10 @@ def get_file_contents(directory, gitignore_file=None, include_exts=None, exclude for root, _, files in os.walk(directory): if '.git' in root: continue + level = root.replace(directory, '').count(os.sep) + 1 # +1 because root is level 1 + # Skip if we've exceeded the maximum depth + if max_depth is not None and level > max_depth: + continue if gitignore: files = [f for f in files if not gitignore(os.path.join(root, f))] for file in files: @@ -138,6 +157,12 @@ def main(): default='cl100k_base', help='Tiktoken encoding to use for token counting (default: cl100k_base)') parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose output') + parser.add_argument('-L', '--max-depth', type=int, default=None, + help='Maximum directory depth to traverse for both tree and files (root=1)') + parser.add_argument('-Lt', '--tree-depth', type=int, default=None, + help='Maximum directory depth for tree view (root=1)') + parser.add_argument('-Lf', '--file-depth', type=int, default=None, + help='Maximum directory depth for file contents (root=1)') args = parser.parse_args() @@ -151,10 +176,15 @@ def main(): include_exts = [".{}".format(ext.lower()) for ext in args.include.split(',')] if args.include else None exclude_exts = [".{}".format(ext.lower()) for ext in args.exclude.split(',')] if args.exclude else None - prompt_md = summarize_directory(args.directory, args.gitignore, include_exts, - exclude_exts, show_docker=args.show_docker, - show_only_docker=args.show_only_docker, - max_lines=args.max_lines) + tree_depth = args.tree_depth if args.tree_depth is not None else args.max_depth + file_depth = args.file_depth if args.file_depth is not None else args.max_depth + + prompt_md = summarize_directory(args.directory, args.gitignore, include_exts, + exclude_exts, show_docker=args.show_docker, + show_only_docker=args.show_only_docker, + max_lines=args.max_lines, + tree_depth=tree_depth, + file_depth=file_depth) prompt_file = os.path.join(args.directory, output_file) try: diff --git a/tests/test_summarize_gpt.py b/tests/test_summarize_gpt.py index be81e20..76acbd6 100644 --- a/tests/test_summarize_gpt.py +++ b/tests/test_summarize_gpt.py @@ -3,7 +3,7 @@ import sys import logging import tempfile -from unittest.mock import patch, MagicMock +from unittest.mock import patch from summarizeGPT.summarizeGPT import ( summarize_directory, get_tree_view, @@ -38,7 +38,7 @@ def test_summarize_directory_basic(self): """Test basic directory summarization""" result = summarize_directory(self.test_dir) self.assertIsInstance(result, str) - self.assertIn(self.test_dir.replace("\\", "/"), result) # Ensure proper path formatting + self.assertIn(os.path.basename(self.test_dir), result) self.assertIn("test.txt", result) def test_get_tree_view_basic(self): @@ -73,20 +73,21 @@ def test_logging_setup_non_verbose(self): @patch('sys.exit') def test_docker_flags_conflict(self, mock_exit, mock_args): """Test handling of conflicting docker flags""" - # Mock the parsed arguments - mock_args.return_value = MagicMock( - directory=self.test_dir, - gitignore=None, - include=None, - exclude=None, - show_docker=True, - show_only_docker=True, - max_lines=None, - encoding='cl100k_base', - verbose=False - ) - - # Capture log messages directly + mock_args.return_value = type('Args', (), { + 'directory': self.test_dir, + 'gitignore': None, + 'include': None, + 'exclude': None, + 'show_docker': True, + 'show_only_docker': True, + 'max_lines': None, + 'encoding': 'cl100k_base', + 'verbose': False, + 'tree_depth': None, + 'file_depth': None, + 'max_depth': None # Added missing attribute + })() + with patch('logging.Logger.error') as mock_logger: main() mock_logger.assert_called_once_with( @@ -98,27 +99,27 @@ def test_docker_flags_conflict(self, mock_exit, mock_args): @patch('sys.exit') def test_file_write_error(self, mock_exit, mock_args): """Test handling of file write errors""" - # Mock the parsed arguments - mock_args.return_value = MagicMock( - directory=self.test_dir, - gitignore=None, - include=None, - exclude=None, - show_docker=False, - show_only_docker=False, - max_lines=None, - encoding='cl100k_base', - verbose=False - ) - - # Create a mock that only affects the output file + mock_args.return_value = type('Args', (), { + 'directory': self.test_dir, + 'gitignore': None, + 'include': None, + 'exclude': None, + 'show_docker': False, + 'show_only_docker': False, + 'max_lines': None, + 'encoding': 'cl100k_base', + 'verbose': False, + 'tree_depth': None, + 'file_depth': None, + 'max_depth': None # Added missing attribute + })() + original_open = open def mock_open_wrapper(*args, **kwargs): if args and isinstance(args[0], str) and args[0].endswith(output_file): raise IOError("Permission denied") return original_open(*args, **kwargs) - # Capture log messages directly with patch('builtins.open', side_effect=mock_open_wrapper): with patch('logging.Logger.error') as mock_logger: main() @@ -127,5 +128,61 @@ def mock_open_wrapper(*args, **kwargs): ) mock_exit.assert_called_once_with(1) + def test_depth_limiting(self): + """Test directory depth limiting functionality""" + # Create a nested directory structure + nested_dir = os.path.join(self.test_dir, "level1", "level2") + os.makedirs(nested_dir) + + # Create files at different levels + with open(os.path.join(self.test_dir, "root.txt"), "w") as f: + f.write("root") + with open(os.path.join(self.test_dir, "level1", "level1.txt"), "w") as f: + f.write("level1") + with open(os.path.join(nested_dir, "level2.txt"), "w") as f: + f.write("level2") + + # Test with tree_depth=2 to ensure we see one level of nesting + result = summarize_directory(self.test_dir, tree_depth=2, file_depth=2) + self.assertIn("root.txt", result) + self.assertIn("level1", result) + self.assertIn("level1.txt", result) + self.assertIn("root", result) # File content + self.assertIn("level1", result) # File content + + def test_depth_edge_cases(self): + """Test edge cases for depth parameters""" + nested_dir = os.path.join(self.test_dir, "level1", "level2") + os.makedirs(nested_dir) + + with open(os.path.join(self.test_dir, "root.txt"), "w") as f: + f.write("root") + + # Test with depth=0 + result = summarize_directory(self.test_dir, tree_depth=0, file_depth=0) + self.assertIn(os.path.basename(self.test_dir), result) + + # Test with unlimited depth + result = summarize_directory(self.test_dir, tree_depth=None, file_depth=None) + self.assertIn("root.txt", result) + + def test_separate_tree_and_file_depths(self): + """Test different depth limits for tree view and file contents""" + with tempfile.TemporaryDirectory() as tmp_dir: + # Create a nested directory structure + os.makedirs(os.path.join(tmp_dir, "level1/level2")) + with open(os.path.join(tmp_dir, "level1/file1.txt"), "w") as f: + f.write("level1 content") + with open(os.path.join(tmp_dir, "level1/level2/file2.txt"), "w") as f: + f.write("level2 content") + + result = summarize_directory(tmp_dir, tree_depth=2, file_depth=2) + + # Check tree view shows both levels and file contents + self.assertIn("level1", result) + self.assertIn("file1.txt", result) + self.assertIn("level1 content", result) + self.assertNotIn("level2 content", result) + if __name__ == '__main__': unittest.main() \ No newline at end of file