add nested output option (piccolo-orm#204)

* add `nested` output option

* fix typo in comment

* fix a typo in docstring

* updated docs

Author: dantownsend

docs/src/piccolo/query_clauses/output.rst

@@ -39,6 +39,16 @@ If you're just querying a single column from a database table, you can use
     >>> Band.select(Band.id).output(as_list=True).run_sync()
     [1, 2]
 
+nested
+~~~~~~
+
+Output any data from related tables in nested dictionaries.
+
+.. code-block:: python
+
+    >>> Band.select(Band.name, Band.manager.name).first().output(nested=True).run_sync()
+    {'name': 'Pythonistas', 'manager': {'name': 'Guido'}}
+
 -------------------------------------------------------------------------------
 
 Select and Objects queries

docs/src/piccolo/query_types/select.rst

@@ -20,7 +20,7 @@ To get certain columns:
     >>> Band.select(Band.name).run_sync()
     [{'name': 'Rustaceans'}, {'name': 'Pythonistas'}]
 
-Or making an alias to make it shorter:
+Or use an alias to make it shorter:
 
 .. code-block:: python
 
@@ -46,24 +46,41 @@ This is equivalent to ``SELECT name AS title FROM band`` in SQL.
 Joins
 -----
 
-One of the most powerful things about select is it's support for joins.
+One of the most powerful things about ``select`` is it's support for joins.
 
 .. code-block:: python
 
-    >>> b = Band
-    >>> b.select(b.name, b.manager.name).run_sync()
+    >>> Band.select(Band.name, Band.manager.name).run_sync()
     [{'name': 'Pythonistas', 'manager.name': 'Guido'}, {'name': 'Rustaceans', 'manager.name': 'Graydon'}]
 
 
 The joins can go several layers deep.
 
 .. code-block:: python
 
-    c = Concert
-    c.select(
-        c.id,
-        c.band_1.manager.name
-    ).run_sync()
+    >>> Concert.select(Concert.id, Concert.band_1.manager.name).run_sync()
+    [{'id': 1, 'band_1.manager.name': 'Guido'}]
+
+If you want all of the columns from a related table, there's a useful shortcut
+which saves you from typing them all out:
+
+.. code-block:: python
+
+    >>> Band.select(Band.name, *Band.manager.all_columns()).run_sync()
+    [
+        {'name': 'Pythonistas', 'manager.id': 1, 'manager.name': 'Guido'},
+        {'name': 'Rustaceans', 'manager.id': 2, 'manager.name': 'Graydon'}
+    ]
+
+You can also get the response as nested dictionaries, which can be very useful:
+
+.. code-block:: python
+
+    >>> Band.select(Band.name, *Band.manager.all_columns()).output(nested=True).run_sync()
+    [
+        {'name': 'Pythonistas', 'manager': {'id': 1, 'name': 'Guido'}},
+        {'name': 'Rustaceans', 'manager': {'id': 2, 'manager.name': 'Graydon'}}
+    ]
 
 String syntax
 -------------
@@ -183,29 +200,29 @@ By default all columns are returned from the queried table.
 
 .. code-block:: python
 
-    b = Band
     # Equivalent to SELECT * from band
-    b.select().run_sync()
+    Band.select().run_sync()
 
 To restrict the returned columns, either pass in the columns into the
 ``select`` method, or use the ``columns`` method.
 
 .. code-block:: python
 
-    b = Band
     # Equivalent to SELECT name from band
-    b.select().columns(b.name).run_sync()
+    Band.select(Band.name).run_sync()
+
+    # Or alternatively:
+    Band.select().columns(Band.name).run_sync()
 
 The ``columns`` method is additive, meaning you can chain it to add additional
 columns.
 
 .. code-block:: python
 
-    b = Band
-    b.select().columns(b.name).columns(b.manager).run_sync()
+    Band.select().columns(Band.name).columns(Band.manager).run_sync()
 
     # Or just define it one go:
-    b.select().columns(b.name, b.manager).run_sync()
+    Band.select().columns(Band.name, Band.manager).run_sync()
 
 
 first

piccolo/query/methods/select.py

@@ -19,6 +19,7 @@
     WhereDelegate,
 )
 from piccolo.querystring import QueryString
+from piccolo.utils.dictionary import make_nested
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.custom_types import Combinable
@@ -203,13 +204,23 @@ def offset(self, number: int) -> Select:
         return self
 
     async def response_handler(self, response):
+        # If no columns were specified, it's a select *, so we know that
+        # no columns were selected from related tables.
+        was_select_star = len(self.columns_delegate.selected_columns) == 0
+
         if self.limit_delegate._first:
             if len(response) == 0:
                 return None
