11import codegen
22from codegen import Codebase
33
4+
45@codegen .function ("freezegun-to-timemachine" )
56def run (codebase : Codebase ):
67 """Convert FreezeGun usage to TimeMachine in test files.
@@ -19,43 +20,43 @@ def run(codebase: Codebase):
1920
2021 # Update imports
2122 for imp in file .imports :
22- if imp .symbol_name and ' freezegun' in imp .source :
23- if imp .name == ' freeze_time' :
23+ if imp .symbol_name and " freezegun" in imp .source :
24+ if imp .name == " freeze_time" :
2425 # required due to Codegen limitations
25- imp .edit (' from time_machine import travel' )
26+ imp .edit (" from time_machine import travel" )
2627 else :
27- imp .set_import_module (' time_machine' )
28+ imp .set_import_module (" time_machine" )
2829
2930 # Find all function calls in the file
3031 for fcall in file .function_calls :
3132 # Skip if not a freeze_time call
32- if ' freeze_time' not in fcall .source :
33+ if " freeze_time" not in fcall .source :
3334 continue
3435
3536 # Get original source and prepare new source
3637 new_source = fcall .source
3738
3839 # Add tick parameter if not present
39- if not fcall .get_arg_by_parameter_name (' tick' ):
40- if new_source .endswith (')' ):
40+ if not fcall .get_arg_by_parameter_name (" tick" ):
41+ if new_source .endswith (")" ):
4142 new_source = new_source [:- 1 ]
42- if not new_source .endswith ('(' ):
43- new_source += ','
44- new_source += ' tick=False)'
43+ if not new_source .endswith ("(" ):
44+ new_source += ","
45+ new_source += " tick=False)"
4546
4647 # Replace freeze_time with travel
47- if '.' in new_source :
48- new_source = new_source .replace (
49- 'freeze_time' , 'travel' ).replace ('freezegun' , 'time_machine' )
48+ if "." in new_source :
49+ new_source = new_source .replace ("freeze_time" , "travel" ).replace ("freezegun" , "time_machine" )
5050 else :
51- new_source = ' travel' + new_source [len (' freeze_time' ) :]
51+ new_source = " travel" + new_source [len (" freeze_time" ) :]
5252
5353 # Make single edit with complete changes
5454 fcall .edit (new_source )
5555
5656 codebase .commit ()
5757 print ("✅ FreezeGun to TimeMachine conversion completed successfully!" )
5858
59+
5960if __name__ == "__main__" :
6061 codebase = Codebase .from_repo ("getmoto/moto" , commit = "786a8ada7ed0c7f9d8b04d49f24596865e4b7901" )
6162 run (codebase )
0 commit comments