adding the recursive schema extension and literal key interpretation for pull request #227

2 files changed, 62 insertions, 2 deletions

diff --git a/voluptuous/schema_builder.py b/voluptuous/schema_builder.py
index 40c59bf..053d951 100644
--- a/voluptuous/schema_builder.py
+++ b/voluptuous/schema_builder.py
@@ -608,8 +608,43 @@ class Schema(object):
         assert type(self.schema) == dict and type(schema) == dict, 'Both schemas must be dictionary-based'
 
         result = self.schema.copy()
-        result.update(schema)
 
+        # returns the key that may have been passed as arugment to Marker constructor
+        def key_literal(key):
+            return (key.schema if isinstance(key, Marker) else key)
+
+        # build a map that takes the key literals to the needed objects
+        # literal -> Required|Optional|literal
+        result_key_map = dict((key_literal(key), key) for key in result)
+
+        # for each item in the extension schema, replace duplicates
+        # or add new keys
+        for key, value in iteritems(schema):
+            
+            # if the key is already in the dictionary, we need to replace it
+            # transform key to literal before checking presence
+            if key_literal(key) in result_key_map:
+
+                result_key = result_key_map[key_literal(key)]
+                result_value = result[result_key]
+
+                # if both are dictionaries, we need to extend recursively
+                # create the new extended sub schema, then remove the old key and add the new one
+                if type(result_value) == dict and type(value) == dict:
+                    new_value = Schema(result_value).extend(value).schema
+                    del result[result_key]
+                    result[key] = new_value
+                # one or the other or both are not sub-schemas, simple replacement is fine
+                # remove old key and add new one
+                else:
+                    del result[result_key]
+                    result[key] = value
+
+            # key is new and can simply be added
+            else:
+                result[key] = value
+
+        # recompile and send old object
         result_required = (required if required is not None else self.required)
         result_extra = (extra if extra is not None else self.extra)
         return Schema(result, required=result_required, extra=result_extra)
diff --git a/voluptuous/tests/tests.py b/voluptuous/tests/tests.py
index 01b7c85..8759b13 100644
--- a/voluptuous/tests/tests.py
+++ b/voluptuous/tests/tests.py
@@ -2,7 +2,7 @@ import copy
 from nose.tools import assert_equal, assert_raises, assert_true
 
 from voluptuous import (
-    Schema, Required, Extra, Invalid, In, Remove, Literal,
+    Schema, Required, Optional, Extra, Invalid, In, Remove, Literal,
     Url, MultipleInvalid, LiteralInvalid, NotIn, Match, Email,
     Replace, Range, Coerce, All, Any, Length, FqdnUrl, ALLOW_EXTRA, PREVENT_EXTRA,
     validate, ExactSequence, Equal, Unordered, Number
@@ -337,6 +337,31 @@ def test_schema_extend_overrides():
     assert extended.extra == ALLOW_EXTRA
 
 
+def test_schema_extend_key_swap():
+    """Verify that Schema.extend can replace keys, even when different markers are used"""
+
+    base = Schema({Optional('a'): int})
+    extension = {Required('a'): int}
+    extended = base.extend(extension)
+
+    assert_equal(len(base.schema), 1)
+    assert_true(isinstance(list(base.schema)[0], Optional))
+    assert_equal(len(extended.schema), 1)
+    assert_true((list(extended.schema)[0], Required))
+
+
+def test_subschema_extension():
+    """Verify that Schema.extend adds and replaces keys in a subschema"""
+
+    base = Schema({'a': {'b': int, 'c': float}})
+    extension = {'d': str, 'a': {'b': str, 'e': int}}
+    extended = base.extend(extension)
+
+    assert_equal(base.schema, {'a': {'b': int, 'c': float}})
+    assert_equal(extension, {'d': str, 'a': {'b': str, 'e': int}})
+    assert_equal(extended.schema, {'a': {'b': str, 'c': float, 'e': int}, 'd': str})
+
+
 def test_repr():
     """Verify that __repr__ returns valid Python expressions"""
     match = Match('a pattern', msg='message')