+
+            if self.output_delegate._output.nested and not was_select_star:
+                return make_nested(response[0])
             else:
                 return response[0]
         else:
-            return response
+            if self.output_delegate._output.nested and not was_select_star:
+                return [make_nested(i) for i in response]
+            else:
+                return response
 
     def order_by(self, *columns: Column, ascending=True) -> Select:
         _columns: t.List[Column] = [
@@ -225,9 +236,13 @@ def output(
         as_list: bool = False,
         as_json: bool = False,
         load_json: bool = False,
+        nested: bool = False,
     ) -> Select:
         self.output_delegate.output(
-            as_list=as_list, as_json=as_json, load_json=load_json
+            as_list=as_list,
+            as_json=as_json,
+            load_json=load_json,
+            nested=nested,
         )
         return self
 

piccolo/query/mixins.py

@@ -78,13 +78,15 @@ class Output:
     as_list: bool = False
     as_objects: bool = False
     load_json: bool = False
+    nested: bool = False
 
     def copy(self) -> Output:
         return self.__class__(
             as_json=self.as_json,
             as_list=self.as_list,
             as_objects=self.as_objects,
             load_json=self.load_json,
+            nested=self.nested,
         )
 
 
@@ -195,6 +197,7 @@ def output(
         as_list: t.Optional[bool] = None,
         as_json: t.Optional[bool] = None,
         load_json: t.Optional[bool] = None,
+        nested: t.Optional[bool] = None,
     ):
         """
         :param as_list:
@@ -216,6 +219,9 @@ def output(
         if load_json is not None:
             self._output.load_json = bool(load_json)
 
+        if nested is not None:
+            self._output.nested = bool(nested)
+
     def copy(self) -> OutputDelegate:
         _output = self._output.copy() if self._output is not None else None
         return self.__class__(_output=_output)

piccolo/utils/dictionary.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import typing as t
+
+
+def make_nested(dictionary: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
+    """
+    Rows are returned from the database as a flat dictionary, with keys such
+    as ``'manager.name'`` if the column belongs to a related table.
+
+    This function puts any values from a related table into a sub dictionary.
+
+    .. code-block::
+
+        response = Band.select(Band.name, Band.manager.name).run_sync()
+        >>> print(response)
+        [{'name': 'Pythonistas', 'band.name': 'Guido'}]
+
+        >>> make_nested(response[0])
+        {'name': 'Pythonistas', 'band': {'name': 'Guido'}}
+
+    """
+    output: t.Dict[str, t.Any] = {}
+
+    for key, value in dictionary.items():
+        path = key.split(".")
+        if len(path) == 1:
+            output[path[0]] = value
+        else:
+            dictionary = output.setdefault(path[0], {})
+            for path_element in path[1:-1]:
+                dictionary = dictionary.setdefault(path_element, {})
+            dictionary[path[-1]] = value
+
+    return output

tests/table/test_output.py

@@ -59,3 +59,30 @@ def test_objects(self):
 
         self.assertEqual(results[0].facilities, json)
         self.assertEqual(results[0].facilities_b, json)
+
+
+class TestOutputNested(DBTestCase):
+    def test_output_nested(self):
+        self.insert_row()
+
+        response = (
+            Band.select(Band.name, Band.manager.name)
+            .output(nested=True)
+            .run_sync()
+        )
+        self.assertEqual(
+            response, [{"name": "Pythonistas", "manager": {"name": "Guido"}}]
+        )
+
+    def test_output_nested_with_first(self):
+        self.insert_row()
+
+        response = (
+            Band.select(Band.name, Band.manager.name)
+            .first()
+            .output(nested=True)
+            .run_sync()
+        )
+        self.assertEqual(
+            response, {"name": "Pythonistas", "manager": {"name": "Guido"}}
+        )

tests/utils/test_dictionary.py

@@ -0,0 +1,28 @@
+from unittest import TestCase
+
+from piccolo.utils.dictionary import make_nested
+
+
+class TestMakeNested(TestCase):
+    def test_nesting(self):
+        response = make_nested(
+            {
+                "id": 1,
+                "name": "Pythonistas",
+                "manager.id": 1,
+                "manager.name": "Guido",
+                "manager.car.colour": "green",
+            }
+        )
+        self.assertEqual(
+            response,
+            {
+                "id": 1,
+                "name": "Pythonistas",
+                "manager": {
+                    "id": 1,
+                    "name": "Guido",
+                    "car": {"colour": "green"},
+                },
+            },
+        )