From cbd645fb7b1bdd926cc5452d997a217b248d6f6e Mon Sep 17 00:00:00 2001
From: Christophe de Vienne <cdevienne@gmail.com>
Date: Fri, 7 Oct 2011 13:14:42 +0200
Subject: [PATCH] Add array support to restjson

---
 wsme/protocols/restjson.py  | 10 ++++++
 wsme/tests/protocol.py      | 67 +++++++++++++++++++++++++++++++++++++
 wsme/tests/test_restjson.py |  7 ++++
 wsme/types.py               |  8 +++++
 4 files changed, 92 insertions(+)

diff --git a/wsme/protocols/restjson.py b/wsme/protocols/restjson.py
index 3a40e3d..7ac5db6 100644
--- a/wsme/protocols/restjson.py
+++ b/wsme/protocols/restjson.py
@@ -24,6 +24,11 @@ def tojson(datatype, value):
     return value
 
 
+@tojson.when_type(list)
+def array_tojson(datatype, value):
+    return [tojson(datatype[0], item) for item in value]
+
+
 @tojson.when_object(decimal.Decimal)
 def decimal_tojson(datatype, value):
     return str(value)
@@ -62,6 +67,11 @@ def fromjson(datatype, value):
     return value
 
 
+@fromjson.when_type(list)
+def array_fromjson(datatype, value):
+    return [fromjson(datatype[0], item) for item in value]
+
+
 @fromjson.when_object(decimal.Decimal)
 def decimal_fromjson(datatype, value):
     return decimal.Decimal(value)
diff --git a/wsme/tests/protocol.py b/wsme/tests/protocol.py
index 5157341..e3b0143 100644
--- a/wsme/tests/protocol.py
+++ b/wsme/tests/protocol.py
@@ -84,6 +84,14 @@ class ReturnTypes(object):
         n = NestedOuter()
         return n
 
+    @expose([str])
+    def getstrarray(self):
+        return ["A", "B", "C"]
+
+    @expose([NestedOuter])
+    def getnestedarray(self):
+        return [NestedOuter(), NestedOuter()]
+
 
 class ArgTypes(object):
     @expose(str)
@@ -156,6 +164,22 @@ class ArgTypes(object):
         assert type(value) == str
         return value
 
+    @expose([str])
+    @validate([str])
+    def setstrarray(self, value):
+        print repr(value)
+        assert type(value) == list
+        assert type(value[0]) == str
+        return value
+
+    @expose([datetime.datetime])
+    @validate([datetime.datetime])
+    def setdatetimearray(self, value):
+        print repr(value)
+        assert type(value) == list
+        assert type(value[0]) == datetime.datetime
+        return value
+
     @expose(NestedOuter)
     @validate(NestedOuter)
     def setnested(self, value):
@@ -163,6 +187,14 @@ class ArgTypes(object):
         assert type(value) == NestedOuter
         return value
 
+    @expose([NestedOuter])
+    @validate([NestedOuter])
+    def setnestedarray(self, value):
+        print repr(value)
+        assert type(value) == list
+        assert type(value[0]) == NestedOuter
+        return value
+
 
 class WithErrors(object):
     @expose()
@@ -252,6 +284,14 @@ class ProtocolTestCase(unittest.TestCase):
         r = self.call('returntypes/getnested', _rt=NestedOuter)
         assert r == {'inner': {'aint': 0}}, r
 
+    def test_return_strarray(self):
+        r = self.call('returntypes/getstrarray', _rt=[str])
+        assert r == ['A', 'B', 'C'], r
+
+    def test_return_strnested(self):
+        r = self.call('returntypes/getnestedarray', _rt=[NestedOuter])
+        assert r == [{'inner': {'aint': 0}}, {'inner': {'aint': 0}}], r
+
     def test_setstr(self):
         assert self.call('argtypes/setstr', value='astring') == 'astring'
 
@@ -301,3 +341,30 @@ class ProtocolTestCase(unittest.TestCase):
                          value=(value, NestedOuter),
                          _rt=NestedOuter)
         assert r == value
+
+    def test_setstrarray(self):
+        value = ["1", "2", "three"]
+        r = self.call('argtypes/setstrarray',
+                         value=(value, [str]),
+                         _rt=[str])
+        assert r == value
+
+    def test_setdatetimearray(self):
+        value = [
+            datetime.datetime(2008, 3, 6, 12, 12, 15),
+            datetime.datetime(2008, 4, 6, 2, 12, 15),
+        ]
+        r = self.call('argtypes/setdatetimearray',
+                         value=(value, [datetime.datetime]),
+                         _rt=[datetime.datetime])
+        assert r == value
+
+    def test_setnestedarray(self):
+        value = [
+            {'inner': {'aint': 54}},
+            {'inner': {'aint': 55}},
+        ]
+        r = self.call('argtypes/setnestedarray',
+                         value=(value, [NestedOuter]),
+                         _rt=[NestedOuter])
+        assert r == value
diff --git a/wsme/tests/test_restjson.py b/wsme/tests/test_restjson.py
index a965155..3aa59d5 100644
--- a/wsme/tests/test_restjson.py
+++ b/wsme/tests/test_restjson.py
@@ -11,7 +11,10 @@ except:
 import wsme.protocols.restjson
 from wsme.utils import *
 
+
 def prepare_value(value, datatype):
+    if isinstance(datatype, list):
+        return [prepare_value(item, datatype[0]) for item in value]
     if datatype in (datetime.date, datetime.time, datetime.datetime):
         return value.isoformat()
     if datatype == decimal.Decimal:
@@ -20,7 +23,10 @@ def prepare_value(value, datatype):
         return base64.encodestring(value)
     return value
 
+
 def prepare_result(value, datatype):
+    if isinstance(datatype, list):
+        return [prepare_result(item, datatype[0]) for item in value]
     if datatype == datetime.date:
         return parse_isodate(value)
     if datatype == datetime.time:
@@ -39,6 +45,7 @@ def prepare_result(value, datatype):
         return datatype(value)
     return value
 
+
 class TestRestJson(wsme.tests.protocol.ProtocolTestCase):
     protocol = 'REST+Json'
 
diff --git a/wsme/types.py b/wsme/types.py
index fae9b49..57c6e41 100644
--- a/wsme/types.py
+++ b/wsme/types.py
@@ -11,6 +11,7 @@ extra_types = [binary, decimal.Decimal]
 native_types = pod_types + dt_types + extra_types
 
 complex_types = []
+array_types = []
 
 
 def iscomplex(datatype):
@@ -104,6 +105,13 @@ def register_type(class_):
             hasattr(class_, '_wsme_attributes'):
         return
 
+    if isinstance(class_, list):
+        if len(class_) != 1:
+            raise ValueError("Cannot register type %s" % repr(class_))
+        register_type(class_[0])
+        array_types.append(class_[0])
+        return
+
     class_._wsme_attributes = inspect_class(class_)
 
     complex_types.append(weakref.ref(class_))