Tweak cancellation output

Author: dandavison

hello/hello_cancellation.py

@@ -1,6 +1,5 @@
 import asyncio
 import time
-import traceback
 from concurrent.futures import ThreadPoolExecutor
 from datetime import timedelta
 from typing import NoReturn
@@ -19,9 +18,11 @@ def never_complete_activity() -> NoReturn:
         while True:
             print("Heartbeating activity")
             activity.heartbeat()
-            time.sleep(1)
-    except CancelledError:
-        print("Activity cancelled")
+            time.sleep(0.5)
+    except CancelledError as err:
+        print(
+            f"Got exception in activity. Cause chain is {format_exception_cause_chain(err)}"
+        )
         raise
 
 
@@ -43,6 +44,11 @@ async def run(self) -> None:
                 # Always set a heartbeat timeout for long-running activities
                 heartbeat_timeout=timedelta(seconds=2),
             )
+        except Exception as err:
+            print(
+                f"Got exception in workflow. Cause chain is {format_exception_cause_chain(err)}"
+            )
+            raise
         finally:
             await workflow.execute_activity(
                 cleanup_activity, start_to_close_timeout=timedelta(seconds=5)
@@ -61,7 +67,6 @@ async def main():
         activities=[never_complete_activity, cleanup_activity],
         activity_executor=ThreadPoolExecutor(5),
     ):
-
         # While the worker is running, use the client to start the workflow.
         # Note, in many production setups, the client would be in a completely
         # separate process from the worker.
@@ -72,16 +77,25 @@ async def main():
         )
 
         # Now that we've started, wait a couple of seconds then cancel it
-        await asyncio.sleep(2)
+        await asyncio.sleep(1)
         await handle.cancel()
 
         # Now wait on the result which we expect will fail since it was
         # cancelled
         try:
             await handle.result()
             raise RuntimeError("Should not succeed")
-        except WorkflowFailureError:
-            print("Got expected exception: ", traceback.format_exc())
+        except WorkflowFailureError as err:
+            print(
+                f"Got expected exception in client. Cause chain is {format_exception_cause_chain(err)}"
+            )
+
+
+def format_exception_cause_chain(err: BaseException) -> str:
+    causes = [err]
+    while causes[-1].__cause__:
+        causes.append(causes[-1].__cause__)
+    return " -> ".join([f"{e.__class__.__name__}" for e in causes])
 
 
 if __name__ == "__main__":