Skip to content

Commit

Permalink
Merge pull request #7 from Maralai/feature/levels
Browse files Browse the repository at this point in the history 
Add max depth parameters to directory summarization functions
  • Loading branch information
@Maralai
Maralai authored Feb 1, 2025
2 parents 0b8c3f5 + d8d3ced commit dcc0eb6
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 43 deletions.
54 changes: 42 additions & 12 deletions summarizeGPT/summarizeGPT.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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()

Expand All @@ -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:
Expand Down
119 changes: 88 additions & 31 deletions tests/test_summarize_gpt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand All @@ -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()
Expand All @@ -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()

0 comments on commit dcc0eb6

Please sign in to comment.