|
| 1 | +import codegen |
1 | 2 | from codegen import Codebase |
2 | | -codebase = Codebase.from_repo("getmoto/moto", commit="786a8ada7ed0c7f9d8b04d49f24596865e4b7901") |
3 | | - |
4 | | -print("🚀 Starting FreezeGun to TimeMachine conversion...") |
5 | | - |
6 | | -for file in codebase.files: |
7 | | - if "tests" not in file.filepath: |
8 | | - continue |
9 | | - print(f"📝 Processing: {file.filepath}") |
10 | | - # Update imports |
11 | | - for imp in file.imports: |
12 | | - if imp.symbol_name and 'freezegun' in imp.source: |
13 | | - if imp.name == 'freeze_time': |
14 | | - # required due to Codegen limitations |
15 | | - imp.edit('from time_machine import travel') |
16 | | - else: |
17 | | - imp.set_import_module('time_machine') |
18 | | - # Find all function calls in the file |
19 | | - for fcall in file.function_calls: |
20 | | - # Skip if not a freeze_time call |
21 | | - if 'freeze_time' not in fcall.source: |
| 3 | + |
| 4 | +@codegen.function("freezegun-to-timemachine") |
| 5 | +def run(codebase: Codebase): |
| 6 | + """Convert FreezeGun usage to TimeMachine in test files. |
| 7 | +
|
| 8 | + This script: |
| 9 | + 1. Identifies test files using FreezeGun. |
| 10 | + 2. Updates imports from FreezeGun to TimeMachine. |
| 11 | + 3. Modifies function calls to include necessary parameters. |
| 12 | + """ |
| 13 | + print("🚀 Starting FreezeGun to TimeMachine conversion...") |
| 14 | + |
| 15 | + for file in codebase.files: |
| 16 | + if "tests" not in file.filepath: |
22 | 17 | continue |
23 | | - # Get original source and prepare new source |
24 | | - new_source = fcall.source |
25 | | - # Add tick parameter if not present |
26 | | - if not fcall.get_arg_by_parameter_name('tick'): |
27 | | - if new_source.endswith(')'): |
28 | | - new_source = new_source[:-1] |
29 | | - if not new_source.endswith('('): |
30 | | - new_source += ',' |
31 | | - new_source += ' tick=False)' |
32 | | - # Replace freeze_time with travel |
33 | | - if '.' in new_source: |
34 | | - new_source = new_source.replace( |
35 | | - 'freeze_time', 'travel').replace('freezegun', 'time_machine') |
36 | | - else: |
37 | | - new_source = 'travel' + new_source[len('freeze_time'):] |
38 | | - # Make single edit with complete changes |
39 | | - fcall.edit(new_source) |
40 | | -codebase.commit() |
41 | | - |
42 | | -print("✅ FreezeGun to TimeMachine conversion completed successfully!") |
| 18 | + print(f"📝 Processing: {file.filepath}") |
| 19 | + |
| 20 | + # Update imports |
| 21 | + for imp in file.imports: |
| 22 | + if imp.symbol_name and 'freezegun' in imp.source: |
| 23 | + if imp.name == 'freeze_time': |
| 24 | + # required due to Codegen limitations |
| 25 | + imp.edit('from time_machine import travel') |
| 26 | + else: |
| 27 | + imp.set_import_module('time_machine') |
| 28 | + |
| 29 | + # Find all function calls in the file |
| 30 | + for fcall in file.function_calls: |
| 31 | + # Skip if not a freeze_time call |
| 32 | + if 'freeze_time' not in fcall.source: |
| 33 | + continue |
| 34 | + |
| 35 | + # Get original source and prepare new source |
| 36 | + new_source = fcall.source |
| 37 | + |
| 38 | + # Add tick parameter if not present |
| 39 | + if not fcall.get_arg_by_parameter_name('tick'): |
| 40 | + if new_source.endswith(')'): |
| 41 | + new_source = new_source[:-1] |
| 42 | + if not new_source.endswith('('): |
| 43 | + new_source += ',' |
| 44 | + new_source += ' tick=False)' |
| 45 | + |
| 46 | + # Replace freeze_time with travel |
| 47 | + if '.' in new_source: |
| 48 | + new_source = new_source.replace( |
| 49 | + 'freeze_time', 'travel').replace('freezegun', 'time_machine') |
| 50 | + else: |
| 51 | + new_source = 'travel' + new_source[len('freeze_time'):] |
| 52 | + |
| 53 | + # Make single edit with complete changes |
| 54 | + fcall.edit(new_source) |
| 55 | + |
| 56 | + codebase.commit() |
| 57 | + print("✅ FreezeGun to TimeMachine conversion completed successfully!") |
| 58 | + |
| 59 | +if __name__ == "__main__": |
| 60 | + codebase = Codebase.from_repo("getmoto/moto", commit="786a8ada7ed0c7f9d8b04d49f24596865e4b7901") |
| 61 | + run(codebase) |
0 commit comments