Add more type annotations (#4657)

* Add some type annotations to __init__.py

* Add type annotations to markdown.py

* Add type annotations to text.py

* Update markdown.py

Import Optional and avoid using DeltaGenerator directly in cast.

* Fix issues in markdown.py

* Fix issues in text.py

* Formatting

* Yet another missing forward ref

* Cast Any to DeltaGenerator

* Relax dg.text body type to Any

* Fix test

* Formatting

* Use typing_extensions.Final for backwards compatibility

* Add special handling for python 3.6

* Update special case to apply to python 3.7

* Compact imports

lib/streamlit/__init__.py

@@ -206,7 +206,7 @@ def _update_logger():
 beta_columns = _main.beta_columns
 
 
-def set_option(key, value):
+def set_option(key: str, value) -> None:
     """Set config option.
 
     Currently, only the following config options can be set within the script itself:
@@ -372,7 +372,7 @@ def experimental_set_query_params(**query_params):
 
 
 @_contextlib.contextmanager
-def spinner(text="In progress..."):
+def spinner(text: str = "In progress..."):
     """Temporarily displays a message while executing a block of code.
 
     Parameters
@@ -446,7 +446,7 @@ def _transparent_write(*args):
 _use_warning_has_been_displayed = False
 
 
-def _maybe_print_use_warning():
+def _maybe_print_use_warning() -> None:
     """Print a warning if Streamlit is imported but not being run with `streamlit run`.
     The warning is printed only once.
     """
@@ -492,7 +492,7 @@ def stop() -> NoReturn:
     raise StopException()
 
 
-def experimental_rerun():
+def experimental_rerun() -> NoReturn:
     """Rerun the script immediately.
 
     When `st.experimental_rerun()` is called, the script is halted - no

lib/streamlit/elements/alert.py

@@ -12,15 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import cast
+from typing import cast, TYPE_CHECKING
 
-import streamlit
 from streamlit.proto.Alert_pb2 import Alert as AlertProto
 from .utils import clean_text
 
+if TYPE_CHECKING:
+    from streamlit.delta_generator import DeltaGenerator
+
 
 class AlertMixin:
-    def error(self, body):
+    def error(self, body: str) -> "DeltaGenerator":
         """Display error message.
 
         Parameters
@@ -36,9 +38,10 @@ def error(self, body):
         alert_proto = AlertProto()
         alert_proto.body = clean_text(body)
         alert_proto.format = AlertProto.ERROR
-        return self.dg._enqueue("alert", alert_proto)
+        dg = self.dg._enqueue("alert", alert_proto)
+        return cast("DeltaGenerator", dg)
 
-    def warning(self, body):
+    def warning(self, body: str) -> "DeltaGenerator":
         """Display warning message.
 
         Parameters
@@ -54,9 +57,10 @@ def warning(self, body):
         alert_proto = AlertProto()
         alert_proto.body = clean_text(body)
         alert_proto.format = AlertProto.WARNING
-        return self.dg._enqueue("alert", alert_proto)
+        dg = self.dg._enqueue("alert", alert_proto)
+        return cast("DeltaGenerator", dg)
 
-    def info(self, body):
+    def info(self, body: str) -> "DeltaGenerator":
         """Display an informational message.
 
         Parameters
@@ -72,9 +76,10 @@ def info(self, body):
         alert_proto = AlertProto()
         alert_proto.body = clean_text(body)
         alert_proto.format = AlertProto.INFO
-        return self.dg._enqueue("alert", alert_proto)
+        dg = self.dg._enqueue("alert", alert_proto)
+        return cast("DeltaGenerator", dg)
 
-    def success(self, body):
+    def success(self, body: str) -> "DeltaGenerator":
         """Display a success message.
 
         Parameters
@@ -90,9 +95,10 @@ def success(self, body):
         alert_proto = AlertProto()
         alert_proto.body = clean_text(body)
         alert_proto.format = AlertProto.SUCCESS
-        return self.dg._enqueue("alert", alert_proto)
+        dg = self.dg._enqueue("alert", alert_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)

lib/streamlit/elements/button.py

@@ -17,6 +17,7 @@
 
 from streamlit.type_util import Key, to_key
 from typing import cast, Optional, Union, BinaryIO, TextIO
+from typing_extensions import Final
 from textwrap import dedent
 
 import streamlit
@@ -34,7 +35,7 @@
 from .utils import check_callback_rules, check_session_state_rules
 
 
-FORM_DOCS_INFO = """
+FORM_DOCS_INFO: Final = """
 
 For more information, refer to the
 [documentation for forms](https://docs.streamlit.io/library/api-reference/control-flow/st.form).

lib/streamlit/elements/markdown.py

@@ -12,16 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import cast
+from typing import cast, Optional, TYPE_CHECKING, Union
 
-import streamlit
 from streamlit import type_util
 from streamlit.proto.Markdown_pb2 import Markdown as MarkdownProto
 from .utils import clean_text
 
+if TYPE_CHECKING:
+    import sympy
+
+    from streamlit.delta_generator import DeltaGenerator
+
 
 class MarkdownMixin:
-    def markdown(self, body, unsafe_allow_html=False):
+    def markdown(self, body: str, unsafe_allow_html: bool = False) -> "DeltaGenerator":
         """Display string formatted as Markdown.
 
         Parameters
@@ -72,9 +76,10 @@ def markdown(self, body, unsafe_allow_html=False):
         markdown_proto.body = clean_text(body)
         markdown_proto.allow_html = unsafe_allow_html
 
-        return self.dg._enqueue("markdown", markdown_proto)
+        dg = self.dg._enqueue("markdown", markdown_proto)
+        return cast("DeltaGenerator", dg)
 
-    def header(self, body, anchor=None):
+    def header(self, body: str, anchor: Optional[str] = None) -> "DeltaGenerator":
         """Display text in header formatting.
 
         Parameters
@@ -97,9 +102,10 @@ def header(self, body, anchor=None):
         else:
             header_proto.body = f'<h2 data-anchor="{anchor}">{clean_text(body)}</h2>'
             header_proto.allow_html = True
-        return self.dg._enqueue("markdown", header_proto)
+        dg = self.dg._enqueue("markdown", header_proto)
+        return cast("DeltaGenerator", dg)
 
-    def subheader(self, body, anchor=None):
+    def subheader(self, body: str, anchor: Optional[str] = None) -> "DeltaGenerator":
         """Display text in subheader formatting.
 
         Parameters
@@ -123,9 +129,10 @@ def subheader(self, body, anchor=None):
             subheader_proto.body = f'<h3 data-anchor="{anchor}">{clean_text(body)}</h3>'
             subheader_proto.allow_html = True
 
-        return self.dg._enqueue("markdown", subheader_proto)
+        dg = self.dg._enqueue("markdown", subheader_proto)
+        return cast("DeltaGenerator", dg)
 
-    def code(self, body, language="python"):
+    def code(self, body: str, language: Optional[str] = "python") -> "DeltaGenerator":
         """Display a code block with optional syntax highlighting.
 
         (This is a convenience wrapper around `st.markdown()`)
@@ -152,9 +159,10 @@ def code(self, body, language="python"):
             "body": body,
         }
         code_proto.body = clean_text(markdown)
-        return self.dg._enqueue("markdown", code_proto)
+        dg = self.dg._enqueue("markdown", code_proto)
+        return cast("DeltaGenerator", dg)
 
-    def title(self, body, anchor=None):
+    def title(self, body: str, anchor: Optional[str] = None) -> "DeltaGenerator":
         """Display text in title formatting.
 
         Each document should have a single `st.title()`, although this is not
@@ -180,9 +188,10 @@ def title(self, body, anchor=None):
         else:
             title_proto.body = f'<h1 data-anchor="{anchor}">{clean_text(body)}</h1>'
             title_proto.allow_html = True
-        return self.dg._enqueue("markdown", title_proto)
+        dg = self.dg._enqueue("markdown", title_proto)
+        return cast("DeltaGenerator", dg)
 
-    def caption(self, body, unsafe_allow_html=False):
+    def caption(self, body: str, unsafe_allow_html: bool = False) -> "DeltaGenerator":
         """Display text in small font.
 
         This should be used for captions, asides, footnotes, sidenotes, and
@@ -223,9 +232,10 @@ def caption(self, body, unsafe_allow_html=False):
         caption_proto.body = clean_text(body)
         caption_proto.allow_html = unsafe_allow_html
         caption_proto.is_caption = True
-        return self.dg._enqueue("markdown", caption_proto)
+        dg = self.dg._enqueue("markdown", caption_proto)
+        return cast("DeltaGenerator", dg)
 
-    def latex(self, body):
+    def latex(self, body: Union[str, "sympy.Expr"]) -> "DeltaGenerator":
         # This docstring needs to be "raw" because of the backslashes in the
         # example below.
         r"""Display mathematical expressions formatted as LaTeX.
@@ -257,9 +267,10 @@ def latex(self, body):
 
         latex_proto = MarkdownProto()
         latex_proto.body = "$$\n%s\n$$" % clean_text(body)
-        return self.dg._enqueue("markdown", latex_proto)
+        dg = self.dg._enqueue("markdown", latex_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)

lib/streamlit/elements/text.py

@@ -12,15 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import cast
+from typing import Any, cast, TYPE_CHECKING
 
-import streamlit
 from streamlit.proto.Text_pb2 import Text as TextProto
 from .utils import clean_text
 
+if TYPE_CHECKING:
+    from streamlit.delta_generator import DeltaGenerator
+
 
 class TextMixin:
-    def text(self, body):
+    def text(self, body: Any) -> "DeltaGenerator":
         """Write fixed-width and preformatted text.
 
         Parameters
@@ -35,9 +37,10 @@ def text(self, body):
         """
         text_proto = TextProto()
         text_proto.body = clean_text(body)
-        return self.dg._enqueue("text", text_proto)
+        dg = self.dg._enqueue("text", text_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)

lib/tests/streamlit/streamlit_test.py

@@ -20,6 +20,7 @@
 import os
 import io
 import re
+import sys
 import time
 import textwrap
 import unittest
@@ -419,7 +420,17 @@ def test_st_help(self):
             el.doc_string.doc_string.startswith("Display text in header formatting.")
         )
         self.assertEqual(el.doc_string.type, "<class 'method'>")
-        self.assertEqual(el.doc_string.signature, "(body, anchor=None)")
+        if sys.version_info[1] == 7:
+            # Python 3.7 represents the signature slightly differently
+            self.assertEqual(
+                el.doc_string.signature,
+                "(body: str, anchor: Union[str, NoneType] = None) -> 'DeltaGenerator'",
+            )
+        else:
+            self.assertEqual(
+                el.doc_string.signature,
+                "(body: str, anchor: Optional[str] = None) -> 'DeltaGenerator'",
+            )
 
     def test_st_image_PIL_image(self):
         """Test st.image with PIL image."""