|
1 | 1 | #!/usr/bin/env python3 |
2 | 2 |
|
3 | | -import re |
| 3 | +import ast |
| 4 | +from collections.abc import Generator |
4 | 5 | from pathlib import Path |
5 | 6 |
|
6 | | -from redbaron import RedBaron |
| 7 | +from griffe import Module, load |
7 | 8 | from utils import sync_to_async_docstring |
8 | 9 |
|
9 | 10 | # Methods where the async docstring is intentionally different from the sync one |
|
12 | 13 | 'with_custom_http_client', |
13 | 14 | } |
14 | 15 |
|
15 | | -# Get the directory of the source files |
16 | | -clients_path = Path(__file__).parent.resolve() / '../src/apify_client' |
17 | | - |
18 | | -# Go through every Python file in that directory |
19 | | -for client_source_path in clients_path.glob('**/*.py'): |
20 | | - with open(client_source_path, 'r+', encoding='utf-8') as source_file: |
21 | | - # Read the source file and parse the code using Red Baron |
22 | | - red = RedBaron(source_code=source_file.read()) |
23 | | - |
24 | | - # Find all classes which end with "ClientAsync" (there should be at most 1 per file) |
25 | | - async_class = red.find('ClassNode', name=re.compile('.*ClientAsync$')) |
26 | | - |
27 | | - if async_class is None: |
28 | | - # No async client class in this file, nothing to fix |
| 16 | +# Load the apify_client package |
| 17 | +src_path = Path(__file__).parent.resolve() / '../src' |
| 18 | +package = load('apify_client', search_paths=[str(src_path)]) |
| 19 | + |
| 20 | + |
| 21 | +def walk_modules(module: Module) -> Generator[Module]: |
| 22 | + """Recursively yield all modules in the package.""" |
| 23 | + yield module |
| 24 | + for submodule in module.modules.values(): |
| 25 | + yield from walk_modules(submodule) |
| 26 | + |
| 27 | + |
| 28 | +def format_docstring(content: str, indent: str) -> str: |
| 29 | + """Format a docstring with proper indentation and triple quotes.""" |
| 30 | + lines = content.split('\n') |
| 31 | + if len(lines) == 1: |
| 32 | + return f'{indent}"""{lines[0]}"""\n' |
| 33 | + |
| 34 | + result_lines = [f'{indent}"""{lines[0]}'] |
| 35 | + for line in lines[1:]: |
| 36 | + if line.strip(): |
| 37 | + result_lines.append(f'{indent}{line}') |
| 38 | + else: |
| 39 | + result_lines.append('') |
| 40 | + result_lines.append(f'{indent}"""') |
| 41 | + return '\n'.join(result_lines) + '\n' |
| 42 | + |
| 43 | + |
| 44 | +def find_docstring_range(tree: ast.AST, class_name: str, method_name: str) -> tuple[int, int | None, int] | None: |
| 45 | + """Find the line range of a method's docstring using ast. |
| 46 | +
|
| 47 | + Returns (start_line, end_line, col_offset) 1-indexed, or None. |
| 48 | + """ |
| 49 | + for node in ast.walk(tree): |
| 50 | + if isinstance(node, ast.ClassDef) and node.name == class_name: |
| 51 | + for item in node.body: |
| 52 | + if ( |
| 53 | + isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) |
| 54 | + and item.name == method_name |
| 55 | + and item.body |
| 56 | + and isinstance(item.body[0], ast.Expr) |
| 57 | + and isinstance(item.body[0].value, ast.Constant) |
| 58 | + and isinstance(item.body[0].value.value, str) |
| 59 | + ): |
| 60 | + expr = item.body[0] |
| 61 | + return expr.lineno, expr.end_lineno, expr.col_offset |
| 62 | + return None |
| 63 | + |
| 64 | + |
| 65 | +def find_method_body_start(tree: ast.AST, class_name: str, method_name: str) -> tuple[int, int] | None: |
| 66 | + """Find where a method's body starts (for inserting a missing docstring). |
| 67 | +
|
| 68 | + Returns (line_number, col_offset) 1-indexed, or None. |
| 69 | + """ |
| 70 | + for node in ast.walk(tree): |
| 71 | + if isinstance(node, ast.ClassDef) and node.name == class_name: |
| 72 | + for item in node.body: |
| 73 | + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == method_name and item.body: |
| 74 | + first_stmt = item.body[0] |
| 75 | + return first_stmt.lineno, first_stmt.col_offset |
| 76 | + return None |
| 77 | + |
| 78 | + |
| 79 | +# Go through every module in the package |
| 80 | +if not isinstance(package, Module): |
| 81 | + raise TypeError('Expected griffe to load a Module') |
| 82 | +for module in walk_modules(package): |
| 83 | + replacements = [] |
| 84 | + |
| 85 | + for async_class in module.classes.values(): |
| 86 | + if not async_class.name.endswith('ClientAsync'): |
29 | 87 | continue |
30 | 88 |
|
31 | | - # Find the corresponding sync classes (same name, but without -Async) |
32 | | - sync_class = red.find('ClassNode', name=async_class.name.replace('ClientAsync', 'Client')) |
| 89 | + # Find the corresponding sync class (same name, but without Async) |
| 90 | + sync_class = module.classes.get(async_class.name.replace('ClientAsync', 'Client')) |
| 91 | + if not sync_class: |
| 92 | + continue |
33 | 93 |
|
34 | 94 | # Go through all methods in the async class |
35 | | - for async_method in async_class.find_all('DefNode'): |
36 | | - # Find corresponding sync method in the sync class |
37 | | - sync_method = sync_class.find('DefNode', name=async_method.name) |
38 | | - |
| 95 | + for async_method in async_class.functions.values(): |
39 | 96 | # Skip methods with @ignore_docs decorator |
40 | | - if len(async_method.decorators) and str(async_method.decorators[0].value) == 'ignore_docs': |
| 97 | + if any(str(d.value) == 'ignore_docs' for d in async_method.decorators): |
41 | 98 | continue |
42 | 99 |
|
43 | 100 | # Skip methods whose docstrings are intentionally different |
44 | 101 | if async_method.name in SKIPPED_METHODS: |
45 | 102 | continue |
46 | 103 |
|
47 | | - # Skip methods that don't exist in the sync class |
48 | | - if sync_method is None: |
| 104 | + # Find corresponding sync method in the sync class |
| 105 | + sync_method = sync_class.functions.get(async_method.name) |
| 106 | + if not sync_method or not sync_method.docstring: |
49 | 107 | continue |
50 | 108 |
|
51 | | - # If the sync method has a docstring, copy it to the async method (with adjustments) |
52 | | - if isinstance(sync_method.value[0].value, str): |
53 | | - sync_docstring = sync_method.value[0].value |
54 | | - async_docstring = async_method.value[0].value |
55 | | - |
56 | | - correct_async_docstring = sync_to_async_docstring(sync_docstring) |
57 | | - if async_docstring == correct_async_docstring: |
| 109 | + correct_docstring = sync_to_async_docstring(sync_method.docstring.value) |
| 110 | + |
| 111 | + if not async_method.docstring: |
| 112 | + print(f'Fixing missing docstring for "{async_class.name}.{async_method.name}"...') |
| 113 | + replacements.append((async_class.name, async_method.name, correct_docstring, False)) |
| 114 | + elif async_method.docstring.value != correct_docstring: |
| 115 | + replacements.append((async_class.name, async_method.name, correct_docstring, True)) |
| 116 | + |
| 117 | + if not replacements: |
| 118 | + continue |
| 119 | + |
| 120 | + # Read the source file and parse with ast for precise locations |
| 121 | + filepath = module.filepath |
| 122 | + if not isinstance(filepath, Path): |
| 123 | + continue |
| 124 | + source = filepath.read_text(encoding='utf-8') |
| 125 | + source_lines = source.splitlines(keepends=True) |
| 126 | + tree = ast.parse(source) |
| 127 | + |
| 128 | + # Collect replacement operations with line numbers |
| 129 | + ops = [] |
| 130 | + for class_name, method_name, correct_docstring, has_existing in replacements: |
| 131 | + if has_existing: |
| 132 | + result = find_docstring_range(tree, class_name, method_name) |
| 133 | + if result: |
| 134 | + start_line, end_line, col_offset = result |
| 135 | + if end_line is None: |
58 | 136 | continue |
59 | | - |
60 | | - # Work around a bug in Red Baron, which indents docstrings too much when you insert them, |
61 | | - # so we have to un-indent it one level first. |
62 | | - correct_async_docstring = re.sub('^ ', '', correct_async_docstring, flags=re.MULTILINE) |
63 | | - |
64 | | - if not isinstance(async_docstring, str): |
65 | | - print(f'Fixing missing docstring for "{async_class.name}.{async_method.name}"...') |
66 | | - async_method.value.insert(0, correct_async_docstring) |
67 | | - else: |
68 | | - async_method.value[0] = correct_async_docstring |
69 | | - |
70 | | - updated_source_code = red.dumps() |
71 | | - |
72 | | - # Work around a bug in Red Baron, which adds indents to docstrings when you insert them (including empty lines), |
73 | | - # so we have to remove the extra whitespace |
74 | | - updated_source_code = re.sub('^ $', '', updated_source_code, flags=re.MULTILINE) |
75 | | - |
76 | | - # Work around a bug in Red Baron, which indents `except` and `finally` statements wrong |
77 | | - # so we have to add some extra whitespace |
78 | | - updated_source_code = re.sub('^except', ' except', updated_source_code, flags=re.MULTILINE) |
79 | | - updated_source_code = re.sub('^ except', ' except', updated_source_code, flags=re.MULTILINE) |
80 | | - updated_source_code = re.sub('^finally', ' finally', updated_source_code, flags=re.MULTILINE) |
81 | | - updated_source_code = re.sub('^ finally', ' finally', updated_source_code, flags=re.MULTILINE) |
82 | | - |
83 | | - # Work around a bug in Red Baron, which sometimes adds an extra new line to the end of a file |
84 | | - updated_source_code = updated_source_code.rstrip() + '\n' |
85 | | - |
86 | | - # Save the updated source code back to the file |
87 | | - source_file.seek(0) |
88 | | - source_file.write(updated_source_code) |
89 | | - source_file.truncate() |
| 137 | + indent = ' ' * col_offset |
| 138 | + formatted = format_docstring(correct_docstring, indent) |
| 139 | + ops.append(('replace', start_line, end_line, formatted)) |
| 140 | + else: |
| 141 | + result = find_method_body_start(tree, class_name, method_name) |
| 142 | + if result: |
| 143 | + insert_line, col_offset = result |
| 144 | + indent = ' ' * col_offset |
| 145 | + formatted = format_docstring(correct_docstring, indent) |
| 146 | + ops.append(('insert', insert_line, None, formatted)) |
| 147 | + |
| 148 | + # Sort by start line descending (process bottom-up to preserve line numbers) |
| 149 | + ops.sort(key=lambda x: x[1], reverse=True) |
| 150 | + |
| 151 | + for op_type, start_line, end_line, formatted in ops: |
| 152 | + formatted_lines = formatted.splitlines(keepends=True) |
| 153 | + if op_type == 'replace': |
| 154 | + source_lines[start_line - 1 : end_line] = formatted_lines |
| 155 | + elif op_type == 'insert': |
| 156 | + source_lines[start_line - 1 : start_line - 1] = formatted_lines |
| 157 | + |
| 158 | + # Save the updated source code back to the file |
| 159 | + filepath.write_text(''.join(source_lines), encoding='utf-8') |
0 commit comments