Merge remote-tracking branch 'origin/feature/new-dataset-management-apis' into feature/new-dataset-management-apis-patch

Author: kgued

dataikuapi/dss/dataset.py

@@ -475,7 +475,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/flow.py

@@ -1,5 +1,6 @@
 from .utils import AnyLoc
-from .recipe import DSSRecipeDefinitionAndPayload
+from .dataset import DSSDataset
+from .recipe import DSSRecipe, DSSRecipeDefinitionAndPayload
 from .future import DSSFuture
 import logging
 
@@ -8,6 +9,9 @@ def __init__(self, client, project):
         self.client = client
         self.project = project
 
+    def get_graph(self):
+        data = self.client._perform_json("GET", "/projects/%s/flow/graph/" % (self.project.project_key))
+        return DSSProjectFlowGraph(self, data)
 
     def replace_input_computable(self, current_ref, new_ref, type="DATASET"):
         """
@@ -81,6 +85,84 @@ def propagate_schema(self, dataset_name, rebuild=False, recipe_update_options={}
         return DSSFuture(self.client,update_future.get('jobId', None), update_future)
 
 
+class DSSProjectFlowGraph(object):
+
+    def __init__(self, flow, data):
+        self.flow = flow
+        self.data = data
+
+    def get_source_computables(self, as_type="dict"):
+        """
+        Returns the list of source computables.
+        :param as_type: How to return the source computables. Possible values are "dict" and "object"
+        
+        :return: if as_type=dict, each computable is returned as a dict containing at least "ref" and "type". 
+                 if as_type=object, each computable is returned as a  :class:`dataikuapi.dss.dataset.DSSDataset`,
+                    :class:`dataikuapi.dss.managedfolder.DSSManagedFolder`,
+                    :class:`dataikuapi.dss.savedmodel.DSSSavedModel`, or streaming endpoint
+        """
+        ret = []
+        for computable in self.data["computables"].values():
+            if len(computable["predecessors"]) == 0:
+                ret.append(computable)
+        return self._convert_computables_list(ret, as_type)
+
+    def get_source_datasets(self):
+        """
+        Returns the list of source datasets for this project.
+        :rtype list of :class:`dataikuapi.dss.dataset.DSSDataset`
+        """
+        return [self._get_object_from_computable(x) for x in self.get_source_computables() if x["type"] == "COMPUTABLE_DATASET"]
+
+    def get_successor_recipes(self, node, as_type="name"):
+        """
+        Returns a list of recipes that are a successor of a graph node
+
+        :param node: Either a name or :class:`dataikuapi.dss.dataset.DSSDataset` object
+        :return if as_type="name", list of strings, recipe names
+                    else list of :class:`dataikuapi.dss.recipe.DSSRecipe`
+        """
+        if isinstance(node, DSSDataset):
+            node = node.dataset_name
+
+        computable = self.data["computables"].get(node, None)
+        if computable is None:
+            raise ValueError("Computable %s not found in Flow graph" % node)
+
+        names= computable["successors"]
+
+        if as_type == "object":
+            return [DSSRecipe(self.flow.client, self.flow.project.project_key, x) for x in names]
+        else:
+            return names
+
+    def get_successor_computables(self, node, as_type="dict"):
+        """
+        Returns a list of computables that are a successor of a given graph node
+        Each computable is returned as a dict containing at least "ref" and "type"
+        """
+        if isinstance(node, DSSRecipe):
+            node = node.recipe_name
+        runnable = self.data["runnables"].get(node, None)
+        if runnable is None:
+            raise ValueError("Runnable %s not found in Flow graph" % node)
+
+        computables = [self.data["computables"][x] for x in runnable["successors"]]
+        return self._convert_computables_list(computables, as_type)
+
+    def _convert_computables_list(self, computables, as_type):
+        if as_type == "object":
+            return [self._get_object_from_computable(computable) for computable in computables]
+        else:
+            return computables
+
+    def _get_object_from_computable(self, computable):
+        if computable["type"] == "COMPUTABLE_DATASET":
+            return DSSDataset(self.flow.client, self.flow.project.project_key, computable["ref"])
+        else:
+            # TODO
+            raise Exception("unsupported %s" % computable["type"])
+
 class DSSFlowTool(object):
     """
     Handle to interact with a flow tool

