Support auto id and paths well for `$` operator

2 files changed, 54 insertions, 4 deletions

diff --git a/jsonpath_rw/jsonpath.py b/jsonpath_rw/jsonpath.py
index ca2cc48..563e147 100644
--- a/jsonpath_rw/jsonpath.py
+++ b/jsonpath_rw/jsonpath.py
@@ -33,10 +33,12 @@ class JSONPath(object):
         """
         Equivalent to Child(self, next) but with some canonicalization
         """
-        if isinstance(self, This):
+        if isinstance(self, This) or isinstance(self, Root):
             return child
         elif isinstance(child, This):
             return self
+        elif isinstance(child, Root):
+            return child
         else:
             return Child(self, child)
 
@@ -165,10 +167,10 @@ class Root(JSONPath):
 
     def find(self, data):
         if not isinstance(data, DatumInContext):
-            return [DatumInContext(data, path=This(), context=None)]
+            return [DatumInContext(data, path=Root(), context=None)]
         else:
             if data.context is None:
-                return data
+                return [DatumInContext(data.value, context=None, path=Root())]
             else:
                 return Root().find(data.context)
 
@@ -190,7 +192,7 @@ class This(JSONPath):
     """
 
     def find(self, datum):
-        return DatumInContext.wrap(datum)
+        return [DatumInContext.wrap(datum)]
 
     def update(self, data, val):
         return val
diff --git a/tests/test_jsonpath.py b/tests/test_jsonpath.py
index 9f341d1..fd7cf32 100644
--- a/tests/test_jsonpath.py
+++ b/tests/test_jsonpath.py
@@ -108,6 +108,22 @@ class TestJsonPath(unittest.TestCase):
         jsonpath.auto_id_field = 'id'
         self.check_cases([ ('*', {'foo': 1, 'baz': 2}, set([1, 2, '@'])) ])
 
+    def test_root_value(self):
+        jsonpath.auto_id_field = None
+        self.check_cases([ 
+            ('$', {'foo': 'baz'}, [{'foo':'baz'}]),
+            ('foo.$', {'foo': 'baz'}, [{'foo':'baz'}]),
+            ('foo.$.foo', {'foo': 'baz'}, ['baz']),
+        ])
+
+    def test_this_value(self):
+        jsonpath.auto_id_field = None
+        self.check_cases([ 
+            ('@', {'foo': 'baz'}, [{'foo':'baz'}]),
+            ('foo.@', {'foo': 'baz'}, ['baz']),
+            ('foo.@.baz', {'foo': {'baz': 3}}, [3]),
+        ])
+
     def test_index_value(self):
         self.check_cases([
             ('[0]', [42], [42]),
@@ -167,6 +183,22 @@ class TestJsonPath(unittest.TestCase):
         jsonpath.auto_id_field = 'id'
         self.check_paths([ ('*', {'foo': 1, 'baz': 2}, set(['foo', 'baz', 'id'])) ])
 
+    def test_root_paths(self):
+        jsonpath.auto_id_field = None
+        self.check_paths([ 
+            ('$', {'foo': 'baz'}, ['$']),
+            ('foo.$', {'foo': 'baz'}, ['$']),
+            ('foo.$.foo', {'foo': 'baz'}, ['foo']),
+        ])
+
+    def test_this_paths(self):
+        jsonpath.auto_id_field = None
+        self.check_paths([ 
+            ('@', {'foo': 'baz'}, ['@']),
+            ('foo.@', {'foo': 'baz'}, ['foo']),
+            ('foo.@.baz', {'foo': {'baz': 3}}, ['foo.baz']),
+        ])
+
     def test_index_paths(self):
         self.check_paths([('[0]', [42], ['[0]']),
                           ('[2]', [34, 65, 29, 59], ['[2]'])])
@@ -197,6 +229,22 @@ class TestJsonPath(unittest.TestCase):
                              'baz': 2},
                              set(['1', 'baz'])) ])
 
+    def test_root_auto_id(self):
+        jsonpath.auto_id_field = 'id'
+        self.check_cases([ 
+            ('$.id', {'foo': 'baz'}, ['$']), # This is a wonky case that is not that interesting
+            ('foo.$.id', {'foo': 'baz', 'id': 'bizzle'}, ['bizzle']), 
+            ('foo.$.baz.id', {'foo': 4, 'baz': 3}, ['baz']),
+        ])
+
+    def test_this_auto_id(self):
+        jsonpath.auto_id_field = 'id'
+        self.check_cases([ 
+            ('id', {'foo': 'baz'}, ['@']), # This is, again, a wonky case that is not that interesting
+            ('foo.@.id', {'foo': 'baz'}, ['foo']),
+            ('foo.@.baz.id', {'foo': {'baz': 3}}, ['foo.baz']),
+        ])
+
     def test_index_auto_id(self):
         jsonpath.auto_id_field = "id"
         self.check_cases([('[0].id', [42], ['[0]']),