Repository: VirusTotal/vt-py

Issue:
Unable to save WhistleBlowerDict obj. (or an object with nested WhistleBlowerDict) to a pickle.
When I'm saving a custom class instance to a pickle file, there's no error, file is being generated in the filesystem. 
I do this with:
`    pickle.dump(proc_db, open('process_db_vted.pickle', 'wb'))`
I have also tried changing default arguments (protocol=pickle.HIGHEST_PROTOCOL and fix_imports=True).

When trying to load the pickle file with:
`    proc_db = pickle.load(open('process_db_vted.pickle', 'rb'))`
I get:
```
    proc_db = pickle.load(open('process_db_vted.pickle', 'rb'))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\venv\Lib\site-packages\vt\object.py", line 39, in __setitem__
    self._on_change_callback()
    ^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'WhistleBlowerDict' object has no attribute '_on_change_callback'
```

please advise
Regards,
Q

Discussion:

Patch:
diff --git a/vt/object.py b/vt/object.py
index 5a03aa8..fd3a668 100644
--- a/vt/object.py
+++ b/vt/object.py
@@ -13,6 +13,7 @@
 
 """Defines a VT object and other helper classes."""
 
+import collections
 import datetime
 import functools
 import re
@@ -21,7 +22,7 @@ import typing
 __all__ = ["Object"]
 
 
-class WhistleBlowerDict(dict):
+class WhistleBlowerDict(collections.UserDict):
   """Helper class for detecting changes in a dictionary.
 
   This class wraps a standard Python dictionary and calls the provided callback
@@ -141,7 +142,7 @@ class Object:
     self._modified_data = {}
     self._error = None
 
-  def __on_attr_change(self, attr: str) -> None:
+  def _on_attr_change(self, attr: str) -> None:
     if hasattr(self, "_modified_attrs"):
       self._modified_attrs.append(attr)
 
@@ -156,12 +157,12 @@ class Object:
   def __setattr__(self, attr: str, value: typing.Any) -> None:
     if isinstance(value, dict):
       value = WhistleBlowerDict(
-          value, functools.partial(self.__on_attr_change, attr)
+          value, functools.partial(self._on_attr_change, attr)
       )
     elif isinstance(value, datetime.datetime):
       value = int(datetime.datetime.timestamp(value))
     if attr not in self.__dict__ or value != self.__dict__[attr]:
-      self.__on_attr_change(attr)
+      self._on_attr_change(attr)
     super().__setattr__(attr, value)
 
   def __repr__(self) -> str:


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 4463afb..4f513ac 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -16,6 +16,7 @@
 import datetime
 import io
 import json
+import pickle
 
 import pytest
 from vt import APIError
@@ -70,6 +71,14 @@ def test_object_date_attrs():
   assert obj.foo_date == datetime.datetime(1970, 1, 1, 0, 0, 0)
 
 
+def test_object_pickle():
+  obj = Object("dummy")
+  obj.whatever = {"1": "2"}
+  new = pickle.loads(pickle.dumps(obj))
+  assert new.whatever == obj.whatever
+  assert new.to_dict() == obj.to_dict()
+
+
 def test_object_to_dict():
   obj = Object.from_dict({
       "type": "dummy_type",