dataikuapi/dss/recipe.py

@@ -1053,80 +1053,69 @@ def _finish_creation_settings(self):
         super(CodeRecipeCreator, self)._finish_creation_settings()
         self.creation_settings['script'] = self.script
 
-class FunctionCodeRecipeCreator(CodeRecipeCreator):
 
-    DEFAULT_RECIPE_CODE_TMPL = """
-import dataiku
-import json
-from {module_name} import {fname}
-
-input_ds_names = {input_ds_names}
-input_ds_list = [dataiku.Dataset(name) for name in input_ds_names]
+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.
 
-output_ds_names = {output_ds_names}
-output_ds_list = [dataiku.Dataset(name) for name in output_ds_names]
+    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
+    """
 
-params = json.loads({params_json})
+    def __init__(self, name, project):
+        DSSRecipeCreator.__init__(self, "python", name, project)
 
-input_df_list = [ds.get_dataframe() for ds in input_ds_list]
-fn_input = input_df_list if len(input_df_list) > 1 else input_df_list[0]
+    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}')
 
-output_df_list = {fname}({input_arg}, **params)
+logging.info("Reading %d input datasets as dataframes" % len(input_datasets))
+input_dataframes = [ds.get_dataframe() for ds in input_datasets]
 
-if not isinstance(output_df_list, list):
-    output_df_list = [output_df_list]
+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 len(output_df_list) == len(output_ds_list):
-    raise Exception("Code function {fname}() returned items len: %d, \\
-                    does not match expected recipe output len: %d" % (len(output_df_list), len(output_ds_list)))
+if not isinstance(output_dataframes, list):
+    output_dataframes = [output_dataframes]
 
-output = list(zip(output_ds_list, output_df_list))
+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 __init__(self, name, type, project):
-        CodeRecipeCreator.__init__(self, name, type, project)
+    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
 
-    def with_function(self, fn, input_arg=None, code_template=None):
+        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):
+        import inspect
         #TODO: add in documentation that relative imports wont work
         module_name = inspect.getmodule(fn).__name__
         fname = fn.__name__
-
-        input_ds_names = self.get_input_refs_for_role(role="main")
-        output_ds_names = self.get_output_refs_for_role(role="main")
-
-        input_arg = 'fn_input' if not input_arg else '{}=fn_input'.format(input_arg)
-
-        script_tmpl = FunctionCodeRecipeCreator.DEFAULT_RECIPE_CODE_TMPL if code_template is None else code_template
-
-        def generate_code_fn(**kwargs):
-            code_tmpl = script_tmpl.format(
-                module_name=module_name,
-                fname=fname,
-                input_ds_names=repr(input_ds_names),
-                output_ds_names=repr(output_ds_names),
-                input_arg=input_arg,
-                params_json = '{params_json}'
-                )
-
-            #TODO: deal with default args
-            argspec = inspect.getargspec(fn)
-            for k in kwargs.keys():
-                if k not in argspec.args:
-                    raise ValueError("Provided key argument {} not an argument of function {}".format(k, fname))
-
-            #params_string = ','.join(["{}='{}'".format(k,v) for k,v in kwargs.items()])
-            params_json = json.dumps(kwargs)
-            code = code_tmpl.format(params_json=repr(params_json))
-
-            self.with_script(code)
-
-            return self
-
-        return generate_code_fn
+        return self.with_function_name(module_name, fname, function_args, custom_template)
+>>>>>>> origin/feature/new-dataset-management-apis
 
 class SQLQueryRecipeCreator(SingleOutputRecipeCreator):
     """