From 61d9b0cba449bb404b90ff4a5e17aee385224759 Mon Sep 17 00:00:00 2001
From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
Date: Wed, 29 May 2024 17:21:12 -0500
Subject: [PATCH] Minor refactor for computed fields in tables (#325)

---
 src/python-fastui/fastui/components/tables.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/src/python-fastui/fastui/components/tables.py b/src/python-fastui/fastui/components/tables.py
index 22328d51..09e709da 100644
--- a/src/python-fastui/fastui/components/tables.py
+++ b/src/python-fastui/fastui/components/tables.py
@@ -43,22 +43,17 @@ def _fill_columns(self) -> _te.Self:
             except IndexError:
                 raise ValueError('Cannot infer model from empty data, please set `Table(..., model=MyModel)`')
 
+        all_model_fields = {**data_model_type.model_fields, **data_model_type.model_computed_fields}
         if self.columns is None:
-            self.columns = []
-            for name, field in data_model_type.model_fields.items():
-                self.columns.append(display.DisplayLookup(field=name, title=field.title))
-            for name, field in data_model_type.model_computed_fields.items():
-                self.columns.append(display.DisplayLookup(field=name, title=field.title))
-
+            self.columns = [
+                display.DisplayLookup(field=name, title=field.title) for name, field in all_model_fields.items()
+            ]
         else:
             # add pydantic titles to columns that don't have them
             for column in (c for c in self.columns if c.title is None):
-                model_field = data_model_type.model_fields.get(column.field)
-                computed_field = data_model_type.model_computed_fields.get(column.field)
-                if model_field and model_field.title:
-                    column.title = model_field.title
-                elif computed_field and computed_field.title:
-                    column.title = computed_field.title
+                field = all_model_fields.get(column.field)
+                if field and field.title:
+                    column.title = field.title
         return self
 
     @classmethod