From 5ed81cb8483f41f132628b5604851ca18d4c54c3 Mon Sep 17 00:00:00 2001
From: Jonathan Reichelt Gjertsen
 <jonathan.gjertsen@disruptive-technologies.com>
Date: Sun, 8 Sep 2019 12:17:47 +0200
Subject: [PATCH] Add support for wrapping classes
 (https://github.com/cool-RR/PySnooper/issues/150)

---
 pysnooper/tracer.py     | 16 +++++++++++++++-
 tests/test_pysnooper.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/pysnooper/tracer.py b/pysnooper/tracer.py
index 0844c59..6584d94 100644
--- a/pysnooper/tracer.py
+++ b/pysnooper/tracer.py
@@ -223,9 +223,23 @@ class Tracer:
         self.custom_repr = custom_repr
         self.last_source_path = None
 
-    def __call__(self, function):
+    def __call__(self, function_or_class):
         if DISABLED:
             return function
+
+        if inspect.isclass(function_or_class):
+            return self._wrap_class(function_or_class)
+        else:
+            return self._wrap_function(function_or_class)
+
+    def _wrap_class(self, cls):
+        for attr_name in dir(cls):
+            attr = getattr(cls, attr_name)
+            if inspect.isfunction(attr):
+                setattr(cls, attr_name, self._wrap_function(attr))
+        return cls
+
+    def _wrap_function(self, function):
         self.target_codes.add(function.__code__)
 
         @functools.wraps(function)
diff --git a/tests/test_pysnooper.py b/tests/test_pysnooper.py
index b6f03e1..65f1d7b 100644
--- a/tests/test_pysnooper.py
+++ b/tests/test_pysnooper.py
@@ -48,6 +48,43 @@ def test_string_io():
     )
 
 
+def test_class():
+    string_io = io.StringIO()
+
+    @pysnooper.snoop(string_io)
+    class MyClass(object):
+        def __init__(self):
+            self.x = 7
+
+        def my_method(self, foo):
+            y = 8
+            return y + self.x
+
+    instance = MyClass()
+    result = instance.my_method('baba')
+    assert result == 15
+    output = string_io.getvalue()
+    assert_output(
+        output,
+        (
+            SourcePathEntry(),
+            VariableEntry('self', value_regex="u?.*<locals>.MyClass object at"),
+            CallEntry('def __init__(self):'),
+            LineEntry('self.x = 7'),
+            ReturnEntry('self.x = 7'),
+            ReturnValueEntry('None'),
+            VariableEntry('self', value_regex="u?.*<locals>.MyClass object at"),
+            VariableEntry('foo', value_regex="u?'baba'"),
+            CallEntry('def my_method(self, foo):'),
+            LineEntry('y = 8'),
+            VariableEntry('y', '8'),
+            LineEntry('return y + self.x'),
+            ReturnEntry('return y + self.x'),
+            ReturnValueEntry('15'),
+        )
+    )
+
+
 def test_thread_info():
 
     @pysnooper.snoop(thread_info=True)