Importing AST modification

Question

I'm attempting dynamic rewrite of pywin32 win32com.client module prior to import, the below seems to work - but I'm not happy with the 6 lines of code for importing (after the last comment). Can anyone recommend a more concise way of creating / importing the resulting module (package)?

import sys, ast, types
import os.path
import importlib.util

import win32com

# import ast
_win32com_client_spec = importlib.util.find_spec("win32com.client")
_win32com_client_ast = ast.parse(open(_win32com_client_spec.origin).read())

(_win32com_client_dispatch_ast,) = (x for x in _win32com_client_ast.body if isinstance(x, ast.FunctionDef) and x.name == 'Dispatch')

# rename Dispatch
_win32com_client_dispatch_ast.name = '_OldDispatch'

# create new Dispatch
_my_dispatch_mod_ast = ast.parse("""
def Dispatch(*args, **kwds):
    base_inst = _OldDispatch(*args, **kwds)
    spec_inst = _OldDispatch(base_inst)
    return spec_inst
""")
(_my_dispatch_ast,) = _my_dispatch_mod_ast.body

# insert new Dispatch in module
_win32com_client_ast.body.insert(_win32com_client_ast.body.index(_win32com_client_dispatch_ast)+1, _my_dispatch_ast)

# import the package
_my_win32com_client_co = compile(_win32com_client_ast, '<patched win32com.client>', 'exec')
_my_win32com_client_mod = types.ModuleType('win32com.client')
_my_win32com_client_mod.__path__ = (os.path.dirname(_win32com_client_spec.origin),)
sys.modules[_my_win32com_client_mod.__name__] = _my_win32com_client_mod
win32com.client = _my_win32com_client_mod
exec(_my_win32com_client_co, _my_win32com_client_mod.__dict__)

Answer 1

further reading of the documentation shows using the ModuleSpec object is preferable to using types.ModuleType.

ref: https://docs.python.org/3/library/importlib.html#importlib.util.module_from_spec

the resulting code implicitly handles the __path__ for the package, and allows to reuse the spec object throughout. Beware if using the spec origin to show patching as below, the file attribute should be set explicitly as below.

The module should be imported prior to being exec'd, this seems to be important for submodule imports to work as found in gencache modules e.g. import win32com.client.CLSISToClass:

import sys, ast
import importlib.util

import win32com

if 'win32com.client' not in sys.modules:
    # import ast
    _win32com_client_spec = importlib.util.find_spec("win32com.client")
    _win32com_client_path = _win32com_client_spec.origin
    _win32com_client_spec.origin = '<patched win32com.client>'
    _win32com_client_ast = ast.parse(open(_win32com_client_path).read())

    # rename Dispatch
    (_win32com_client_dispatch_ast,) = (x for x in _win32com_client_ast.body if isinstance(x, ast.FunctionDef) and x.name == 'Dispatch')
    _win32com_client_dispatch_ast.name = '_OldDispatch'

    # create new Dispatch (disregard line numbers)
    _my_dispatch_mod_ast = ast.parse("""
def Dispatch(*args, **kwds):
    base_inst = _OldDispatch(*args, **kwds)
    spec_inst = _OldDispatch(base_inst)
    return spec_inst
""")
    (_my_dispatch_ast,) = _my_dispatch_mod_ast.body

    # insert new Dispatch in module
    _win32com_client_ast.body.insert(_win32com_client_ast.body.index(_win32com_client_dispatch_ast)+1, _my_dispatch_ast)

    # compile (original path for debugging), create, import and exec the module
    _my_win32com_client_co = compile(_win32com_client_ast, _win32com_client_path, 'exec')
    _my_win32com_client_mod = importlib.util.module_from_spec(_win32com_client_spec)
    _my_win32com_client_mod.__file__ = _win32com_client_path
    win32com.client = sys.modules['win32com.client'] = _my_win32com_client_mod
    try:
        exec(_my_win32com_client_co, _my_win32com_client_mod.__dict__)
    except:
        del win32com.client, sys.modules['win32com.client']
        raise