Adjustments and tests for jsonpath.py command line script

7 files changed, 132 insertions, 34 deletions

diff --git a/jsonpath_rw/bin/__init__.py b/jsonpath_rw/bin/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/jsonpath_rw/bin/__init__.py
diff --git a/jsonpath_rw/bin/jsonpath.py b/jsonpath_rw/bin/jsonpath.py
index 364615f..853b8b7 100755
--- a/jsonpath_rw/bin/jsonpath.py
+++ b/jsonpath_rw/bin/jsonpath.py
@@ -5,47 +5,67 @@
 # terms of the Do What The Fuck You Want To Public License, Version 2,
 # as published by Sam Hocevar. See the COPYING file for more details.
 
-from jsonpath_rw import parse
+# Use modern Python
+from __future__ import unicode_literals, print_function, absolute_import
+
+# Standard Library imports
 import json
 import sys
 import glob
-if len(sys.argv) < 2:
-    print("""usage: jsonpath.py expression [files]
+import argparse
+
+# JsonPath-RW imports
+from jsonpath_rw import parse
 
-The expression is JSONPath and can be:
+def find_matches_for_file(expr, f):
+    return expr.find(json.load(f))
 
-    atomics:
-        $              - root object
-        `this`         - current object
+def print_matches(matches):
+    print('\n'.join(['{0}'.format(match.value) for match in matches]))
 
-    operators:
-        path1.path2    - same as xpath /
-        path1|path2    - union
-        path1..path2   - somewhere in between
 
-    fiels:
-        fieldname       - field with name
-        *               - any field
-        [_start_?:_end_?] - array slice
-        [*]             - any array index
-""")
-    sys.exit(1)
+def main(*argv):
+    parser = argparse.ArgumentParser(
+        description='Search JSON files (or stdin) according to a JSONPath expression.',
+        formatter_class=argparse.RawTextHelpFormatter,
+        epilog="""
+        Quick JSONPath reference (see more at https://github.com/kennknowles/python-jsonpath-rw)
 
-expr = parse(sys.argv[1])
+        atomics:
+            $              - root object
+            `this`         - current object
 
-def find_matches_for_file(f):
-    return [unicode(match.value) for match in expr.find(json.load(f))]
+        operators:
+            path1.path2    - same as xpath /
+            path1|path2    - union
+            path1..path2   - somewhere in between
 
-def print_matches(matches):
-    print(u"\n".join(matches).encode("utf-8"))
-
-if len(sys.argv) < 3:
-    # stdin mode
-    print_matches(find_matches_for_file(sys.stdin))
-else:
-    # file paths mode
-    for pattern in sys.argv[2:]:
-        for filename in glob.glob(pattern):
-            with open(filename) as f:
-                print_matches(find_matches_for_file(f))
+        fields:
+            fieldname       - field with name
+            *               - any field
+            [_start_?:_end_?] - array slice
+            [*]             - any array index
+    """)
+
+
+
+    parser.add_argument('expression', help='A JSONPath expression.')
+    parser.add_argument('files', metavar='file', nargs='*', help='Files to search (if none, searches stdin)')
+
+    args = parser.parse_args(argv[1:])
+
+    expr = parse(args.expression)
+    glob_patterns = args.files
+
+    if len(glob_patterns) == 0:
+        # stdin mode
+        print_matches(find_matches_for_file(expr, sys.stdin))
+    else:
+        # file paths mode
+        for pattern in glob_patterns:
+            for filename in glob.glob(pattern):
+                with open(filename) as f:
+                    print_matches(find_matches_for_file(expr, f))
 
+def entry_point():
+    main(*sys.argv)
diff --git a/setup.py b/setup.py
index c37c0f6..8d896ef 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,9 @@ setuptools.setup(
     license='Apache 2.0',
     long_description=io.open('README.rst', encoding='utf-8').read(),
     packages = ['jsonpath_rw'],
-    scripts = ['jsonpath_rw/bin/jsonpath.py'],
+    entry_points = {
+        'console_scripts':  ['jsonpath.py = jsonpath_rw.bin.jsonpath:entry_point'],
+    },
     test_suite = 'tests',
     install_requires = [ 'ply', 'decorator', 'six' ],
     classifiers = [
diff --git a/tests/bin/__init__.py b/tests/bin/__init__.py
new file mode 100644
index 0000000..0c0485e
--- /dev/null
+++ b/tests/bin/__init__.py
@@ -0,0 +1,2 @@
+# Use modern python
+from __future__ import absolute_import, print_function, unicode_literals
diff --git a/tests/bin/test1.json b/tests/bin/test1.json
new file mode 100644
index 0000000..726a29f
--- /dev/null
+++ b/tests/bin/test1.json
@@ -0,0 +1,8 @@
+{
+    "foo": {
+        "baz": 1,
+        "bizzle": {
+            "baz": 2
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/bin/test2.json b/tests/bin/test2.json
new file mode 100644
index 0000000..afc7561
--- /dev/null
+++ b/tests/bin/test2.json
@@ -0,0 +1,10 @@
+{
+    "foo": {
+        "foo": {
+            "baz": 3,
+            "merp": {
+                "baz": 4
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/bin/test_jsonpath.py b/tests/bin/test_jsonpath.py
new file mode 100644
index 0000000..3724044
--- /dev/null
+++ b/tests/bin/test_jsonpath.py
@@ -0,0 +1,56 @@
+# Use modern Python
+from __future__ import unicode_literals, print_function, absolute_import, division, generators, nested_scopes
+
+# Standard library imports
+import unittest
+import logging
+import io
+import sys
+import os
+import json
+
+from jsonpath_rw.bin.jsonpath import main
+
+class TestJsonPathScript(unittest.TestCase):
+    """
+    Tests for the jsonpath.py command line interface.
+    """
+    
+    @classmethod
+    def setup_class(cls):
+        logging.basicConfig()
+
+    def setUp(self):
+        self.input = io.StringIO()
+        self.output = io.StringIO()
+        self.saved_stdout = sys.stdout
+        self.saved_stdin = sys.stdin
+        sys.stdout = self.output
+        sys.stdin = self.input
+
+    def tearDown(self):
+        self.output.close()
+        self.input.close()
+        sys.stdout = self.saved_stdout
+        sys.stdin = self.saved_stdin
+
+    def test_stdin_mode(self):
+        # 'format' is a benign Python 2/3 way of ensuring it is a text type rather than binary
+        self.input.write('{0}'.format(json.dumps({
+            'foo': {
+                'baz': 1,
+                'bizzle': {
+                    'baz': 2
+                }
+            }
+        })))
+        self.input.seek(0)
+        main('jsonpath.py', 'foo..baz')
+        self.assertEqual(self.output.getvalue(), '1\n2\n')
+
+    def test_filename_mode(self):
+        test1 = os.path.join(os.path.dirname(__file__), 'test1.json')
+        test2 = os.path.join(os.path.dirname(__file__), 'test2.json')
+        main('jsonpath.py', 'foo..baz', test1, test2)
+        self.assertEqual(self.output.getvalue(), '1\n2\n3\n4\n')
+