Skip to content

Commit

Permalink
Summon dataframe support
Browse files Browse the repository at this point in the history
  • Loading branch information
dmpetrov committed Jan 26, 2020
1 parent 8f79921 commit d5a3b6f
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 19 deletions.
52 changes: 37 additions & 15 deletions dvc/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,14 @@
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,
Expand Down Expand Up @@ -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:
Expand All @@ -156,6 +159,16 @@ def _read_summon_content(self):
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):
Expand Down Expand Up @@ -198,17 +211,20 @@ def pull(self, dobj):
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(
Expand All @@ -221,18 +237,24 @@ def get_dobject(self, name):

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())
Expand Down
8 changes: 4 additions & 4 deletions tests/func/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
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
Expand Down Expand Up @@ -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>"},
Expand All @@ -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")
Expand Down

0 comments on commit d5a3b6f

Please sign in to comment.