Support for set_pointer and indexing arbitrary objects via __getitem__/__setitem__

3 files changed, 185 insertions, 7 deletions

diff --git a/doc/tutorial.rst b/doc/tutorial.rst
index 4cdb75a..4519b29 100644
--- a/doc/tutorial.rst
+++ b/doc/tutorial.rst
@@ -7,7 +7,7 @@ method is basically a deep ``get``.
 
 .. code-block:: python
 
-    >>> import jsonpointer
+    >>> from jsonpointer import resolve_pointer
     >>> obj = {"foo": {"anArray": [ {"prop": 44}], "another prop": {"baz": "A string" }}}
 
     >>> resolve_pointer(obj, '') == obj
@@ -29,6 +29,32 @@ method is basically a deep ``get``.
     True
 
 
+The ``set_pointer`` method allows modifying a portion of an object using
+JSON pointer notation:
+
+.. code-block:: python
+
+    >>> from jsonpointer import set_pointer
+    >>> obj = {"foo": {"anArray": [ {"prop": 44}], "another prop": {"baz": "A string" }}}
+
+    >>> set_pointer(obj, '/foo/anArray/0/prop', 55)
+    {'foo': {'another prop': {'baz': 'A string'}, 'anArray': [{'prop': 55}]}}
+
+    >>> obj
+    {'foo': {'another prop': {'baz': 'A string'}, 'anArray': [{'prop': 55}]}}
+
+By default ``set_pointer`` modifies the original object.  Pass ``inplace=False``
+to create a copy and modify the copy instead:
+
+    >>> from jsonpointer import set_pointer
+    >>> obj = {"foo": {"anArray": [ {"prop": 44}], "another prop": {"baz": "A string" }}}
+
+    >>> set_pointer(obj, '/foo/anArray/0/prop', 55, inplace=False)
+    {'foo': {'another prop': {'baz': 'A string'}, 'anArray': [{'prop': 55}]}}
+
+    >>> obj
+    {'foo': {'another prop': {'baz': 'A string'}, 'anArray': [{'prop': 44}]}}
+
 The ``JsonPointer`` class wraps a (string) path and can be used to access the
 same path on several objects.
 
diff --git a/jsonpointer.py b/jsonpointer.py
index fe7914a..648a377 100644
--- a/jsonpointer.py
+++ b/jsonpointer.py
@@ -50,6 +50,7 @@ except ImportError: # Python 3
 
 from itertools import tee
 import re
+import copy
 
 
 # array indices must not contain leading zeros, signs, spaces, decimals, etc
@@ -104,6 +105,28 @@ def resolve_pointer(doc, pointer, default=_nothing):
     pointer = JsonPointer(pointer)
     return pointer.resolve(doc, default)
 
+def set_pointer(doc, pointer, value, inplace=True):
+    """
+    Resolves pointer against doc and sets the value of the target within doc.
+
+    With inplace set to true, doc is modified as long as pointer is not the
+    root.
+
+    >>> obj = {"foo": {"anArray": [ {"prop": 44}], "another prop": {"baz": "A string" }}}
+
+    >>> set_pointer(obj, '/foo/anArray/0/prop', 55) == \
+    {'foo': {'another prop': {'baz': 'A string'}, 'anArray': [{'prop': 55}]}}
+    True
+
+    >>> set_pointer(obj, '/foo/yet%20another%20prop', 'added prop') == \
+    {'foo': {'another prop': {'baz': 'A string'}, 'yet another prop': 'added prop', 'anArray': [{'prop': 55}]}}
+    True
+
+    """
+
+    pointer = JsonPointer(pointer)
+    return pointer.set(doc, value, inplace)
+
 
 class JsonPointer(object):
     """ A JSON Pointer that can reference parts of an JSON document """
@@ -149,6 +172,21 @@ class JsonPointer(object):
 
     get = resolve
 
+    def set(self, doc, value, inplace=True):
+        """ Resolve the pointer against the doc and replace the target with value. """
+
+        if len(self.parts) == 0:
+            if inplace:
+                raise JsonPointerException('cannot set root in place')
+            return value
+
+        if not inplace:
+            doc = copy.deepcopy(doc)
+
+        (parent, part) = self.to_last(doc)
+
+        parent[part] = value
+        return doc
 
     def get_part(self, doc, part):
         """ Returns the next step in the correct type """
@@ -166,8 +204,13 @@ class JsonPointer(object):
 
             return int(part)
 
+        elif hasattr(doc, '__getitem__'):
+            # Allow indexing via ducktyping if the target has defined __getitem__
+            return part
+
         else:
-            raise JsonPointerException("Unknown document type '%s'" % (doc.__class__,))
+            raise JsonPointerException("Document '%s' does not support indexing, "
+                                       "must be dict/list or support __getitem__" % type(doc))
 
 
     def walk(self, doc, part):
@@ -175,9 +218,7 @@ class JsonPointer(object):
 
         part = self.get_part(doc, part)
 
-        # type is already checked in get_part, so we assert here
-        # for consistency
-        assert type(doc) in (dict, list), "invalid document type %s" (type(doc),)
+        assert (type(doc) in (dict, list) or hasattr(doc, '__getitem__')), "invalid document type %s" (type(doc))
 
         if isinstance(doc, dict):
             try:
@@ -197,9 +238,12 @@ class JsonPointer(object):
             except IndexError:
                 raise JsonPointerException("index '%s' is out of bounds" % (part, ))
 
