-
Notifications
You must be signed in to change notification settings - Fork 32
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #278 from materialsproject/source_loader
Add custom source loading
- Loading branch information
Showing
17 changed files
with
414 additions
and
16 deletions.
There are no files selected for viewing
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
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,6 @@ | ||
uvicorn==0.11.8 | ||
boto3==1.14.60 | ||
hvac==0.10.5 | ||
IPython==7.16.1 | ||
nbformat==5.0.7 | ||
regex==2020.6.8 |
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
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
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
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,178 @@ | ||
import importlib.util | ||
import sys | ||
from glob import glob | ||
from importlib.abc import Loader, MetaPathFinder | ||
from importlib.machinery import ModuleSpec, SourceFileLoader | ||
from pathlib import Path | ||
from typing import List | ||
|
||
|
||
from maggma.core import Builder | ||
|
||
try: | ||
import nbformat | ||
from IPython import get_ipython | ||
from IPython.core.interactiveshell import InteractiveShell | ||
from regex import match | ||
except ModuleNotFoundError: | ||
pass | ||
|
||
_BASENAME = "maggma.cli.sources" | ||
|
||
|
||
class ScriptFinder(MetaPathFinder): | ||
""" | ||
Special Finder designed to find custom script builders | ||
""" | ||
|
||
@classmethod | ||
def find_spec(cls, fullname, path, target=None): | ||
if not (str(fullname).startswith(f"{_BASENAME}.")): | ||
return None | ||
|
||
# The last module is what we want to find the path for | ||
sub_path = str(fullname).split(".")[-1] | ||
segments = sub_path.split("_") | ||
|
||
file_path = next(find_matching_file(segments)) | ||
|
||
if file_path is None: | ||
return None | ||
|
||
return spec_from_source(file_path) | ||
|
||
|
||
class NotebookLoader(Loader): | ||
"""Module Loader for Jupyter Notebooks or Source Files""" | ||
|
||
def __init__(self, name=None, path=None): | ||
|
||
self.shell = InteractiveShell.instance() | ||
|
||
self.name = name | ||
self.path = path | ||
|
||
def create_module(self, spec): | ||
return None | ||
|
||
def exec_module(self, module): | ||
|
||
module.__dict__["get_ipython"] = get_ipython | ||
module.__path__ = self.path | ||
|
||
# load the notebook object | ||
with open(self.path, "r", encoding="utf-8") as f: | ||
nb = nbformat.read(f, 4) | ||
|
||
# extra work to ensure that magics that would affect the user_ns | ||
# actually affect the notebook module's ns | ||
save_user_ns = self.shell.user_ns | ||
self.shell.user_ns = module.__dict__ | ||
|
||
try: | ||
for cell in nb.cells: | ||
if cell.cell_type == "code": | ||
# transform the input to executable Python | ||
code = self.shell.input_transformer_manager.transform_cell( | ||
cell.source | ||
) | ||
# run the code in themodule | ||
exec(code, module.__dict__) | ||
finally: | ||
self.shell.user_ns = save_user_ns | ||
return module | ||
|
||
|
||
def spec_from_source(file_path: str) -> ModuleSpec: | ||
""" | ||
Returns a ModuleSpec from a filepath for importlib loading | ||
Specialized for loading python source files and notebooks into | ||
a temporary maggma cli package to run as a builder | ||
""" | ||
file_path_obj = Path(file_path).resolve().relative_to(Path(".").resolve()) | ||
file_path_str = str(file_path_obj) | ||
|
||
if file_path_obj.parts[-1][-3:] == ".py": | ||
# Gets module name from the filename without the .py extension | ||
module_name = "_".join(file_path_obj.parts).replace(" ", "_").replace(".py", "") | ||
|
||
spec = ModuleSpec( | ||
name=f"{_BASENAME}.{module_name}", | ||
loader=SourceFileLoader( | ||
fullname=f"{_BASENAME}.{module_name}", path=file_path_str | ||
), | ||
origin=file_path_str, | ||
) | ||
# spec._set_fileattr = True | ||
elif file_path_obj.parts[-1][-6:] == ".ipynb": | ||
# Gets module name from the filename without the .ipnb extension | ||
module_name = ( | ||
"_".join(file_path_obj.parts).replace(" ", "_").replace(".ipynb", "") | ||
) | ||
|
||
spec = ModuleSpec( | ||
name=f"{_BASENAME}.{module_name}", | ||
loader=NotebookLoader( | ||
name=f"{_BASENAME}.{module_name}", path=file_path_str | ||
), | ||
origin=file_path_str, | ||
) | ||
# spec._set_fileattr = True | ||
else: | ||
raise Exception( | ||
"Can't load {file_path}. Must provide a python source file such as a .py or .ipynb file" | ||
) | ||
|
||
return spec | ||
|
||
|
||
def load_builder_from_source(file_path: str) -> List[Builder]: | ||
""" | ||
Loads Maggma Builders from a Python source file | ||
""" | ||
file_path = str(Path(file_path).resolve()) | ||
spec = spec_from_source(file_path) | ||
module_object = importlib.util.module_from_spec(spec) | ||
spec.loader.exec_module(module_object) # type: ignore | ||
|
||
sys.modules[spec.name] = module_object | ||
|
||
if hasattr(module_object, "__builders__"): | ||
return getattr(module_object, "__builders__") | ||
elif hasattr(module_object, "__builder__"): | ||
return getattr(module_object, "__builder__") | ||
else: | ||
raise Exception( | ||
f"No __builders__ or __builder__ attribute found in {file_path}" | ||
) | ||
|
||
|
||
def find_matching_file(segments, curr_path="./"): | ||
""" | ||
Finds file that has the right sequence of segments | ||
in the path relative to the current path | ||
Requires all segments match the file path | ||
""" | ||
|
||
# If we've gotten to the end of the segment match check to see if a file exists | ||
if len(segments) == 0: | ||
if Path(curr_path + ".py").exists(): | ||
yield curr_path + ".py" | ||
if Path(curr_path + ".ipynb").exists(): | ||
yield curr_path + ".ipynb" | ||
else: | ||
# Recurse down the segment tree some more | ||
current_segment = segments[0] | ||
remainder = segments[1:] | ||
|
||
re = fr"({curr_path}[\s_]*{current_segment})" | ||
pos_matches = [match(re, pos_path) for pos_path in glob(curr_path + "*")] | ||
pos_matches = {pmatch.group(1) for pmatch in pos_matches if pmatch} | ||
for new_path in pos_matches: | ||
if Path(new_path).exists() and Path(new_path).is_dir: | ||
for sub_match in find_matching_file( | ||
remainder, curr_path=new_path + "/" | ||
): | ||
yield sub_match | ||
for sub_match in find_matching_file(remainder, curr_path=new_path): | ||
yield sub_match |
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
""" Dummy module to allow for loading dynamic source files """ |
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
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
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
Oops, something went wrong.