Port a simplified version of FunctionRecipeCreator

Author: cstenac

dataikuapi/dss/dataset.py

@@ -471,7 +471,11 @@ def new_code_recipe(self, type, code=None, recipe_name=None):
 
         if recipe_name is None:
             recipe_name = "%s_recipe_from_%s" % (type, self.dataset_name)
-        builder = recipe.CodeRecipeCreator(recipe_name, type, self.project)
+
+        if type == "python":
+            builder = recipe.PythonRecipeCreator(recipe_name, self.project)
+        else:
+            builder = recipe.CodeRecipeCreator(recipe_name, type, self.project)
         builder.with_input(self.dataset_name)
         if code is not None:
             builder.with_script(code)

dataikuapi/dss/recipe.py

@@ -969,12 +969,77 @@ def with_new_output_dataset(self, name, connection,
         self.with_output(name, append=append)
         return self
 
-    
-
     def _finish_creation_settings(self):
         super(CodeRecipeCreator, self)._finish_creation_settings()
         self.creation_settings['script'] = self.script
 
+class PythonRecipeCreator(CodeRecipeCreator):
+    """
+    Creates a Python recipe.
+    A Python recipe can be defined either by its complete code, like a normal Python recipe, or
+    by a function signature.
+
+    When using a function, the function must take as arguments:
+     * A list of dataframes corresponding to the dataframes of the input datasets
+     * Optional named arguments corresponding to arguments passed to the creator
+    """
+
+    def __init__(self, name, project):
+        DSSRecipeCreator.__init__(self, "python", name, project)
+
+    DEFAULT_RECIPE_CODE_TMPL = """
+# This code is autogenerated by PythonRecipeCreator function mode
+import dataiku, dataiku.recipe, json
+from {module_name} import {fname}
+input_datasets = dataiku.recipe.get_inputs_as_datasets()
+output_datasets = dataiku.recipe.get_outputs_as_datasets()
+params = json.loads('{params_json}')
+
+logging.info("Reading %d input datasets as dataframes" % len(input_datasets))
+input_dataframes = [ds.get_dataframe() for ds in input_datasets]
+
+logging.info("Calling user function {fname}")
+function_input = input_dataframes if len(input_dataframes) > 1 else input_dataframes[0]
+output_dataframes = {fname}(function_input, **params)
+
+if not isinstance(output_dataframes, list):
+    output_dataframes = [output_dataframes]
+
+if not len(output_dataframes) == len(output_datasets):
+    raise Exception("Code function {fname}() returned %d dataframes but recipe expects %d output datasets", \\
+                                            (len(output_dataframes), len(output_datasets)))
+output = list(zip(output_datasets, output_dataframes))
+for ds, df in output:
+    logging.info("Writing function result to dataset %s" % ds.name)
+    ds.write_with_schema(df)
+"""
+
+    def with_function_name(self, module_name, function_name, function_args=None, custom_template=None):
+        """
+        Defines this recipe as being a functional recipe calling a function name from a module name
+        """
+        script_tmpl = PythonRecipeCreator.DEFAULT_RECIPE_CODE_TMPL if custom_template is None else custom_template
+
+        if function_args is None:
+            function_args = {}
+
+        code = script_tmpl.format(module_name=module_name, fname=function_name, params_json = json.dumps(function_args))
+        self.with_script(code)
+
+        return self
+
+    def with_function(self, fn, function_args=None, custom_template=None):
+        #TODO: add in documentation that relative imports wont work
+        module_name = inspect.getmodule(fn).__name__
+        fname = fn.__name__
+
+        # Validate that function_args apply to fn
+        argspec = inspect.getargspec(fn)
+        for k in function_args.keys():
+            if k not in argspec.args:
+                raise ValueError("Provided key argument {} not an argument of function {}".format(k, fname))
+
+        return with_function_name(module_name, fname, function_args, custom_template)
 
 class SQLQueryRecipeCreator(SingleOutputRecipeCreator):
     """