From d5a3b6f4f69082623f8421499e2b32eec2fd848d Mon Sep 17 00:00:00 2001
From: Dmitry Petrov <dmitry.petrov@nevesomo.com>
Date: Sun, 26 Jan 2020 02:50:42 -0800
Subject: [PATCH] Summon dataframe support

---
 dvc/api.py             | 52 ++++++++++++++++++++++++++++++------------
 tests/func/test_api.py |  8 +++----
 2 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/dvc/api.py b/dvc/api.py
index 4736b3fa7..d22c70091 100644
--- a/dvc/api.py
+++ b/dvc/api.py
@@ -14,12 +14,14 @@ from dvc.exceptions import DvcException, NotDvcRepoError
 from dvc.external_repo import external_repo
 
 DEF_SUMMON = "Summon.yaml"
+DOBJ_SECTION = "d-objects"
 
 SUMMON_FILE_SCHEMA = Schema(
     {
-        Required("objects"): [
+        Required(DOBJ_SECTION): [
             {
                 Required("name"): str,
+                "description": str,
                 "meta": dict,
                 Required("summon"): {
                     Required("type"): str,
@@ -142,8 +144,9 @@ def summon(name, repo=None, rev=None, summon_file=DEF_SUMMON, args=None):
 class SummonDesc(object):
     def __init__(self, repo_obj, summon_file=DEF_SUMMON):
         self.repo = repo_obj
+        self.filename = summon_file
         self.path = os.path.join(self.repo.root_dir, summon_file)
-        self.summon_content = self._read_summon_content()
+        self.content = self._read_summon_content()
 
     def _read_summon_content(self):
         try:
@@ -156,6 +159,16 @@ class SummonDesc(object):
         except Invalid as exc:
             raise SummonError(str(exc)) from exc
 
+    def _write_summon_content(self):
+        try:
+            with builtin_open(self.path, "w") as fobj:
+                content = SUMMON_FILE_SCHEMA(self.content)
+                ruamel.yaml.serialize_all(content, fobj)
+        except ruamel.yaml.YAMLError as exc:
+            raise SummonError("Summon file schema error") from exc
+        except Exception as exc:
+            raise SummonError(str(exc)) from exc
+
     @staticmethod
     @contextmanager
     def prepare_summon(repo=None, rev=None, summon_file=DEF_SUMMON):
@@ -198,17 +211,20 @@ class SummonDesc(object):
                 self.repo.cloud.pull(out.get_used_cache())
                 out.checkout()
 
-    # def to_abs_paths(self, paths):
-    #     return [self.repo.find_out_by_relpath(d) for d in paths]
+    def push(self, dobj):
+        paths = self.deps_abs_paths(dobj)
 
-    def get_dobject(self, name):
+        with self.repo.state:
+            for path in paths:
+                self.repo.add(path)
+                self.repo.add(path)
+
+    def get_dobject(self, name, default=False):
         """
         Given a summonable object's name, search for it on the given content
         and return its description.
         """
-        objects = [
-            x for x in self.summon_content["objects"] if x["name"] == name
-        ]
+        objects = [x for x in self.content[DOBJ_SECTION] if x["name"] == name]
 
         if not objects:
             raise SummonErrorNoObjectFound(
@@ -221,18 +237,24 @@ class SummonDesc(object):
 
         return objects[0]
 
-    def set_dobject(self, obj_new, overwrite=False):
+    def update_dobj(self, new_dobj, overwrite=False):
         try:
-            name = obj_new["name"]
-            obj = self.get_dobject(name)
+            name = new_dobj["name"]
+            dobj = self.get_dobject(name)
 
             if overwrite:
-                idx = self.summon_content["objects"].index(obj)
-                self.summon_content["objects"][idx] = obj_new
+                idx = self.content[DOBJ_SECTION].index(dobj)
+                self.content[DOBJ_SECTION][idx] = new_dobj
             else:
-                raise SummonError("Object '{}' already exist".format(name))
+                raise SummonError(
+                    "D-object '{}' already exist in '{}'".format(
+                        name, self.filename
+                    )
+                )
         except SummonErrorNoObjectFound:
-            self.summon_content["objects"].append(obj_new)
+            self.content[DOBJ_SECTION].append(new_dobj)
+
+        self._write_summon_content()
 
 
 @wrap_with(threading.Lock())
diff --git a/tests/func/test_api.py b/tests/func/test_api.py
index a9d665075..ad078fed4 100644
--- a/tests/func/test_api.py
+++ b/tests/func/test_api.py
@@ -6,7 +6,7 @@ import ruamel.yaml
 import pytest
 
 from dvc import api
-from dvc.api import SummonError, UrlNotDvcRepoError, DEF_SUMMON
+from dvc.api import SummonError, UrlNotDvcRepoError, DEF_SUMMON, DOBJ_SECTION
 from dvc.compat import fspath
 from dvc.exceptions import FileMissingError
 from dvc.main import main
@@ -145,7 +145,7 @@ def test_open_not_cached(dvc):
 
 def test_summon(tmp_dir, dvc, erepo_dir):
     objects = {
-        "objects": [
+        DOBJ_SECTION: [
             {
                 "name": "sum",
                 "meta": {"description": "Add <x> to <number>"},
@@ -160,10 +160,10 @@ def test_summon(tmp_dir, dvc, erepo_dir):
     }
 
     other_objects = copy.deepcopy(objects)
-    other_objects["objects"][0]["summon"]["args"]["x"] = 100
+    other_objects[DOBJ_SECTION][0]["summon"]["args"]["x"] = 100
 
     dup_objects = copy.deepcopy(objects)
-    dup_objects["objects"] *= 2
+    dup_objects[DOBJ_SECTION] *= 2
 
     with erepo_dir.chdir():
         erepo_dir.dvc_gen("number", "100", commit="Add number.dvc")