+        else:
+            # Object supports __getitem__, assume custom indexing
+            return doc[part]
 
     def contains(self, ptr):
-        """" Returns True if self contains the given ptr """
+        """ Returns True if self contains the given ptr """
         return len(self.parts) > len(ptr.parts) and \
              self.parts[:len(ptr.parts)] == ptr.parts
 
diff --git a/tests.py b/tests.py
index b75723a..9dc5f1b 100755
--- a/tests.py
+++ b/tests.py
@@ -4,8 +4,9 @@
 import doctest
 import unittest
 import sys
+import copy
 from jsonpointer import resolve_pointer, EndOfList, JsonPointerException, \
-         JsonPointer
+         JsonPointer, set_pointer
 
 class SpecificationTests(unittest.TestCase):
     """ Tests all examples from the JSON Pointer specification """
@@ -110,11 +111,118 @@ class ToLastTests(unittest.TestCase):
         self.assertEqual(nxt, 'b')
 
 
+class SetTests(unittest.TestCase):
+
+    def test_set(self):
+        doc =   {
+            "foo": ["bar", "baz"],
+            "": 0,
+            "a/b": 1,
+            "c%d": 2,
+            "e^f": 3,
+            "g|h": 4,
+            "i\\j": 5,
+            "k\"l": 6,
+            " ": 7,
+            "m~n": 8
+        }
+        origdoc = copy.deepcopy(doc)
+
+        # inplace=False
+        newdoc = set_pointer(doc, "/foo/1", "cod", inplace=False)
+        self.assertEqual(resolve_pointer(newdoc, "/foo/1"), "cod")
+
+        newdoc = set_pointer(doc, "/", 9, inplace=False)
+        self.assertEqual(resolve_pointer(newdoc, "/"), 9)
+
+        newdoc = set_pointer(doc, "/fud", {}, inplace=False)
+        newdoc = set_pointer(newdoc, "/fud/gaw", [1, 2, 3], inplace=False)
+        self.assertEqual(resolve_pointer(newdoc, "/fud"), {'gaw' : [1, 2, 3]})
+
+        newdoc = set_pointer(doc, "", 9, inplace=False)
+        self.assertEqual(newdoc, 9)
+
+        self.assertEqual(doc, origdoc)
+
+        # inplace=True
+        set_pointer(doc, "/foo/1", "cod")
+        self.assertEqual(resolve_pointer(doc, "/foo/1"), "cod")
+
+        set_pointer(doc, "/", 9)
+        self.assertEqual(resolve_pointer(doc, "/"), 9)
+
+        self.assertRaises(JsonPointerException, set_pointer, doc, "/fud/gaw", 9)
+
+        set_pointer(doc, "/fud", {})
+        set_pointer(doc, "/fud/gaw", [1, 2, 3] )
+        self.assertEqual(resolve_pointer(doc, "/fud"), {'gaw' : [1, 2, 3]})
+
+        self.assertRaises(JsonPointerException, set_pointer, doc, "", 9)
+
+class AltTypesTests(unittest.TestCase):
+
+    def test_alttypes(self):
+        JsonPointer.alttypes = True
+
+        class Node(object):
+            def __init__(self, name, parent=None):
+                self.name = name
+                self.parent = parent
+                self.left = None
+                self.right = None
+
+            def set_left(self, node):
+                node.parent = self
+                self.left = node
+
+            def set_right(self, node):
+                node.parent = self
+                self.right = node
+
+            def __getitem__(self, key):
+                if key == 'left':
+                    return self.left
+                if key == 'right':
+                    return self.right
+
+                raise KeyError("Only left and right supported")
+
+            def __setitem__(self, key, val):
+                if key == 'left':
+                    return self.set_left(val)
+                if key == 'right':
+                    return self.set_right(val)
+
+                raise KeyError("Only left and right supported: %s" % key)
+
+
+        root = Node('root')
+        root.set_left(Node('a'))
+        root.left.set_left(Node('aa'))
+        root.left.set_right(Node('ab'))
+        root.set_right(Node('b'))
+        root.right.set_left(Node('ba'))
+        root.right.set_right(Node('bb'))
+
+        self.assertEqual(resolve_pointer(root, '/left').name, 'a')
+        self.assertEqual(resolve_pointer(root, '/left/right').name, 'ab')
+        self.assertEqual(resolve_pointer(root, '/right').name, 'b')
+        self.assertEqual(resolve_pointer(root, '/right/left').name, 'ba')
+
+        newroot = set_pointer(root, '/left/right', Node('AB'), inplace=False)
+        self.assertEqual(resolve_pointer(root, '/left/right').name, 'ab')
+        self.assertEqual(resolve_pointer(newroot, '/left/right').name, 'AB')
+
+        set_pointer(root, '/left/right', Node('AB'))
+        self.assertEqual(resolve_pointer(root, '/left/right').name, 'AB')
+
 suite = unittest.TestSuite()
 suite.addTest(unittest.makeSuite(SpecificationTests))
 suite.addTest(unittest.makeSuite(ComparisonTests))
 suite.addTest(unittest.makeSuite(WrongInputTests))
 suite.addTest(unittest.makeSuite(ToLastTests))
+suite.addTest(unittest.makeSuite(SetTests))
+suite.addTest(unittest.makeSuite(AltTypesTests))
 
 modules = ['jsonpointer']