Merge pull request #582 from efiop/master

dvc: validate stage files

dvc/output.py

@@ -2,6 +2,7 @@ import os
 import stat
 import json
 import shutil
+import schema
 from checksumdir import dirhash
 
 from dvc.system import System
@@ -50,6 +51,11 @@ class Dependency(object):
     PARAM_MD5 = 'md5'
     MD5_DIR_SUFFIX = '.dir'
 
+    SCHEMA = {
+        PARAM_PATH: str,
+        schema.Optional(PARAM_MD5): schema.Or(str, None),
+    }
+
     def __init__(self, project, path, md5=None):
         self.project = project
         self.path = os.path.abspath(os.path.realpath(path))
@@ -126,6 +132,9 @@ class Dependency(object):
 class Output(Dependency):
     PARAM_CACHE = 'cache'
 
+    SCHEMA = Dependency.SCHEMA
+    SCHEMA[schema.Optional(PARAM_CACHE)] = bool
+
     def __init__(self, project, path, md5=None, use_cache=True):
         super(Output, self).__init__(project, path, md5=md5)
         self.use_cache = use_cache

dvc/stage.py

@@ -2,9 +2,11 @@ import os
 import yaml
 import itertools
 import subprocess
+import schema
 
 from dvc.exceptions import DvcException
 from dvc.output import Output, Dependency, OutputError
+from dvc.logger import Logger
 
 
 class StageCmdFailedError(DvcException):
@@ -13,6 +15,11 @@ class StageCmdFailedError(DvcException):
         super(StageCmdFailedError, self).__init__(msg)
 
 
+class StageFileFormatError(DvcException):
+    def __init__(self):
+        super(StageFileFormatError, self).__init__('Stage file format error')
+
+
 class MissingDataSource(OutputError):
     def __init__(self, missing_files):
         assert len(missing_files) > 0
@@ -33,6 +40,12 @@ class Stage(object):
     PARAM_DEPS = 'deps'
     PARAM_OUTS = 'outs'
 
+    SCHEMA = {
+        schema.Optional(PARAM_CMD): schema.Or(str, None),
+        schema.Optional(PARAM_DEPS): schema.Or(schema.And(list, schema.Schema([Dependency.SCHEMA])), None),
+        schema.Optional(PARAM_OUTS): schema.Or(schema.And(list, schema.Schema([Output.SCHEMA])), None),
+    }
+
     def __init__(self, project, path=None, cmd=None, cwd=None, deps=[], outs=[]):
         self.project = project
         self.path = path
@@ -94,8 +107,18 @@ class Stage(object):
 
         return self
 
+    @staticmethod
+    def validate(d):
+        try:
+            schema.Schema(Stage.SCHEMA).validate(d)
+        except schema.SchemaError as exc:
+            Logger.debug(str(exc))
+            raise StageFileFormatError()
+
     @staticmethod
     def loadd(project, d, path):
+        Stage.validate(d)
+
         path = os.path.abspath(path)
         cwd = os.path.dirname(path)
         cmd = d.get(Stage.PARAM_CMD, None)

requirements.txt

@@ -15,3 +15,4 @@ checksumdir>=1.1.4
 setuptools>=34.0.0
 nanotime>=0.5.2
 pyasn1>=0.4.1
+schema>=0.6.7

setup.py

@@ -20,6 +20,7 @@ install_requires = [
     "setuptools>=34.0.0",
     "nanotime>=0.5.2",
     "pyasn1>=0.4.1",
+    "schema>=0.6.7",
 ]
 
 setup(

tests/test_stage.py

@@ -0,0 +1,63 @@
+from dvc.output import Output
+from dvc.stage import Stage, StageFileFormatError
+
+from tests.basic_env import TestDvc
+
+
+class TestSchema(TestDvc):
+    def _validate_fail(self, d):
+        with self.assertRaises(StageFileFormatError):
+            Stage.validate(d)
+
+
+class TestSchemaCmd(TestSchema):
+    def test_cmd_object(self):
+        d = {Stage.PARAM_CMD: {}}
+        self._validate_fail(d)
+
+    def test_cmd_none(self):
+        d = {Stage.PARAM_CMD: None}
+        Stage.validate(d)
+
+    def test_no_cmd(self):
+        d = {}
+        Stage.validate(d)
+
+    def test_cmd_str(self):
+        d = {Stage.PARAM_CMD: 'cmd'}
+        Stage.validate(d)
+
+
+class TestSchemaDepsOuts(TestSchema):
+    def test_object(self):
+        d = {Stage.PARAM_DEPS: {}}
+        self._validate_fail(d)
+
+        d = {Stage.PARAM_OUTS: {}}
+        self._validate_fail(d)
+
+    def test_none(self):
+        d = {Stage.PARAM_DEPS: None}
+        Stage.validate(d)
+
+        d = {Stage.PARAM_OUTS: None}
+        Stage.validate(d)
+
+    def test_empty_list(self):
+        d = {Stage.PARAM_DEPS: []}
+        Stage.validate(d)
+
+        d = {Stage.PARAM_OUTS: []}
+        Stage.validate(d)
+
+    def test_list(self):
+        l = [{Output.PARAM_PATH: 'foo', Output.PARAM_MD5: '123'},
+             {Output.PARAM_PATH: 'bar', Output.PARAM_MD5: None},
+             {Output.PARAM_PATH: 'baz'}]
+        d = {Stage.PARAM_DEPS: l}
+        Stage.validate(d)
+
+        l[0][Output.PARAM_CACHE] = True
+        l[1][Output.PARAM_CACHE] = False
+        d = {Stage.PARAM_OUTS: l}
+        Stage.validate(d)