Refactored decorators #4

mateuszrezler · Dec 10, 2020 · 3f61ab4 · 3f61ab4
1 parent c389313
commit 3f61ab4
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/checkmyxl/predefined.py b/checkmyxl/predefined.py
@@ -2,19 +2,19 @@
 RED = (255, 0, 0)
 
 
-def iterate(func):
-    def wrapper(sheet_range, *args):
+def iterate(function):
+    def wrapper(sheet_range, *args, **kwargs):
         for cell in sheet_range:
-            func(cell, *args)
+            function(cell, *args, **kwargs)
     return wrapper
 
 
-def highlight(func):
-    def wrapper(cell, *args):
-        if func(cell, *args):
-            cell.color = GREEN
-        else:
-            cell.color = RED
+def highlight(function):
+    def wrapper(cell, *args, correct_color=GREEN, incorrect_color=RED,
+                **kwargs):
+        result = function(cell, *args, **kwargs)
+        cell.color = (incorrect_color, correct_color)[int(result)]
+        return result
     return wrapper