fixint output

Author: dantownsend

aragorm/query.py

@@ -87,14 +87,16 @@ async def run(self, as_dict=True, credentials=None):
         output = getattr(self, '_output')
         if output:
             if output.as_list:
-                if output[0].keys().length != 1:
+                if len(raw[0].keys()) != 1:
                     raise ValueError(
                         'Each row returned more than on value'
                     )
                 else:
-                    output = [i for i in itertools.chain(*output.values())]
+                    raw = list(
+                        itertools.chain(*[j.values() for j in raw])
+                    )
             if output.as_json:
-                output = json.dumps(raw)
+                raw = json.dumps(raw)
 
         return raw
 
@@ -234,6 +236,8 @@ def output(
         if type(as_json) is bool:
             self._output.as_json = as_json
 
+        return self
+
 
 ###############################################################################
 

tests/table/test_output.py

@@ -0,0 +1,31 @@
+from ..base import DBTestCase
+from ..example_project.tables import Pokemon
+
+
+class TestOutput(DBTestCase):
+
+    def test_output_as_list(self):
+        self.insert_row()
+
+        response = Pokemon.select(
+            'name'
+        ).output(
+            as_list=True
+        ).run_sync()
+
+        self.assertTrue(
+            response == ['pikachu']
+        )
+
+    def test_output_as_json(self):
+        self.insert_row()
+
+        response = Pokemon.select(
+            'name'
+        ).output(
+            as_json=True
+        ).run_sync()
+
+        self.assertTrue(
+            response == '[{"name":"pikachu"}]'
+        )