Start adding "list item" types of objects with nicer accessors than raw dicts

Author: cstenac

dataikuapi/dss/dataset.py

@@ -3,12 +3,52 @@
 from ..utils import DataikuStreamedHttpUTF8CSVReader
 from .future import DSSFuture
 import json, warnings
+from .utils import DSSTaggableObjectListItem
 from .future import DSSFuture
 from .metrics import ComputedMetrics
 from .discussion import DSSObjectDiscussions
 from .statistics import DSSStatisticsWorksheet
 from . import recipe
 
+class DSSDatasetListItem(DSSTaggableObjectListItem):
+    """An item in a list of datasets. Do not instantiate this class"""
+    def __init__(self, client, data):
+        super(DSSDatasetListItem, self).__init__(data)
+        self.client = client
+
+    def to_dataset(self):
+        """Gets the :class:`DSSDataset` corresponding to this dataset"""
+        return DSSDataset(self.client, self._data["projectKey"], self._data["name"])
+
+    @property
+    def name(self):
+        return self._data["name"]
+    @property
+    def id(self):
+        return self._data["name"]
+    @property
+    def type(self):
+        return self._data["type"]
+    @property
+    def schema(self):
+        return self._data["schema"]
+
+    @property
+    def connection(self):
+        """Returns the connection on which this dataset is attached, or None if there is no connection for this dataset"""
+        if not "params" in self._data:
+            return None
+        return self._data["params"].get("connection", None)
+
+    def get_column(self, column):
+        """
+        Returns the schema column given a name.
+        :param str column: Column to find
+        :return a dict of the column settings or None if column does not exist
+        """
+        matched = [col for col in self.schema["columns"] if col["name"] == column]
+        return None if len(matched) == 0 else matched[0]
+
 class DSSDataset(object):
     """
     A dataset on the DSS instance

dataikuapi/dss/project.py

@@ -1,5 +1,5 @@
 import time, warnings, sys, os.path as osp
-from .dataset import DSSDataset, DSSManagedDatasetCreationHelper
+from .dataset import DSSDataset, DSSDatasetListItem, DSSManagedDatasetCreationHelper
 from .recipe import DSSRecipe
 from . import recipe
 from .managedfolder import DSSManagedFolder
@@ -208,15 +208,22 @@ def set_permissions(self, permissions):
     # Datasets
     ########################################################
 
-    def list_datasets(self):
+    def list_datasets(self, rtype="listitems"):
         """
-        List the datasets in this project
-        
-        :returns: The list of the datasets, each one as a dictionary. Each dataset dict contains at least a `name` field which is the name of the dataset
-        :rtype: list of dicts
+        List the datasets in this project.
+
+        :param str rtype: How to return the list. Supported values are "listitems" and "objects".
+        :returns: The list of the datasets. If "rtype" is "listitems", each one as a :class:`dataset.DSSDatasetListItem`.
+                  If "rtype" is "objects", each one as a :class:`dataset.DSSDataset`
+        :rtype: list
         """
-        return self.client._perform_json(
-            "GET", "/projects/%s/datasets/" % self.project_key)
+        items = self.client._perform_json("GET", "/projects/%s/datasets/" % self.project_key)
+        if rtype == "listitems":
+            return [DSSDatasetListItem(self.client, item) for item in items]
+        elif rtype == "objects":
+            return [DSSDataset(self.client, self.project_key, item["name"]) for item in items]
+        else:
+            raise ValueError("Unknown rtype")
 
     def get_dataset(self, dataset_name):
         """

dataikuapi/dss/utils.py

@@ -88,3 +88,14 @@ def from_full(ref):
             return AnyLoc(elts[0], elts[1])
         else:
             raise Exception("Cannot parse object id, it's not a full id")
+
+
+class DSSTaggableObjectListItem(dict):
+    """An item in a list of taggable objects. Do not instantiate this class"""
+    def __init__(self, data):
+        super(DSSTaggableObjectListItem, self).__init__(data)
+        self._data = data
+
+    @property
+    def tags(self):
+        return self._data["tags"]