def _prepare_transformer_assets(fn: Callable, assets: Dict = None):
notebook_path = jputils.get_notebook_path()
processor = NotebookProcessor(nb_path=notebook_path, skip_validation=True)
fn_source = astutils.get_function_source(fn, strip_signature=False)
missing_names = flakeutils.pyflakes_report(
processor.get_imports_and_functions() + "\n" + fn_source)
if not assets:
assets = dict()
if not isinstance(assets, dict):
ValueError("Please provide preprocessing assets as a dictionary"
" mapping variables *names* to their objects")
missing_assets = [x not in assets.keys() for x in missing_names]
if any(missing_assets):
raise RuntimeError(
"The following abjects are a dependency for the"
" provided preprocessing function. Please add the"
" to the `preprocessing_assets` dictionary: %s" %
[a for a, m in zip(missing_names, missing_assets) if m])
# save function and assets
utils.clean_dir(TRANSFORMER_ASSETS_DIR)
marshal.set_data_dir(TRANSFORMER_ASSETS_DIR)
marshal.save(fn, TRANSFORMER_FN_ASSET_NAME)
for asset_name, asset_value in assets.items():
marshal.save(asset_value, asset_name)
# save notebook as well
shutil.copy(
notebook_path,
os.path.join(TRANSFORMER_ASSETS_DIR, TRANSFORMER_SRC_NOTEBOOK_NAME))
def test_get_function_source_with_signature():
"""Test that the function source is properly retrieved with signature."""
def foo():
a = 1 # noqa: F841
b = 2 # noqa: F841
target = "def foo():\n a = 1\n b = 2\n"
assert kale_ast.get_function_source(foo, strip_signature=False) == target
def test_get_function_source():
"""Test that the function source is properly retrieved."""
def foo():
a = 1 # noqa: F841
b = 2 # noqa: F841
target = "a = 1\nb = 2\n"
assert kale_ast.get_function_source(foo) == target
def rendered_source(self):
"""Source to be rendered in the template."""
# FIXME: I don't like this approach. Currently step.source is either
# a list of strings (if processed from the notebook) or a callable
# object (function) (if processed from the sdk). This means that when
# rendering the sdk template, we need to get the function's source.
# It would be great to, in some way, unify how we treat the "source"
# both for the notebook of the SDK all the way from the step object
# to the template.
return astutils.get_function_source(self.source, strip_signature=False)
def test_get_function_source_with_signature_and_decorator():
"""Test that the function source is properly retrieved with signature."""
def foodec(**kwargs):
def _(fn):
return fn
return _
@foodec(a=1, b=2, c=3, d=4)
def foo():
a = 1 # noqa: F841
b = 2 # noqa: F841
target = "def foo():\n a = 1\n b = 2\n"
assert kale_ast.get_function_source(foo, strip_signature=False) == target
def __init__(self,
pipeline_function: Callable,
config: PipelineConfig = None,
**kwargs):
self.pipeline_fn = pipeline_function
# self.pipeline.origin = PipelineOrigin.PYTHON
# Will be populated during processing
self._steps_input_vars = None
self._steps_return_vars = None
super().__init__(config=config, **kwargs)
fn_source = astutils.get_function_source(self.pipeline_fn)
self.validate(fn_source)
self.process(fn_source)
