Python: Force import to prefer .py over .so

Question

I have a situation where the same Python module is present in the same directory in two different versions; mymodule.py and mymodule.so (I obtain the latter from the first via Cython, but that's irrelevant to my question). When from Python I do

import mymodule

it always chooses mymodule.so. Sometimes I really want to import mymodule.py instead. I could temporarily move mymodule.so to another location, but that does not play well if I simultaneously have another Python instance running which needs to import mymodule.so.

The question is how to make import prefer .py files over .so, rather than vice versa?

Here's my thoughts on a solution: I imagine performing some magic using importlib and possibly edit sys.meta_path. Specifically I see that sys.meta_path[2] holds _frozen_importlib_external.PathFinder which is used to import external modules, i.e. this is used for both mymodule.py and mymodule.so. If I could just replace this with a similar PathFinder which used the reverse ordering for file types, I would have a solution.

I'm using Python 3.7, if that affects the solution.

Edit

Note that simply reading in the source lines of mymodule.py and exec'ing them won't do, as mymodule.py may itself import other modules which again exist in both a .py and .so version (I want to import the .py version of these as well).

Comus

Answer 1

Using these notes I came up with this. It's not too pretty, but it seems to work.

import glob, importlib, sys

def hook(name):
    if name != '.':
        raise ImportError()
    modnames = set(f.rstrip('.py') for f in glob.glob('*.py'))
    return Finder(modnames)
sys.path_hooks.insert(1, hook)
sys.path.insert(0, '.')

class Finder(object):
    def __init__(self, modnames):
        self.modnames = modnames
    def find_spec(self, modname, target=None):
        if modname in self.modnames:
            origin = './' + modname + '.py'
            loader = Loader()
            return importlib.util.spec_from_loader(modname, loader, origin=origin)
        else:
            return None

class Loader(object):
    def create_module(self, target):
        return None
    def exec_module(self, module):
        with open(module.__spec__.origin, 'r', encoding='utf-8') as f:
            code = f.read()
        compile(code, module.__spec__.origin, 'exec')
        exec(code, module.__dict__)

Comus

Answer 2

Here is another solution, that works just by tweaking the finders that the runtime generates by default. This uses hidden implementation details (FileFinder._loaders), but I've tested on CPythons 3.7, 3.8, and 3.9.

from contextlib import contextmanager
from dataclasses import dataclass
from importlib.machinery import FileFinder
from importlib.abc import Finder
import sys
from typing import Callable


@dataclass
class PreferPureLoaderHook:
    orig_hook: Callable[[str], Finder]

    def __call__(self, path: str) -> Finder:
        finder = self.orig_hook(path)
        if isinstance(finder, FileFinder):
            # Move pure python file loaders to the front
            finder._loaders.sort(key=lambda pair: 0 if pair[0] in (".py", ".pyc") else 1)  # type: ignore
        return finder


@contextmanager
def prefer_pure_python_imports():
    sys.path_hooks = [PreferPureLoaderHook(h) for h in sys.path_hooks]
    sys.path_importer_cache.clear()
    yield
    assert all(isinstance(h, PreferPureLoaderHook) for h in sys.path_hooks)
    sys.path_hooks = [h.orig_hook for h in sys.path_hooks]
    sys.path_importer_cache.clear()

with prefer_pure_python_imports():
    ...

Avrilavrit