From 37c602b0140666c2129bbb887d05327b5a17c318 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marta=20G=C3=B3mez=20Mac=C3=ADas?= <mgmacias@google.com>
Date: Tue, 19 Dec 2023 17:39:25 +0100
Subject: [PATCH] fix(object): Load a pickled vt object (#175)

* fix(object): Load a pickled vt object

* improve test
---
 tests/test_client.py | 9 +++++++++
 vt/object.py         | 9 +++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

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",
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 @@
 __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 @@ def __init__(
     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 @@ def __getattribute__(self, attr: str) -> typing.Any:
   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: