fix code display for temporary modules
diff --git a/hamilton/ad_hoc_utils.py b/hamilton/ad_hoc_utils.py
index 773d078..5b2dab3 100644
--- a/hamilton/ad_hoc_utils.py
+++ b/hamilton/ad_hoc_utils.py
@@ -5,7 +5,7 @@
import types
import uuid
from types import ModuleType
-from typing import Callable
+from typing import Callable, Optional
def _copy_func(f):
@@ -60,9 +60,9 @@
return module
-def module_from_source(source: str) -> ModuleType:
+def module_from_source(source: str, module_name: Optional[str] = None) -> ModuleType:
"""Create a temporary module from source code"""
- module_name = _generate_unique_temp_module_name()
+ module_name = module_name if module_name else _generate_unique_temp_module_name()
module_object = ModuleType(module_name)
code_object = compile(source, module_name, "exec")
sys.modules[module_name] = module_object
diff --git a/ui/sdk/src/hamilton_sdk/driver.py b/ui/sdk/src/hamilton_sdk/driver.py
index 0087af7..7677ce8 100644
--- a/ui/sdk/src/hamilton_sdk/driver.py
+++ b/ui/sdk/src/hamilton_sdk/driver.py
@@ -2,6 +2,7 @@
import hashlib
import inspect
import json
+import linecache
import logging
import operator
import os
@@ -67,15 +68,19 @@
f"attribute or it is None. This happens with lazy loaders."
)
continue
- # Check if the module is in the same top level package
- if value.__package__ != module.__package__ and not value.__package__.startswith(
- module.__package__
- ):
- logger.debug(
- f"Skipping hash for module {value.__name__} because it is in a different "
- f"package {value.__package__} than {module.__package__}"
- )
- continue
+
+ # Modules imported in a temporary module have no `__package__` attribute
+ if module.__package__:
+ # Check if the module is in the same top level package
+ if value.__package__ != module.__package__ and not value.__package__.startswith(
+ module.__package__
+ ):
+ logger.debug(
+ f"Skipping hash for module {value.__name__} because it is in a different "
+ f"package {value.__package__} than {module.__package__}"
+ )
+ continue
+
# Recursively hash the sub-module
hash_object = _hash_module(value, hash_object, seen_modules)
@@ -688,6 +693,11 @@
def _slurp_code(fg: graph.FunctionGraph, repo_base: str) -> List[dict]:
+ """Get the source code from modules. Returns a list with a dictionary for each module.
+
+ The `path` attribute needs to match the `path` of code artifacts generated by
+ `extract_code_artifacts_from_function_graph()`
+ """
modules = set()
for node_ in fg.nodes.values():
originating_functions = node_.originating_functions
@@ -702,6 +712,14 @@
module_path = os.path.relpath(module.__file__, repo_base)
with open(module.__file__, "r") as f:
out.append({"path": module_path, "contents": f.read()})
+ # for temporary modules registed via `module_from_source`
+ else:
+ # get source code from the linecache; returns a tuple (size, mtime, lines, fullname)
+ source_lines = linecache.cache[module.__name__][2]
+ source = "".join(source_lines)
+ # the path won't have a `.py` suffix to match `extract_code_artifacts_from_function_grap()`
+ module_path = os.path.relpath(module.__name__, repo_base)
+ out.append({"path": module_path, "contents": source})
return out
