From 56d64bd37724415d55316cf9f302fa6e7d0c19c6 Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Thu, 18 Apr 2019 17:24:20 +0700
Subject: [PATCH] dvc: add meta key to stages to store user data

---
 dvc/stage.py             |  3 +++
 tests/func/test_stage.py | 16 ++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/dvc/stage.py b/dvc/stage.py
index 115bbd34bc..05691de8ae 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -125,6 +125,7 @@ class Stage(object):
     PARAM_DEPS = "deps"
     PARAM_OUTS = "outs"
     PARAM_LOCKED = "locked"
+    PARAM_META = "meta"
 
     SCHEMA = {
         Optional(PARAM_MD5): Or(str, None),
@@ -133,6 +134,7 @@ class Stage(object):
         Optional(PARAM_DEPS): Or(And(list, Schema([dependency.SCHEMA])), None),
         Optional(PARAM_OUTS): Or(And(list, Schema([output.SCHEMA])), None),
         Optional(PARAM_LOCKED): bool,
+        Optional(PARAM_META): object,
     }
 
     TAG_REGEX = r"^(?P<path>.*)@(?P<tag>[^\\/@:]*)$"
@@ -612,6 +614,7 @@ def dumpd(self):
                 Stage.PARAM_LOCKED: self.locked,
                 Stage.PARAM_DEPS: [d.dumpd() for d in self.deps],
                 Stage.PARAM_OUTS: [o.dumpd() for o in self.outs],
+                Stage.PARAM_META: self._state.get("meta"),
             }.items()
             if value
         }
diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py
index 1095eefe25..598ee83248 100644
--- a/tests/func/test_stage.py
+++ b/tests/func/test_stage.py
@@ -167,3 +167,19 @@ def test_md5_ignores_comments(repo_dir, dvc):
 
     new_stage = Stage.load(dvc, stage.path)
     assert not new_stage.changed_md5()
+
+
+def test_meta_is_preserved(dvc):
+    stage, = dvc.add("foo")
+
+    # Add meta to stage file
+    data = load_stage_file(stage.path)
+    data["meta"] = {"custom_key": 42}
+    dump_stage_file(stage.path, data)
+
+    # Loading and dumping to test that it works and meta is retained
+    new_stage = Stage.load(dvc, stage.path)
+    new_stage.dump()
+
+    new_data = load_stage_file(stage.path)
+    assert new_data["meta"] == data["meta"]