Enable non-ascii support when dumping YAML; expand test suite

pemensik · May 2, 2017 · fc2fe20 · fc2fe20
1 parent 4fbfbb9
commit fc2fe20
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/test/test.py b/test/test.py
@@ -8,16 +8,34 @@
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 from yq import main
 
+USING_PYTHON2 = True if sys.version_info < (3, 0) else False
+
 class TestYq(unittest.TestCase):
-    def test_basic_statements(self):
+    def run_yq(self, input_data, args, expect_exit_code=os.EX_OK):
+        stdin, stdout = sys.stdin, sys.stdout
         try:
-            orig_stdin = sys.stdin
-            sys.stdin = io.StringIO("foo:\n bar: 1\n baz: {bat: 3}")
-            main([".foo.baz.bat"])
+            sys.stdin = io.StringIO(input_data)
+            sys.stdout = io.BytesIO() if USING_PYTHON2 else io.StringIO()
+            main(args)
         except SystemExit as e:
-            self.assertEqual(e.code, os.EX_OK)
+            self.assertEqual(e.code, expect_exit_code)
         finally:
-            sys.stdin = orig_stdin
+            result = sys.stdout.getvalue()
+            if USING_PYTHON2:
+                result = result.decode("utf-8")
+            sys.stdin, sys.stdout = stdin, stdout
+        return result
+
+    def test_yq(self):
+        self.assertEqual(self.run_yq("{}", ["."]), "")
+        self.assertEqual(self.run_yq("foo:\n bar: 1\n baz: {bat: 3}", [".foo.baz.bat"]), "")
+        self.assertEqual(self.run_yq("[1, 2, 3]", ["--yaml-output", "-M", "."]), "- 1\n- 2\n- 3\n")
+        self.assertEqual(self.run_yq("foo:\n bar: 1\n baz: {bat: 3}", ["-y", ".foo.baz.bat"]), "3\n...\n")
+        self.assertEqual(self.run_yq("[aaaaaaaaaa bbb]", ["-y", "."]), "- aaaaaaaaaa bbb\n")
+        self.assertEqual(self.run_yq("[aaaaaaaaaa bbb]", ["-y", "-w8", "."]), "- aaaaaaaaaa\n  bbb\n")
+        self.assertEqual(self.run_yq('{"понедельник": 1}', ['.["понедельник"]']), "")
+        self.assertEqual(self.run_yq('{"понедельник": 1}', ["-y", '.["понедельник"]']), "1\n...\n")
+        self.assertEqual(self.run_yq("- понедельник\n- вторник\n", ["-y", "."]), "- понедельник\n- вторник\n")
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/yq/__init__.py b/yq/__init__.py
@@ -61,7 +61,8 @@ def main(args=None):
         if args.yaml_output:
             out, err = jq.communicate(json.dumps(yaml.load(input_stream, Loader=OrderedLoader)))
             out = json.loads(out, object_pairs_hook=OrderedDict)
-            yaml.dump(out, Dumper=OrderedDumper, stream=sys.stdout, default_flow_style=False, width=args.width)
+            yaml.dump(out, stream=sys.stdout, Dumper=OrderedDumper, width=args.width,
+                      allow_unicode=True, default_flow_style=False)
         else:
             json.dump(yaml.load(input_stream, OrderedLoader), jq.stdin)
             jq.stdin.close()