diff --git a/tests/connectors/test_recipe.py b/tests/connectors/test_recipe.py index 10799612..fbfd5aca 100644 --- a/tests/connectors/test_recipe.py +++ b/tests/connectors/test_recipe.py @@ -46,6 +46,37 @@ def test_write_recipe_connector(): assert df['Find2'].iloc[0] == 'brg' +def test_write_recipe_connector_returns_nested_dataframe(): + """ + A nested recipe used as a write: - recipe: step can shape the + return value of the outer recipe via its own write: - dataframe: step. + https://github.com/wrangleworks/WranglesPY/issues/1157 + """ + recipe = """ + read: + - test: + rows: 1 + values: + Col1: Val1 + Col2: Val2 + + write: + - recipe: + wrangles: + - create.column: + output: Col3 + value: Val3 + write: + - dataframe: + columns: + - Col2 + - Col3 + """ + df = wrangles.recipe.run(recipe) + assert list(df.columns) == ['Col2', 'Col3'] + assert df['Col3'].iloc[0] == 'Val3' + + # Running recipe def test_run_recipe_connector(): recipe = """ diff --git a/wrangles/connectors/recipe.py b/wrangles/connectors/recipe.py index 71181848..0df95665 100644 --- a/wrangles/connectors/recipe.py +++ b/wrangles/connectors/recipe.py @@ -109,7 +109,7 @@ def write( columns: list = None, functions: _Union[_types.FunctionType, list] = [], **kwargs -) -> None: +) -> _pd.DataFrame: """ Run a recipe, from a recipe! Recipe-ception. This will trigger a new recipe with the contents of the current recipe. @@ -121,6 +121,9 @@ def write( :param variables: (Optional) A dictionary of custom variables to override placeholders in the recipe. Variables can be indicated as ${MY_VARIABLE}. Variables can also be overwritten by Environment Variables. :param columns: (Optional) A list of the columns to pass to the recipe. If omitted, all columns will be included. :param functions: Pass in a custom function or list of custom functions that can be called in the recipe. + :return: The nested recipe's resulting dataframe, so it can be used to \ + shape the return value of the outer recipe (e.g. via the nested \ + recipe's own write: - dataframe: step). """ if variables is None: variables = {} @@ -130,7 +133,7 @@ def write( columns = _wildcard_expansion(df.columns, columns) df = df[columns] - _recipe.run(name, dataframe=df, variables=variables, functions=functions) + return _recipe.run(name, dataframe=df, variables=variables, functions=functions) _schema['write'] = """ diff --git a/wrangles/recipe.py b/wrangles/recipe.py index c35f19e4..99328747 100644 --- a/wrangles/recipe.py +++ b/wrangles/recipe.py @@ -1122,8 +1122,13 @@ def _write_data( # Validate with a placeholder for df _validate_function_args(func, {"df": None, **args}, export_type) - # Execute the function - func(df_temp, **args) + # Execute the function. If it returns a dataframe + # (e.g. the recipe connector reflecting a nested + # write: - dataframe: step), use that as the + # dataframe returned by this recipe. + result = func(df_temp, **args) + if result is not None: + df_return = result except Exception as e: # Wrap with enhanced error information _wrap_and_raise('WRITE', export_type, None, e)