Adds ability to capture return value from the cases.

mikeckennedy · mikeckennedy · commit 7da8e53ad71f · 2017-09-14T09:53:25.000-07:00
diff --git a/switchlang.py b/switchlang.py
@@ -1,11 +1,15 @@
 # Here is a first pass implementation at adding switch
+import uuid
 from typing import Callable, Any
 
 
 class switch:
+    __no_result = uuid.uuid4()
+
     def __init__(self, value):
         self.value = value
         self.cases = {}
+        self.__result = switch.__no_result
 
     def default(self, func: Callable[[], Any]):
         self.case('__default__', func)
@@ -42,6 +46,15 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             func = self.cases.get('__default__')
 
         if not func:
-            raise Exception("Value does not match any case and there is no default case: value {}".format(self.value))
+            raise Exception("Value does not match any case and there "
+                            "is no default case: value {}".format(self.value))
+
+        # noinspection PyCallingNonCallable
+        self.__result = func()
+
+    @property
+    def result(self):
+        if self.__result == switch.__no_result:
+            raise Exception("No result has been computed.")
 
-        func()
+        return self.__result
diff --git a/tests/coretests.py b/tests/coretests.py
@@ -132,3 +132,18 @@ def get_set_case(val):
             s.default(lambda: get_set_case('default'))
 
         self.assertEqual(executed_case, "even")
+
+    def test_return_value_from_case(self):
+        value = 4
+        with switch(value) as s:
+            s.case([1, 3, 5, 7], lambda: value + 1)
+            s.case([0, 2, 4, 6, 8], lambda: value * value)
+            s.default(lambda: 0)
+
+        self.assertEqual(s.result, 16)
+
+    # noinspection PyStatementEffect
+    def test_result_inaccessible_if_hasnt_run(self):
+        with self.assertRaises(Exception):
+            s = switch(7)
+            s.result
