Further simplifcation and testing of contentviews.

3 files changed, 37 insertions, 64 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index c968fb1c..9bac5773 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -255,7 +255,7 @@ class ConsoleState(flow.State):
         flow.State.__init__(self)
         self.focus = None
         self.follow_focus = None
-        self.default_body_view = contentview.ViewAuto
+        self.default_body_view = contentview.get("Auto")
         self.view_flow_mode = common.VIEW_FLOW_REQUEST
         self.last_script = ""
         self.last_saveload = ""
diff --git a/libmproxy/console/contentview.py b/libmproxy/console/contentview.py
index ddc82540..763baab1 100644
--- a/libmproxy/console/contentview.py
+++ b/libmproxy/console/contentview.py
@@ -42,6 +42,18 @@ class ViewAuto:
     name = "Auto"
     prompt = ("auto", "a")
     content_types = []
+    def __call__(self, hdrs, content, limit):
+        ctype = hdrs.get("content-type")
+        if ctype:
+            ctype = ctype[0]
+        ct = utils.parse_content_type(ctype) if ctype else None
+        if ct:
+            ct = "%s/%s"%(ct[0], ct[1])
+            if ct in content_types_map:
+                return content_types_map[ct][0](hdrs, content, limit)
+            elif utils.isXML(content):
+                return get("XML")(hdrs, content, limit)
+        return get("Raw")(hdrs, content, limit)
 
 
 class ViewRaw:
@@ -324,26 +336,6 @@ def get(name):
             return i
 
 
-def get_view_func(viewmode, hdrs, content):
-    """
-        Returns a function object.
-    """
-    if viewmode.name == "Auto":
-        ctype = hdrs.get("content-type")
-        if ctype:
-            ctype = ctype[0]
-        ct = utils.parse_content_type(ctype) if ctype else None
-        if ct:
-            ct = "%s/%s"%(ct[0], ct[1])
-            if ct in content_types_map:
-                return content_types_map[ct][0]
-            elif utils.isXML(content):
-                return ViewXML
-        return ViewRaw
-    else:
-        return viewmode
-
-
 def get_content_view(viewmode, hdrItems, content, limit):
     """
         Returns a (msg, body) tuple.
@@ -358,13 +350,12 @@ def get_content_view(viewmode, hdrItems, content, limit):
         if decoded:
             content = decoded
             msg.append("[decoded %s]"%enc[0])
-    func = get_view_func(viewmode, hdrs, content)
     try:
-        ret = func(hdrs, content, limit)
+        ret = viewmode(hdrs, content, limit)
     # Third-party viewers can fail in unexpected ways...
     except Exception, e:
-        #s = traceback.format_exc()
-        #return "", _view_text(s, len(s), len(s))
+        s = traceback.format_exc()
+        return "", _view_text(s, len(s), len(s))
         ret = None
     if not ret:
         ret = get("Raw")(hdrs, content, limit)
@@ -372,5 +363,3 @@ def get_content_view(viewmode, hdrItems, content, limit):
     else:
         msg.append(ret[0])
     return " ".join(msg), ret[1]
-
-
diff --git a/test/test_console_contentview.py b/test/test_console_contentview.py
index 04f72a17..d2ab72dc 100644
--- a/test/test_console_contentview.py
+++ b/test/test_console_contentview.py
@@ -11,61 +11,42 @@ class TestContentView:
         cv.trailer(cv.VIEW_CUTOFF + 10, txt, cv.VIEW_CUTOFF)
         assert txt
 
-    def test_get_view_func(self):
-        f = cv.get_view_func(
-                cv.get("Hex"),
+    def test_view_auto(self):
+        v = cv.ViewAuto()
+        f = v(
                 flow.ODictCaseless(),
-                "foo"
-              )
-        assert f.name == "Hex"
-
-        f = cv.get_view_func(
-                cv.get("Auto"),
-                flow.ODictCaseless(),
-                "foo"
+                "foo",
+                1000
               )
-        assert f.name == "Raw"
+        assert f[0] == "Raw"
 
-        f = cv.get_view_func(
-                cv.get("Auto"),
+        f = v(
                 flow.ODictCaseless(
                     [["content-type", "text/html"]],
                 ),
-                "foo"
+                "<html></html>",
+                1000
               )
-        assert f.name == "HTML"
+        assert f[0] == "HTML"
 
-        f = cv.get_view_func(
-                cv.get("Auto"),
+        f = v(
                 flow.ODictCaseless(
                     [["content-type", "text/flibble"]],
                 ),
-                "foo"
+                "foo",
+                1000
               )
-        assert f.name == "Raw"
+        assert f[0] == "Raw"
 
-        f = cv.get_view_func(
-                cv.get("Auto"),
+        f = v(
                 flow.ODictCaseless(
                     [["content-type", "text/flibble"]],
                 ),
-                "<xml></xml>"
+                "<xml></xml>",
+                1000
               )
-        assert f.name == "XML" 
-
-        try:
-            import pyamf
+        assert f[0].startswith("XML")
 
-            f = cv.get_view_func(
-                    cv.get("Auto"),
-                    flow.ODictCaseless(
-                        [["content-type", "application/x-amf"]],
-                    ),
-                    ""
-                  )
-            assert f.name == "AMF"
-        except ImportError:
-            pass
 
     def test_view_urlencoded(self):
         d = utils.urlencode([("one", "two"), ("three", "four")])
@@ -223,3 +204,6 @@ Larry
         assert "decoded gzip" in r[0]
         assert "Raw" in r[0]
 
+
+def test_get_by_shortcut():
+    assert cv.get_by_shortcut("h")