Merge pull request #15 from plumdog/inherit_when_creating_table

Inherit columns when inheriting a parent table class

flask_table/table.py

@@ -21,8 +21,21 @@ def __new__(meta, name, bases, attrs):
 
         """
         cls = type.__new__(meta, name, bases, attrs)
-        cols = [(k, v) for k, v in attrs.items() if isinstance(v, Col)]
-        cls._cols = OrderedDict(sorted(cols, key=lambda x: x[1]._counter_val))
+        cls._cols = OrderedDict()
+        # If there are any base classes with a `_cols` attribute, add
+        # them to the columns for this table.
+        for parent in bases:
+            try:
+                parent_cols = parent._cols
+            except AttributeError:
+                continue
+            else:
+                cls._cols.update(parent_cols)
+        # Then add the columns from this class.
+        this_cls_cols = sorted(
+            ((k, v) for k, v in attrs.items() if isinstance(v, Col)),
+            key=lambda x: x[1]._counter_val)
+        cls._cols.update(OrderedDict(this_cls_cols))
         return cls
 
 
@@ -131,9 +144,16 @@ def add_column(cls, name, col):
         return cls
 
 
-def create_table(name=str('_Table')):
+def create_table(name=str('_Table'), base=Table):
     """Creates and returns a new table class. You can specify a name for
-    you class if you wish.
+    you class if you wish. You can also set the base class (or
+    classes) that should be used when creating the class.
 
     """
-    return type(name, (Table,), {})
+    try:
+        base = tuple(base)
+    except TypeError:
+        # Then assume that what we have is a single class, so make it
+        # into a 1-tuple.
+        base = (base,)
+    return TableMeta(name, base, {})

tests/__init__.py

@@ -149,6 +149,21 @@ def test_ten(self):
             'dynamic_cols_num_cols_test', 'test_ten', items)
 
 
+class DynamicColsInheritTest(TableTest):
+    def setUp(self):
+        # Start with MyTable
+        class MyTable(Table):
+            name = Col('Name')
+        # Then dynamically extend it.
+        self.table_cls = create_table(base=MyTable)
+        self.table_cls.add_column('number', Col('Number'))
+
+    def test_one(self):
+        items = [{'name': 'TestName', 'number': 10}]
+        self.assert_html_equivalent_from_file(
+            'dynamic_cols_inherit_test', 'test_one', items)
+
+
 class OverrideTrTest(TableTest):
     def setUp(self):
         class MyTable(Table):

tests/html/dynamic_cols_inherit_test/test_one.html

@@ -0,0 +1,14 @@
+<table>
+    <thead>
+	<tr>
+	    <th>Name</th>
+	    <th>Number</th>
+	</tr>
+    </thead>
+    <tbody>
+	<tr>
+	    <td>TestName</td>
+	    <td>10</td>
+	</tr>
+    </tbody>
+</table>