Repository: jsonpickle/jsonpickle
Pull request #517 (2024-07-13 02:20:51)

Define _jsonpickle_exclude to omit attributes from encoding
Initial PR to talk about concept.


Theelx: The test failures are from Ruff, from a previous commit having an unused variable. I'll make a PR to fix that now.
Theelx: @dwagon If you could rebase this on the main branch which'll fix the tests, and if @davvid approves this, then I think this'll be ready to merge!

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index efd9fd5..23c8410 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -3,6 +3,9 @@ v3.3.0
     * The unpickler was updated to avoid using ``eval``, which helps improve its
       security. Users can still pass ``safe=False`` to ``decode`` to use the old
       behavior, though this is not recommended. (+513)
+    * Objects can now exclude specific attributes from pickling by providing a
+      ``_jsonpickle_exclude`` class or instance attribute. This attribute should contain
+      the list of attribute names to exclude when pickling the object.
 
 v3.2.2
 ======
diff --git a/docs/api.rst b/docs/api.rst
index c90f9c7..b126ec2 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -91,6 +91,12 @@ protocol for representing object instances.
    :meth:`__getstate__` and :meth:`__setstate__`, the state object needn't be a
    dictionary and these methods can do what they want.
 
+.. py:attribute:: object._jsonpickle_exclude = set()
+
+    Classes can specify attributes that they want to exclude from being pickled
+    by defining an attribute named `_jsonpickle_exclude`. This set should contain
+    the names of attributes of that class to exclude from being pickled.
+
 
 :mod:`jsonpickle.handlers` -- Custom Serialization Handlers
 -----------------------------------------------------------
diff --git a/jsonpickle/pickler.py b/jsonpickle/pickler.py
index 9be34c3..75a8e92 100644
--- a/jsonpickle/pickler.py
+++ b/jsonpickle/pickler.py
@@ -539,6 +539,7 @@ class Pickler(object):
         has_getnewargs_ex = util.has_method(obj, '__getnewargs_ex__')
         has_getinitargs = util.has_method(obj, '__getinitargs__')
         has_reduce, has_reduce_ex = util.has_reduce(obj)
+        exclude = set(getattr(obj, '_jsonpickle_exclude', ()))
 
         # Support objects with __getstate__(); this ensures that
         # both __setstate__() and __getstate__() are implemented
@@ -675,7 +676,7 @@ class Pickler(object):
             return data
 
         if util.is_dictionary_subclass(obj):
-            self._flatten_dict_obj(obj, data)
+            self._flatten_dict_obj(obj, data, exclude=exclude)
             return data
 
         if util.is_sequence_subclass(obj):
@@ -693,7 +694,7 @@ class Pickler(object):
 
             # hack for zope persistent objects; this unghostifies the object
             getattr(obj, '_', None)
-            return self._flatten_dict_obj(obj.__dict__, data)
+            return self._flatten_dict_obj(obj.__dict__, data, exclude=exclude)
 
         if has_slots:
             return self._flatten_newstyle_with_slots(obj, data)
@@ -770,7 +771,7 @@ class Pickler(object):
         data[k] = self._flatten(v)
         return data
 
-    def _flatten_dict_obj(self, obj, data=None):
+    def _flatten_dict_obj(self, obj, data=None, exclude=()):
         """Recursively call flatten() and return json-friendly dict"""
         if data is None:
             data = obj.__class__()
@@ -780,17 +781,17 @@ class Pickler(object):
         if self.keys:
             # Phase 1: serialize regular objects, ignore fancy keys.
             flatten = self._flatten_string_key_value_pair
-            for k, v in util.items(obj):
+            for k, v in util.items(obj, exclude=exclude):
                 flatten(k, v, data)
 
             # Phase 2: serialize non-string keys.
             flatten = self._flatten_non_string_key_value_pair
-            for k, v in util.items(obj):
+            for k, v in util.items(obj, exclude=exclude):
                 flatten(k, v, data)
         else:
             # If we have string keys only then we only need a single pass.
             flatten = self._flatten_key_value_pair
-            for k, v in util.items(obj):
+            for k, v in util.items(obj, exclude=exclude):
                 flatten(k, v, data)
 
         # the collections.defaultdict protocol
@@ -817,7 +818,7 @@ class Pickler(object):
         if hasattr(obj, '__dict__') and self.unpicklable and obj != obj.__dict__:
             if self._mkref(obj.__dict__):
                 dict_data = {}
-                self._flatten_dict_obj(obj.__dict__, dict_data)
+                self._flatten_dict_obj(obj.__dict__, dict_data, exclude=exclude)
                 data['__dict__'] = dict_data
             else:
                 data['__dict__'] = self._getref(obj.__dict__)
diff --git a/jsonpickle/util.py b/jsonpickle/util.py
index 6b3f551..199a21c 100644
--- a/jsonpickle/util.py
+++ b/jsonpickle/util.py
@@ -577,9 +577,11 @@ def itemgetter(obj, getter=operator.itemgetter(0)):
     return compat.ustr(getter(obj))
 
 
-def items(obj):
+def items(obj, exclude=()):
     """
     TODO: Replace all calls to this with plain dict.items()
     """
     for k, v in obj.items():
+        if k in exclude:
+            continue
         yield k, v


Test patch:
diff --git a/tests/object_test.py b/tests/object_test.py
index 2434057..04610e8 100644
--- a/tests/object_test.py
+++ b/tests/object_test.py
@@ -289,6 +289,20 @@ class ExceptionWithArguments(Exception):
         self.value = value
 
 
+class ThingWithExclusion:
+    _jsonpickle_exclude = ["foo"]
+
+    def __init__(self, a):
+        self.foo = 1
+        self.bar = a
+
+
+class ThingWithExcludeSubclass:
+    def __init__(self, foo):
+        self.foo = foo
+        self.thing = ThingWithExclusion(3)
+
+
 class AdvancedObjectsTestCase(SkippableTest):
     def setUp(self):
         self.pickler = jsonpickle.pickler.Pickler()
@@ -1141,5 +1155,23 @@ def test_datetime_with_tz_copies_refs():
     assert decoded[1] == d1
 
 
+def test_with_exclude():
+    """Does the _jsonpickle_exclude work"""
+    obj = ThingWithExclusion('baz')
+    encoded = jsonpickle.encode(obj)
+    decoded = jsonpickle.decode(encoded)
+    assert decoded.bar == 'baz'
+    assert not hasattr(decoded, 'foo')
+
+
+def test_contained_exclusion():
+    """_jsonpickle_exclude should work only on the class it is defined in"""
+    obj = ThingWithExcludeSubclass('alpha')
+    encoded = jsonpickle.encode(obj)
+    decoded = jsonpickle.decode(encoded)
+    assert decoded.foo == 'alpha'
+    assert not hasattr(decoded.thing, 'foo')
+
+
 if __name__ == '__main__':
     pytest.main([__file__])