Allow editing toml in place (#171)

2 files changed, 10 insertions, 3 deletions

diff --git a/test/test.py b/test/test.py
index 759e9f4..b674c1a 100755
--- a/test/test.py
+++ b/test/test.py
@@ -185,7 +185,7 @@ class TestYq(unittest.TestCase):
         with io.open(cfn_filename) as fh:
             self.assertEqual(self.run_yq("", ["-Y", ".", cfn_filename]), fh.read())
 
-    def test_in_place(self):
+    def test_in_place_yaml(self):
         with tempfile.NamedTemporaryFile() as tf, tempfile.NamedTemporaryFile() as tf2:
             tf.write(b"- foo\n- bar\n")
             tf.seek(0)
@@ -202,6 +202,13 @@ class TestYq(unittest.TestCase):
             self.assertEqual(tf.read(), b"foo\n...\n")
             self.assertEqual(tf2.read(), b"foo\n...\n")
 
+    def test_in_place_toml(self):
+        with tempfile.NamedTemporaryFile() as tf:
+            tf.write(b'[GLOBAL]\nversion="1.0.0"\n')
+            tf.seek(0)
+            self.run_yq("", ["-i", "-t", '.GLOBAL.version="1.0.1"', tf.name], input_format="toml")
+            self.assertEqual(tf.read(), b'[GLOBAL]\nversion = "1.0.1"\n')
+
     def test_explicit_doc_markers(self):
         test_doc = os.path.join(os.path.dirname(__file__), "doc.yml")
         self.assertTrue(self.run_yq("", ["-y", ".", test_doc]).startswith("yaml_struct"))
diff --git a/yq/__init__.py b/yq/__init__.py
index 7669610..9545b60 100644
--- a/yq/__init__.py
+++ b/yq/__init__.py
@@ -131,8 +131,8 @@ def cli(args=None, input_format="yaml", program_name="yq"):
 
     yq_args = dict(input_format=input_format, program_name=program_name, jq_args=jq_args, **vars(args))
     if in_place:
-        if args.output_format not in {"yaml", "annotated_yaml"}:
-            sys.exit("{}: -i/--in-place can only be used with -y/-Y".format(program_name))
+        if args.output_format not in {"yaml", "annotated_yaml", "toml"}:
+            sys.exit("{}: -i/--in-place can only be used with -y/-Y/-t".format(program_name))
         input_streams = yq_args.pop("input_streams")
         if len(input_streams) == 1 and input_streams[0].name == "<stdin>":
             msg = "{}: -i/--in-place can only be used with filename arguments, not on standard input"