diff --git a/scripts/insert_venom_pragma.py b/scripts/insert_venom_pragma.py
index 4f700bba..ae74b59d 100644
--- a/scripts/insert_venom_pragma.py
+++ b/scripts/insert_venom_pragma.py
@@ -1,29 +1,24 @@
 import os
 
 
-def add_venom_pragma(filepath):
-    with open(filepath, "r") as file:
-        lines = file.readlines()
-
-    # Insert `pragma experimental-codegen` in the second line to activate the `venom` backend.
-    if len(lines) >= 2:
+def insert_venom_pragma(filepath):
+    with open(filepath, "r+") as f:
+        lines = f.readlines()
+        if len(lines) < 2:
+            lines.append("")
+        # Insert `pragma experimental-codegen` on the second line to activate the `venom` backend.
         lines.insert(1, "# pragma experimental-codegen\n")
-    else:
-        lines.append("# pragma experimental-codegen\n")
-
-    # Write the modified lines back to the file.
-    with open(filepath, "w") as file:
-        file.writelines(lines)
+        # Move the file pointer back to the beginning of the file.
+        f.seek(0)
+        f.writelines(lines)
 
 
 def process_directory(directory):
     for root, _, files in os.walk(directory):
         for file in files:
-            if file.endswith(".vy") or file.endswith(".vyi"):
-                filepath = os.path.join(root, file)
-                add_venom_pragma(filepath)
+            if file.endswith((".vy", ".vyi")):
+                insert_venom_pragma(os.path.join(root, file))
 
 
 if __name__ == "__main__":
-    src_directory = "src/snekmate"
-    process_directory(src_directory)
+    process_directory("src/snekmate")