Merge pull request #23 from uncountableinc/yagmin/divide-zero

Custom error for divide by zero

hotxlfp/formulas/error.py

@@ -6,8 +6,11 @@
 class XLError(RuntimeError):
     pass
 
+class ComputationValueError(XLError):
+    pass
+
 ERROR = XLError('#ERROR!')
-DIV_ZERO = XLError('#DIV/0!')
+DIV_ZERO = ComputationValueError("#DIV/0!")
 NAME = XLError('#NAME?')
 NOT_AVAILABLE = XLError('#N/A')
 NULL = XLError('#NULL!')

hotxlfp/formulas/logic.py

@@ -17,10 +17,27 @@ def AND(*args):
 
 @dispatcher.register_for('IF')
 def IF(test, then, otherwise):
+    if isinstance(test, error.XLError):
+        return error.XLError
+    test_condition = torch.tensor(test, dtype=torch.bool)
+    if test_condition.all():
+        if isinstance(then, error.XLError):
+            return then
+        return torch.tensor(then, dtype=torch.double)
+    elif (~test_condition).all():
+        if isinstance(otherwise, error.XLError):
+            return otherwise
+        return torch.tensor(otherwise, dtype=torch.double)
+
+    if isinstance(then, error.XLError):
+        return then
+    if isinstance(otherwise, error.XLError):
+        return otherwise
+
     return torch.where(
-        torch.tensor(test, dtype=torch.bool),
+        test_condition,
         torch.tensor(then, dtype=torch.double),
-        torch.tensor(otherwise, dtype=torch.double)
+        torch.tensor(otherwise, dtype=torch.double),
     )
 
 

setup.py

@@ -4,7 +4,7 @@
 
 setup(
     name="hotxlfp",
-    version="0.0.11-unc19",
+    version="0.0.11-unc20",
     packages=[
         "hotxlfp",
         "hotxlfp._compat",