⚡️ Parse the module once instead of _codemod_ times (#22)

pydantic · Jun 20, 2023 · 9546c33 · 9546c33
1 parent 3097a3c
commit 9546c33
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/bump_pydantic/main.py b/bump_pydantic/main.py
@@ -105,27 +105,26 @@ def run_codemods(
         code = fp.read()
         fp.seek(0)
 
-        input_code = str(code)
+        input_tree = cst.parse_module(code)
 
         for codemod in codemods:
             transformer = codemod(context=context)
 
-            input_tree = cst.parse_module(input_code)
             output_tree = transformer.transform_module(input_tree)
+            input_tree = output_tree
 
-            input_code = output_tree.code
-
-        if code != input_code:
+        output_code = input_tree.code
+        if code != output_code:
             if diff:
                 lines = difflib.unified_diff(
                     code.splitlines(keepends=True),
-                    input_code.splitlines(keepends=True),
+                    output_code.splitlines(keepends=True),
                     fromfile=filename,
                     tofile=filename,
                 )
                 return list(lines)
             else:
-                fp.write(input_code)
+                fp.write(output_code)
                 fp.truncate()
 
     return None