improvements to to_dict to support nested objects (piccolo-orm#222)

Author: dantownsend

piccolo/table.py

@@ -430,16 +430,34 @@ def to_dict(self, *columns: Column) -> t.Dict[str, t.Any]:
             {'id': 1, 'title': 'Guido'}
 
         """
+        # Make sure we're only looking at columns for the current table. If
+        # someone passes in a column for a sub table (for example
+        # `Band.manager.name`), we need to add `Band.manager` so the nested
+        # value appears in the output.
+        filtered_columns = []
+        for column in columns:
+            if column._meta.table == self.__class__:
+                filtered_columns.append(column)
+            else:
+                for parent_column in column._meta.call_chain:
+                    if parent_column._meta.table == self.__class__:
+                        filtered_columns.append(parent_column)
+                        break
+
         alias_names = {
             column._meta.name: getattr(column, "alias", None)
-            for column in columns
+            for column in filtered_columns
         }
 
         output = {}
-        for column in columns or self._meta.columns:
+        for column in filtered_columns if columns else self._meta.columns:
+            value = getattr(self, column._meta.name)
+            if isinstance(value, Table):
+                value = value.to_dict(*columns)
+
             output[
                 alias_names.get(column._meta.name) or column._meta.name
-            ] = getattr(self, column._meta.name)
+            ] = value
         return output
 
     def __setitem__(self, key: str, value: t.Any):

tests/table/instance/test_to_dict.py

@@ -1,5 +1,5 @@
 from tests.base import DBTestCase
-from tests.example_app.tables import Manager
+from tests.example_app.tables import Band, Manager
 
 
 class TestToDict(DBTestCase):
@@ -11,7 +11,26 @@ def test_to_dict(self):
 
         instance = Manager.objects().first().run_sync()
         dictionary = instance.to_dict()
-        self.assertEqual(dictionary, {"id": 1, "name": "Guido"})
+        self.assertDictEqual(dictionary, {"id": 1, "name": "Guido"})
+
+    def test_nested(self):
+        """
+        Make sure that `to_dict` works correctly, when the object contains
+        nested objects.
+        """
+        self.insert_row()
+
+        instance = Band.objects(Band.manager).first().run_sync()
+        dictionary = instance.to_dict()
+        self.assertDictEqual(
+            dictionary,
+            {
+                "id": 1,
+                "name": "Pythonistas",
+                "manager": {"id": 1, "name": "Guido"},
+                "popularity": 1000,
+            },
+        )
 
     def test_filter_rows(self):
         """
@@ -21,9 +40,26 @@ def test_filter_rows(self):
 
         instance = Manager.objects().first().run_sync()
         dictionary = instance.to_dict(Manager.name)
-        self.assertEqual(dictionary, {"name": "Guido"})
+        self.assertDictEqual(dictionary, {"name": "Guido"})
+
+    def test_nested_filter(self):
+        """
+        Make sure that `to_dict` works correctly with nested objects and
+        filtering.
+        """
+        self.insert_row()
+
+        instance = Band.objects(Band.manager).first().run_sync()
+        dictionary = instance.to_dict(Band.name, Band.manager.id)
+        self.assertDictEqual(
+            dictionary,
+            {
+                "name": "Pythonistas",
+                "manager": {"id": 1},
+            },
+        )
 
-    def test_to_dict_aliases(self):
+    def test_aliases(self):
         """
         Make sure that `to_dict` works correctly with aliases.
         """
@@ -33,4 +69,4 @@ def test_to_dict_aliases(self):
         dictionary = instance.to_dict(
             Manager.id, Manager.name.as_alias("title")
         )
-        self.assertEqual(dictionary, {"id": 1, "title": "Guido"})
+        self.assertDictEqual(dictionary, {"id": 1, "title": "Guido"})