sim_edit.py

import sys
import os
import threading
import sqlite3
import uuid

# add parent dir
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import env

e = env.ETLEnv(difficulty="hard")
obs = e.reset()

edits = [
    {"target_line": 6, "new_code": 'conn = sqlite3.connect("etl_target.db", check_same_thread=False)\nglobal_lock = threading.Lock()\n'},
    {"target_line": 14, "new_code": '        try:\n            with global_lock:\n                cursor.execute("INSERT INTO events (worker_id, event_no) VALUES (?, ?)", (worker_id, i))\n                conn.commit()\n'},
    {"target_line": 15, "new_code": ""},
    {"target_line": 16, "new_code": ""},
    {"target_line": 17, "new_code": ""},
    {"target_line": 18, "new_code": ""},
    {"target_line": 19, "new_code": ""}
]

for ed in edits:
    action = env.EditCodeAction(filepath="pipeline.py", target_line=ed["target_line"], new_code=ed["new_code"])
    e.step(action)

with open(e._workspace / "pipeline.py", "r") as f:
    for i, line in enumerate(f, 1):
        print(f"{i:2}: {repr(line)}")

e.close()