Add type annotations for exception (#4681)

* Add type annotations for exception

* Add some defensive code

* Formatting

* Add comment

* Conditionally import from typing_extensions

* Import sys

* Simplify constant definition and add type

* Unify cast calls

* Revert branching import

lib/streamlit/__init__.py

@@ -311,8 +311,17 @@ def experimental_show(*args: Any) -> None:
             markdown("**%s**" % escaped)
             write(args[idx])
 
-    except Exception:
+    except Exception as raised_exc:
         _, exc, exc_tb = _sys.exc_info()
+        if exc is None:
+            # Presumably, exc should never be None, but it is typed as
+            # Optional, and I don't know the internals of sys.exc_info() well
+            # enough to just use a cast here. Hence, the runtime check.
+            raise RuntimeError(
+                "Unexpected state: exc was None. If you see this message, "
+                "please create an issue at "
+                "https://github.com/streamlit/streamlit/issues"
+            ) from raised_exc
         exception(exc)
 
 

lib/streamlit/elements/exception.py

@@ -14,8 +14,8 @@
 
 import os
 import traceback
-import typing
-from typing import Optional, cast, List
+from typing import cast, List, Optional, TYPE_CHECKING
+from typing_extensions import Final
 
 import streamlit
 from streamlit.proto.Exception_pb2 import Exception as ExceptionProto
@@ -26,22 +26,24 @@
 from streamlit.errors import UncaughtAppException
 from streamlit.logger import get_logger
 
-LOGGER = get_logger(__name__)
+if TYPE_CHECKING:
+    from streamlit.delta_generator import DeltaGenerator
+
+LOGGER: Final = get_logger(__name__)
 
 # When client.showErrorDetails is False, we show a generic warning in the
 # frontend when we encounter an uncaught app exception.
-_GENERIC_UNCAUGHT_EXCEPTION_TEXT = "This app has encountered an error. The original error message is redacted to prevent data leaks.  Full error details have been recorded in the logs (if you're on Streamlit Cloud, click on 'Manage app' in the lower right of your app)."
-
-# Extract the streamlit package path
-_STREAMLIT_DIR = os.path.dirname(streamlit.__file__)
+_GENERIC_UNCAUGHT_EXCEPTION_TEXT: Final = "This app has encountered an error. The original error message is redacted to prevent data leaks.  Full error details have been recorded in the logs (if you're on Streamlit Cloud, click on 'Manage app' in the lower right of your app)."
 
-# Make it absolute, resolve aliases, and ensure there's a trailing path
-# separator
-_STREAMLIT_DIR = os.path.join(os.path.realpath(_STREAMLIT_DIR), "")
+# Extract the streamlit package path. Make it absolute, resolve aliases, and
+# ensure there's a trailing path separator
+_STREAMLIT_DIR: Final = os.path.join(
+    os.path.realpath(os.path.dirname(streamlit.__file__)), ""
+)
 
 
 class ExceptionMixin:
-    def exception(self, exception):
+    def exception(self, exception: BaseException) -> "DeltaGenerator":
         """Display an exception.
 
         Parameters
@@ -57,12 +59,13 @@ def exception(self, exception):
         """
         exception_proto = ExceptionProto()
         marshall(exception_proto, exception)
-        return self.dg._enqueue("exception", exception_proto)
+        dg = self.dg._enqueue("exception", exception_proto)
+        return cast("DeltaGenerator", dg)
 
     @property
-    def dg(self) -> "streamlit.delta_generator.DeltaGenerator":
+    def dg(self) -> "DeltaGenerator":
         """Get our DeltaGenerator."""
-        return cast("streamlit.delta_generator.DeltaGenerator", self)
+        return cast("DeltaGenerator", self)
 
 
 def marshall(exception_proto: ExceptionProto, exception: BaseException) -> None:
@@ -140,7 +143,7 @@ def marshall(exception_proto: ExceptionProto, exception: BaseException) -> None:
         )
 
     if is_uncaught_app_exception:
-        uae = typing.cast(UncaughtAppException, exception)
+        uae = cast(UncaughtAppException, exception)
         exception_proto.message = _GENERIC_UNCAUGHT_EXCEPTION_TEXT
         type_str = str(type(uae.exc))
         exception_proto.type = type_str.replace("<class '", "").replace("'>", "")