diff --git a/dbzero/dbzero/decorators.py b/dbzero/dbzero/decorators.py index f33c92a1..3e0279fa 100644 --- a/dbzero/dbzero/decorators.py +++ b/dbzero/dbzero/decorators.py @@ -34,19 +34,28 @@ def wrapper(*args, **kwargs): wrapper.__db0_fulltext = True return wrapper -def table_view(f, operator=None): +def table_view(f=None, operator=None): """The table_view decorator marks a specific function or method as generating a table view. The following properties apply: - First result represents the table header - All other results (rows) are key-decorated (i.e. return tuples with unique identifiers) - see Cell Editor - Optional operator may be defined to allow cell editions (see Cell Operator) """ - @functools.wraps(f) - def wrapper(*args, **kwargs): - retval = f(*args, **kwargs) - return retval - # Immutable query if True and immutable parameter if false - wrapper.__db0_table_view = True - return wrapper + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + retval = func(*args, **kwargs) + return retval + # Immutable query if True and immutable parameter if false + wrapper.__db0_table_view = True + wrapper.__db0_table_view_operator = operator + return wrapper + + if f is None: + # Called with arguments: @table_view(operator=...) + return decorator + else: + # Called without arguments: @table_view + return decorator(f) def is_immutable(f): return hasattr(f, '__db0_immutable_query') diff --git a/python_tests/test_decorators.py b/python_tests/test_decorators.py index e914c617..5b28ae73 100644 --- a/python_tests/test_decorators.py +++ b/python_tests/test_decorators.py @@ -158,3 +158,26 @@ def action(self, has, param): assert db0.has_complete_action(InnerClass.query) with pytest.raises(RuntimeError): db0.get_complete_action(InnerClass.query) + + +def test_table_view_with_operator(): + def my_operator(row_id, col_id, new_value): + return f"Updated {row_id}, {col_id} to {new_value}" + + @db0.table_view(operator=my_operator) + def my_table_view_with_op(): + return [("header1", "header2"), ("row1", "row2")] + + assert db0.is_table_view(my_table_view_with_op) + assert hasattr(my_table_view_with_op, '__db0_table_view') + assert my_table_view_with_op.__db0_table_view == True + assert hasattr(my_table_view_with_op, '__db0_table_view_operator') + assert my_table_view_with_op.__db0_table_view_operator == my_operator + + # Test the decorator preserves function behavior + result = my_table_view_with_op() + assert result == [("header1", "header2"), ("row1", "row2")] + + # Test the operator works correctly + assert my_table_view_with_op.__db0_table_view_operator("row1", "col1", "new_val") == "Updated row1, col1 to new_val" +