Created
February 19, 2024 08:47
-
-
Save alvinwan/f86d7c3d2ba331ce082a1c52b36ed794 to your computer and use it in GitHub Desktop.
How to mock submodules from just a single file
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
"Hello world" example of a single-file 'library'. | |
After importing this file, you can then use this file to mock submodules for | |
a library. For example, | |
import main # this file | |
import torch.nn.functional # 'imports' the mock module below instead | |
This helps with debugging, logging, and more. | |
""" | |
import sys | |
from importlib.abc import MetaPathFinder, Loader | |
from importlib.util import spec_from_file_location | |
import types | |
class MyCustomFinder(MetaPathFinder): | |
def find_spec(self, fullname, path, target=None): | |
# Implement your logic to locate the module specification. Return a | |
# ModuleSpec (module found) or None (module not found, causing import | |
# error). | |
return spec_from_file_location( | |
name=fullname, | |
location=None, | |
loader=MyLoader(fullname), | |
submodule_search_locations=[] | |
) | |
class MyLoader(Loader): | |
def __init__(self, name): | |
self.name = name | |
def create_module(self, spec): | |
# This is the module that is returned. You can attach module attributes | |
# to this object. | |
return types.ModuleType(self.name) | |
def exec_module(self, module): | |
pass | |
# Register your custom finder. Inserting at the start of the list ensures that | |
# our mock modules above are used, *even if a package is actually installed. If | |
# you instead want to mock a module only if it's *not installed, use .append | |
# instead. | |
sys.meta_path.insert(0, MyCustomFinder()) | |
# Import whatever random modules you want. We can now mock any submodules using | |
# the above. | |
import a.b.c |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment