Make watchpoints work with dataframe (#31)

* Make watchpoints work with dataframe

Author: gaogaotiantian

src/watchpoints/watch_element.py

@@ -6,6 +6,11 @@
 from .watch_print import WatchPrint
 import copy
 
+try:
+    import pandas as pd
+except ImportError:  # pragma: no cover
+    pd = None
+
 
 class WatchElement:
     def __init__(self, frame, node, **kwargs):
@@ -94,6 +99,8 @@ def changed(self, frame):
     def obj_changed(self, other):
         if not isinstance(self.obj, type(other)):
             return True
+        elif pd is not None and isinstance(self.obj, pd.DataFrame):
+            return not self.obj.equals(other)
         elif self.cmp:
             return self.cmp(self.obj, other)
         elif self.obj.__class__.__module__ == "builtins":
@@ -110,7 +117,9 @@ def obj_changed(self, other):
                 return not guess
 
     def update(self):
-        if self.copy:
+        if pd is not None and isinstance(self.obj, pd.DataFrame):
+            self.prev_obj = self.obj.copy(True)
+        elif self.copy:
             self.prev_obj = self.copy(self.obj)
         elif self.deepcopy:
             self.prev_obj = copy.deepcopy(self.obj)

tests/test_pandas.py

@@ -43,7 +43,7 @@ def __comparison_series__(obj1, obj2):
 
         unwatch()
 
-    def test_dataframe(self):
+    def test_dataframe_cmp(self):
         def __comparison_dataframe__(obj1, obj2):
             return not obj1.equals(obj2)
 
@@ -66,3 +66,24 @@ def __comparison_dataframe__(obj1, obj2):
         self.assertEqual(cb.counter, 1)
 
         unwatch()
+
+    def test_dataframe(self):
+        cb = CB()
+
+        df = pd.DataFrame(
+            data=[[1, 2], [3, 4], [5, 6]], index=list("abc"), columns=list("AB")
+        )
+
+        watch(df, callback=cb)
+
+        # Other stuff happens
+        a = 2
+        _ = a + 5
+
+        self.assertEqual(cb.counter, 0)
+
+        df.loc["a", "B"] = 10
+
+        self.assertEqual(cb.counter, 1)
+
+        unwatch()

tests/test_watch_element.py

@@ -5,6 +5,7 @@
 import unittest
 import inspect
 import os
+import pandas as pd
 from watchpoints.watch_element import WatchElement
 from watchpoints.util import getargnodes
 
@@ -169,6 +170,15 @@ def __eq__(self, other):
         obj_eq.a["a"] = 3
         self.assertTrue(wobj_eq.changed(frame)[0])
 
+    def test_dataframe(self):
+        frame = inspect.currentframe()
+        df = pd.DataFrame([1, 2, 3])
+        lst = self.helper(df)
+        wedf = lst[0]
+        self.assertFalse(wedf.changed(frame)[0])
+        df.iloc[0] = 0
+        self.assertTrue(wedf.changed(frame)[0])
+
     def test_global_module(self):
         os.environ['a'] = 'test'
         self.helper(os.environ['a'])