From 7c7e7eb81a44be21cb28a938f5b505ebd2804b7e Mon Sep 17 00:00:00 2001
From: Xiang Gao <qasdfgtyuiop@gmail.com>
Date: Wed, 29 May 2019 11:10:17 -0400
Subject: [PATCH] Changes according to review

---
 README.md               |  8 ++++++--
 pysnooper/tracer.py     |  4 ++++
 pysnooper/utils.py      | 17 +++++++++++------
 setup.py                |  1 +
 tests/test_pysnooper.py | 15 ++++++++++++++-
 5 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 64f2c1a..26815c5 100644
--- a/README.md
+++ b/README.md
@@ -160,15 +160,19 @@ def large(l):
 def print_list_size(l):
     return 'list(size={})'.format(len(l))
 
-@pysnooper.snoop(custom_repr=((large, print_list_size),))
+def print_ndarray(a):
+    return 'ndarray(shape={}, dtype={})'.format(a.shape, a.dtype)
+
+@pysnooper.snoop(custom_repr=((large, print_list_size), (numpy.ndarray, print_ndarray)))
 def sum_to_x(x):
     l = list(range(x))
+    a = numpy.zeros((10,10))
     return sum(l)
 
 sum_to_x(10000)
 ```
 
-You will get `l = list(size=10000)` for the list
+You will get `l = list(size=10000)` for the list, and `a = ndarray(shape=(10, 10), dtype=float64)` for the ndarray.
 
 # Installation #
 
diff --git a/pysnooper/tracer.py b/pysnooper/tracer.py
index a478d00..98a2d3c 100644
--- a/pysnooper/tracer.py
+++ b/pysnooper/tracer.py
@@ -179,6 +179,10 @@ class Tracer:
 
         @pysnooper.snoop(thread_info=True)
 
+    Customize how values are represented as strings::
+
+        @pysnooper.snoop(custom_repr=((type1, custom_repr_func1), (condition2, custom_repr_func2), ...))
+
     '''
     def __init__(
             self,
diff --git a/pysnooper/utils.py b/pysnooper/utils.py
index a1a6f80..cab633f 100644
--- a/pysnooper/utils.py
+++ b/pysnooper/utils.py
@@ -49,14 +49,19 @@ def shitcode(s):
     )
 
 
+def get_repr_function(item, custom_repr):
+    for condition, action in custom_repr:
+        if isinstance(condition, type):
+            condition = lambda x, y=condition: isinstance(x, y)
+        if condition(item):
+            return action
+    return repr
+
+
 def get_shortish_repr(item, custom_repr=()):
     try:
-        for condition, action in custom_repr:
-            if condition(item):
-                r = action(item)
-                break
-        else:
-            r = repr(item)
+        repr_function = get_repr_function(item, custom_repr)
+        r = repr_function(item)
     except Exception:
         r = 'REPR FAILED'
     r = r.replace('\r', '').replace('\n', '')
diff --git a/setup.py b/setup.py
index e01e6e6..1186979 100644
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@ setuptools.setup(
         'tests': {
             'pytest',
             'python-toolbox',
+            'numpy',
         },
     },
     classifiers=[
diff --git a/tests/test_pysnooper.py b/tests/test_pysnooper.py
index e4e74e9..8ec1037 100644
--- a/tests/test_pysnooper.py
+++ b/tests/test_pysnooper.py
@@ -6,6 +6,7 @@ import textwrap
 import threading
 import types
 import sys
+import numpy
 
 from pysnooper.utils import truncate
 from python_toolbox import sys_tools, temp_file_tools
@@ -1140,9 +1141,19 @@ def test_custom_repr():
     def print_list_size(l):
         return 'list(size={})'.format(len(l))
 
-    @pysnooper.snoop(string_io, custom_repr=((large, print_list_size),))
+    def print_ndarray(a):
+        return 'ndarray(shape={}, dtype={})'.format(a.shape, a.dtype)
+
+    def evil_condition(x):
+        return large(x) or isinstance(x, numpy.ndarray)
+
+    @pysnooper.snoop(string_io, custom_repr=(
+        (large, print_list_size),
+        (numpy.ndarray, print_ndarray),
+        (evil_condition, lambda x: 'I am evil')))
     def sum_to_x(x):
         l = list(range(x))
+        a = numpy.zeros((10,10))
         return sum(l)
 
     result = sum_to_x(10000)
@@ -1156,6 +1167,8 @@ def test_custom_repr():
             LineEntry(),
             VariableEntry('l', 'list(size=10000)'),
             LineEntry(),
+            VariableEntry('a', 'ndarray(shape=(10, 10), dtype=float64)'),
+            LineEntry(),
             ReturnEntry(),
             ReturnValueEntry('49995000'),
         )