[READ-ONLY] Mirror of https://github.com/mrgnw/scripts.
17 kB
447 lines
1# /// script
2# requires-python = ">=3.12"
3# dependencies = [
4# "requests",
5# "tqdm",
6# "argparse",
7# "aiohttp",
8# "asyncio",
9# "rich"
10# ]
11# ///
12import os
13import requests
14import json
15import argparse
16import asyncio
17import aiohttp
18from pathlib import Path
19import sys
20from fnmatch import fnmatch
21from rich.console import Console
22from rich.panel import Panel
23from rich.syntax import Syntax
24from rich.markdown import Markdown
25from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn
26from rich.tree import Tree
27from rich import print as rprint
28from tqdm.asyncio import tqdm as async_tqdm
29
30# Configuration
31DEFAULT_BRANCH = 'main'
32GITHUB_TOKEN = os.getenv('GITHUB_TOKEN') # Set your GitHub token in env variables
33OUTPUT_FILE = 'repo_files.md' # Updated default file name
34HEADERS = (
35 {'Authorization': f'bearer {GITHUB_TOKEN}', 'Content-Type': 'application/json'} if GITHUB_TOKEN else {}
36)
37
38# Initialize rich console
39console = Console()
40
41# File patterns to fetch
42INCLUDE_PATTERNS = [
43 'package.json',
44 '**/*.config.js',
45 '**/*.config.ts',
46 'src/routes/**/*.svelte',
47 'src/routes/**/*.js',
48 'src/routes/**/*.ts',
49 # Adding more general patterns for all JavaScript, TypeScript, and Svelte files
50 '**/*.js',
51 '**/*.ts',
52 '**/*.svelte',
53 '**/*.json'
54]
55
56
57def graphql_query(query, variables=None):
58 url = 'https://api.github.com/graphql'
59 payload = {'query': query, 'variables': variables or {}}
60 response = requests.post(url, headers=HEADERS, json=payload)
61 response.raise_for_status()
62 return response.json()
63
64
65async def async_graphql_query(session, query, variables=None):
66 url = 'https://api.github.com/graphql'
67 payload = {'query': query, 'variables': variables or {}}
68 async with session.post(url, json=payload, headers=HEADERS) as response:
69 response.raise_for_status()
70 return await response.json()
71
72
73def get_all_files(repo, branch='main'):
74 """Get all files in the repository using a single GraphQL query with Git tree."""
75 owner, name = repo.split('/')
76
77 # First, try to get the default branch
78 with console.status(f"[bold green]Checking repository {repo}...", spinner="dots"):
79 query = """
80 query($owner: String!, $name: String!) {
81 repository(owner: $owner, name: $name) {
82 defaultBranchRef {
83 name
84 }
85 }
86 }
87 """
88
89 data = graphql_query(query, {'owner': owner, 'name': name})
90
91 if 'data' not in data or data['data'] is None or not data['data'].get('repository'):
92 raise ValueError(f"Error fetching repository data. API response: {data}")
93
94 repository = data['data'].get('repository')
95
96 # Use default branch if main is not available
97 if branch == 'main' and repository.get('defaultBranchRef'):
98 default_branch = repository['defaultBranchRef']['name']
99 if default_branch != branch:
100 console.print(f"[yellow]Using default branch:[/] [bold cyan]{default_branch}[/]")
101 branch = default_branch
102
103 # Use Git Trees API directly - more reliable for full repository listing
104 rest_api_url = f"https://api.github.com/repos/{owner}/{name}/git/trees/{branch}?recursive=1"
105 console.print(f"[bold]Fetching file list from:[/] [blue]{rest_api_url}[/]")
106
107 with console.status("[bold green]Downloading repository structure...", spinner="dots"):
108 response = requests.get(rest_api_url, headers=HEADERS)
109 response.raise_for_status()
110
111 all_files = []
112 tree_data = response.json()
113
114 if tree_data.get('truncated', False):
115 console.print("[bold yellow]Warning:[/] Repository tree is too large and was truncated. Some files may be missing.")
116
117 # Print total number of files found
118 total_files = len(tree_data.get('tree', []))
119 console.print(f"Total files in repository: [bold cyan]{total_files}[/]")
120
121 # Process all blob entries (files)
122 for item in tree_data.get('tree', []):
123 if item['type'] == 'blob': # Only include files, not directories
124 all_files.append({
125 'path': item['path'],
126 'type': item['type']
127 })
128
129 console.print(f"Total blob entries: [bold cyan]{len(all_files)}[/]")
130
131 # Filter files based on patterns
132 matching_files = []
133 for file in all_files:
134 for pattern in INCLUDE_PATTERNS:
135 if fnmatch(file['path'], pattern):
136 matching_files.append(file)
137 break # No need to check other patterns
138
139 # Display matching files as a tree
140 file_tree = Tree("[bold]Matching files found:[/]")
141
142 # Group files by directory for better visualization
143 file_paths = [file['path'] for file in matching_files]
144 file_paths.sort()
145
146 # Build directory tree
147 path_dict = {}
148 for path in file_paths:
149 parts = path.split('/')
150 current_dict = path_dict
151 for i, part in enumerate(parts):
152 if i == len(parts) - 1: # If it's the last part (file)
153 if 'files' not in current_dict:
154 current_dict['files'] = []
155 current_dict['files'].append(part)
156 else: # If it's a directory
157 if part not in current_dict:
158 current_dict[part] = {}
159 current_dict = current_dict[part]
160
161 # Function to build tree display
162 def build_tree(tree_node, path_dict, is_root=False):
163 # Add directories first
164 for key, value in sorted(path_dict.items()):
165 if key != 'files':
166 dir_node = tree_node.add(f"[bold blue]{key}/[/]")
167 build_tree(dir_node, value)
168
169 # Then add files
170 if 'files' in path_dict:
171 for file in sorted(path_dict['files']):
172 ext = file.split('.')[-1] if '.' in file else ''
173 if ext in ['js', 'ts']:
174 tree_node.add(f"[yellow]{file}[/]")
175 elif ext == 'json':
176 tree_node.add(f"[green]{file}[/]")
177 elif ext == 'svelte':
178 tree_node.add(f"[orange1]{file}[/]")
179 else:
180 tree_node.add(f"[white]{file}[/]")
181
182 # Build and display the tree
183 build_tree(file_tree, path_dict, True)
184 console.print(file_tree)
185
186 console.print(f"\nFound [bold green]{len(matching_files)}[/] matching files. Fetching content...")
187
188 # Generate a markdown representation of the tree for the output file
189 # Use GitHub-flavored markdown tree format
190 md_tree_content = []
191
192 def build_md_tree(path_dict, prefix=""):
193 items = sorted([(k, v) for k, v in path_dict.items() if k != "files"])
194 file_items = path_dict.get("files", [])
195
196 # Process all items except the last one with ├── prefix
197 for i, (key, value) in enumerate(items[:-1] if items else []):
198 md_tree_content.append(f"{prefix}├── 📁 {key}/")
199 new_prefix = f"{prefix}│ "
200 build_md_tree(value, new_prefix)
201
202 # Process the last directory item with └── prefix
203 if items:
204 key, value = items[-1]
205 md_tree_content.append(f"{prefix}└── 📁 {key}/")
206 new_prefix = f"{prefix} "
207 build_md_tree(value, new_prefix)
208
209 # Process files with appropriate prefixes
210 sorted_files = sorted(file_items)
211 for i, file in enumerate(sorted_files):
212 ext = file.split('.')[-1] if '.' in file else ''
213 icon = "📄"
214 if ext in ['js', 'ts']:
215 icon = "🟨"
216 elif ext == 'json':
217 icon = "🔧"
218 elif ext == 'svelte':
219 icon = "🔥"
220 elif ext in ['md', 'markdown']:
221 icon = "📝"
222
223 # Use different prefix for last item
224 if i == len(sorted_files) - 1:
225 md_tree_content.append(f"{prefix}└── {icon} {file}")
226 else:
227 md_tree_content.append(f"{prefix}├── {icon} {file}")
228
229 build_md_tree(path_dict)
230 md_tree = "\n".join(md_tree_content)
231
232 return matching_files, md_tree
233
234
235async def fetch_content_async(repo, files, branch):
236 """Fetch file contents asynchronously with individual progress bars for each file."""
237 owner, name = repo.split('/')
238
239 # Set up a rich progress display with multiple tasks
240 results = []
241
242 # Create multi-progress display
243 with Progress(
244 TextColumn("[bold blue]{task.description}"),
245 BarColumn(),
246 TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
247 TextColumn("•"),
248 TimeRemainingColumn(),
249 console=console,
250 ) as progress:
251 # Create a task for each file
252 tasks = {}
253 for i, file in enumerate(files):
254 # Truncate file name if it's too long for display
255 display_name = file['path']
256 if len(display_name) > 40:
257 display_name = "..." + display_name[-37:]
258
259 tasks[file['path']] = progress.add_task(
260 f"[cyan]{display_name}[/]",
261 total=1.0,
262 completed=0.0
263 )
264
265 # Process files with proper progress tracking
266 async with aiohttp.ClientSession() as session:
267 # Process in batches to avoid hitting rate limits
268 batch_size = 5 # Adjust based on API rate limits
269 for i in range(0, len(files), batch_size):
270 batch = files[i:i+batch_size]
271
272 # Create tasks for each file in the batch
273 batch_tasks = []
274 for file in batch:
275 task = asyncio.create_task(
276 fetch_file_with_progress(
277 session, repo, file['path'], branch,
278 progress, tasks[file['path']]
279 )
280 )
281 batch_tasks.append(task)
282
283 # Wait for this batch to complete
284 batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
285
286 # Process results from this batch
287 for j, result in enumerate(batch_results):
288 file_path = batch[j]['path']
289 if isinstance(result, Exception):
290 console.print(f"[red]Error fetching {file_path}: {result}[/]")
291 # Ensure progress bar is completed even on error
292 progress.update(tasks[file_path], completed=1.0)
293 else:
294 results.append({
295 'path': file_path,
296 'content': result
297 })
298
299 return results
300
301
302async def fetch_file_with_progress(session, repo, path, branch, progress, task_id):
303 """Fetch a single file with progress tracking."""
304 # First update to show we're starting
305 progress.update(task_id, completed=0.1)
306
307 try:
308 # Fetch the file content
309 content = await fetch_file_content_async(session, repo, path, branch)
310 # Mark as complete
311 progress.update(task_id, completed=1.0)
312 return content
313 except Exception as e:
314 progress.update(task_id, completed=1.0)
315 raise e
316
317
318async def fetch_file_content_async(session, repo, path, branch='main'):
319 """Fetch content of a specific file asynchronously."""
320 query = """
321 query($owner: String!, $name: String!, $expression: String!) {
322 repository(owner: $owner, name: $name) {
323 object(expression: $expression) {
324 ... on Blob {
325 text
326 }
327 }
328 }
329 }
330 """
331 owner, name = repo.split('/')
332 expression = f"{branch}:{path}"
333 data = await async_graphql_query(session, query, {'owner': owner, 'name': name, 'expression': expression})
334
335 # Check if file exists
336 if ('data' not in data or
337 not data['data'].get('repository') or
338 not data['data']['repository'].get('object')):
339 raise ValueError(f"File '{path}' not found")
340
341 return data['data']['repository']['object'].get('text', '')
342
343
344def parse_args():
345 parser = argparse.ArgumentParser(description='Generate a structure document from GitHub repository files')
346 parser.add_argument('repo', nargs='?', help='GitHub repository in format owner/repo')
347 parser.add_argument('output', nargs='?', default=OUTPUT_FILE, help='Output markdown file path')
348 parser.add_argument('--repo', dest='repo_option', help='GitHub repository in format owner/repo')
349 parser.add_argument('--branch', default='main', help='Repository branch (default: main)')
350 parser.add_argument('--patterns', nargs='+', help='File patterns to include')
351 parser.add_argument('--all', action='store_true', help='Include all js/ts/svelte/json files')
352 parser.add_argument('--preview', action='store_true', help='Preview the markdown in terminal')
353 args = parser.parse_args()
354
355 # Use the positional argument if provided, otherwise use the named argument
356 if not args.repo and not args.repo_option:
357 parser.error("Repository name is required. Provide it as a positional argument or with --repo")
358
359 # Positional argument takes precedence
360 final_args = args
361 final_args.repo = args.repo or args.repo_option
362 delattr(final_args, 'repo_option')
363
364 return final_args
365
366
367async def async_main():
368 args = parse_args()
369
370 if not args.repo or '/' not in args.repo:
371 console.print("[bold red]Error:[/] Please provide a valid repository in the format 'owner/repo'")
372 sys.exit(1)
373
374 # Show title
375 console.print(Panel.fit(
376 f"[bold cyan]GitHub Repository Explorer[/]\n[green]{args.repo}[/] (branch: [yellow]{args.branch}[/])",
377 border_style="cyan"
378 ))
379
380 global INCLUDE_PATTERNS
381 if args.patterns:
382 INCLUDE_PATTERNS = args.patterns
383 elif args.all:
384 # Use simplified patterns if --all flag is provided
385 INCLUDE_PATTERNS = [
386 '**/*.js',
387 '**/*.ts',
388 '**/*.svelte',
389 '**/*.json'
390 ]
391
392 output_file = args.output
393
394 try:
395 console.print(f"[bold]Fetching repository structure for [green]{args.repo}[/] (branch: [yellow]{args.branch}[/])...")
396 matching_files, md_tree = get_all_files(args.repo, args.branch)
397
398 if not matching_files:
399 console.print("[yellow]No files matched the specified patterns.[/]")
400 return
401
402 # No need to repeat the file count here - removed redundant message
403
404 # Process files in sorted order for consistent output
405 matching_files.sort(key=lambda f: f['path'])
406
407 # Fetch file contents asynchronously
408 file_contents = await fetch_content_async(args.repo, matching_files, args.branch)
409
410 # Check if we actually got content for all files - simplified message
411 console.print(f"Successfully fetched [bold green]{len(file_contents)}/{len(matching_files)}[/] files")
412
413 # Extract repository name for the header
414 repo_name = args.repo.split('/')[1]
415
416 # Generate markdown with tree structure at the top - more GitHub-friendly format
417 md_content = f'# {repo_name} Project Structure\n\n'
418 md_content += "## File Tree\n\n"
419 md_content += md_tree
420 md_content += "\n\n## File Contents\n\n"
421
422 for file_data in file_contents:
423 file_path = file_data['path']
424 content = file_data['content']
425 # Add language-specific syntax highlighting
426 extension = file_path.split('.')[-1] if '.' in file_path else ''
427 md_content += f"### {file_path}\n\n```{extension}\n{content}\n```\n\n"
428
429 Path(args.output).write_text(md_content, encoding='utf-8')
430 console.print(f"[bold green]Generated:[/] {args.output}")
431
432 if args.preview:
433 console.print("\n[bold]Preview of the generated document:[/]")
434 md = Markdown(md_content[:1000] + "...\n\n(Output truncated for preview)")
435 console.print(md)
436
437 except Exception as e:
438 console.print(f"[bold red]Error:[/] {e}")
439 sys.exit(1)
440
441
442def main():
443 asyncio.run(async_main())
444
445
446if __name__ == '__main__':
447 main()