Cant toggle edit mode after using bpy.ops.wm.open_mainfile

I need to reload my .blend file and then toggle edit mode to edit polygons repeatedly in my script.
Code to reproduce the behaviour (code is executed from the text editor in blender):

import bpy

bpy.ops.wm.open_mainfile(filepath=bpy.data.filepath)
bpy.ops.object.mode_set(mode='EDIT')

Error message: File "C:\Program Files\Blender Foundation\Blender 2.93\2.93\scripts\modules\bpy\ops.py", line 132, in __call__ ret = _op_call(self.idname_py(), None, kw) RuntimeError: Operator bpy.ops.object.mode_set.poll() failed, context is incorrect

There is an active object in the bpy.context, so when I toggle edit mode without reloading it works fine. It seems that after reloading some attributes of the context are missing (bpy.context.active_object, bpy.context.screen is None), but after script execution stops it’s fine, and edit mode toggle also works.

Solved it. The workaround is to use callback function with persistent content

Example:

# Function that enters edit mode and updates mesh after loading blend file
@bpy.app.handlers.persistent
def do_stuff(dummy):
    ...

# How to use it:
# activate do_stuff
bpy.app.handlers.load_post.append(do_stuff)
bpy.ops.wm.open_mainfile(filepath=filepath)  
# disable do_stuff (so that it doesnt get called every other time you load blend file)
bpy.app.handlers.load_post.remove(do_stuff)

Also, because I couldn’t pass variables to the function, I needed to use global context. global doesnt work in blender session, so I used my custom imported module to store global attributes (python - Declare a global variable in a Blender session - Blender Stack Exchange)