st.echo: improved readability
diff --git a/lib/streamlit/__init__.py b/lib/streamlit/__init__.py
index cd925b6762bf..2b8452d53579 100644
--- a/lib/streamlit/__init__.py
+++ b/lib/streamlit/__init__.py
@@ -52,7 +52,6 @@
 
 # Give the package a version.
 import pkg_resources as _pkg_resources
-from typing import List
 from typing import NoReturn
 
 # This used to be pkg_resources.require('streamlit') but it would cause
@@ -60,11 +59,8 @@
 __version__ = _pkg_resources.get_distribution("streamlit").version
 
 import contextlib as _contextlib
-import re as _re
 import sys as _sys
-import textwrap as _textwrap
 import threading as _threading
-import traceback as _traceback
 import urllib.parse as _parse
 
 import click as _click
@@ -84,6 +80,8 @@
 
 # Modules that the user should have access to. These are imported with "as"
 # syntax pass mypy checking with implicit_reexport disabled.
+
+from streamlit.echo import echo as echo
 from streamlit.legacy_caching import cache as cache
 from streamlit.caching import singleton as experimental_singleton
 from streamlit.caching import memo as experimental_memo
@@ -435,61 +433,6 @@ def set_message():
                 message.empty()
 
 
-_SPACES_RE = _re.compile("\\s*")
-
-
-@_contextlib.contextmanager
-def echo(code_location="above"):
-    """Use in a `with` block to draw some code on the app, then execute it.
-
-    Parameters
-    ----------
-    code_location : "above" or "below"
-        Whether to show the echoed code before or after the results of the
-        executed code block.
-
-    Example
-    -------
-
-    >>> with st.echo():
-    >>>     st.write('This code will be printed')
-
-    """
-    if code_location == "below":
-        show_code = code
-        show_warning = warning
-    else:
-        placeholder = empty()
-        show_code = placeholder.code
-        show_warning = placeholder.warning
-
-    try:
-        frame = _traceback.extract_stack()[-3]
-        filename, start_line = frame.filename, frame.lineno
-        yield
-        frame = _traceback.extract_stack()[-3]
-        end_line = frame.lineno
-        lines_to_display = []  # type: List[str]
-        with _source_util.open_python_file(filename) as source_file:
-            source_lines = source_file.readlines()
-            lines_to_display.extend(source_lines[start_line:end_line])
-            match = _SPACES_RE.match(lines_to_display[0])
-            initial_spaces = match.end() if match else 0
-            for line in source_lines[end_line:]:
-                match = _SPACES_RE.match(line)
-                indentation = match.end() if match else 0
-                # The != 1 is because we want to allow '\n' between sections.
-                if indentation != 1 and indentation < initial_spaces:
-                    break
-                lines_to_display.append(line)
-        line_to_display = _textwrap.dedent("".join(lines_to_display))
-
-        show_code(line_to_display, "python")
-
-    except FileNotFoundError as err:
-        show_warning("Unable to display code. %s" % err)
-
-
 def _transparent_write(*args):
     """This is just st.write, but returns the arguments you passed to it."""
     write(*args)
diff --git a/lib/streamlit/echo.py b/lib/streamlit/echo.py
new file mode 100644
index 000000000000..ac48625e30f0
--- /dev/null
+++ b/lib/streamlit/echo.py
@@ -0,0 +1,95 @@
+# Copyright 2018-2021 Streamlit Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import re
+import textwrap
+import traceback
+from typing import List
+
+_SPACES_RE = re.compile("\\s*")
+
+
+@contextlib.contextmanager
+def echo(code_location="above"):
+    """Use in a `with` block to draw some code on the app, then execute it.
+
+    Parameters
+    ----------
+    code_location : "above" or "below"
+        Whether to show the echoed code before or after the results of the
+        executed code block.
+
+    Example
+    -------
+
+    >>> with st.echo():
+    >>>     st.write('This code will be printed')
+
+    """
+
+    from streamlit import code, warning, empty, source_util
+
+    if code_location == "below":
+        show_code = code
+        show_warning = warning
+    else:
+        placeholder = empty()
+        show_code = placeholder.code
+        show_warning = placeholder.warning
+
+    try:
+        # Get stack frame *before* running the echoed code. The frame's
+        # line number will point to the `st.echo` statement we're running.
+        frame = traceback.extract_stack()[-3]
+        filename, start_line = frame.filename, frame.lineno
+
+        # Run the echoed code.
+        yield
+
+        # Get stack frame *after* running code. This frame's line number will
+        # point to the last line in the echoed code.
+        frame = traceback.extract_stack()[-3]
+        end_line = frame.lineno
+
+        # Open the file containing the source code of the echoed statement,
+        # and extract the lines inside the `with st.echo` block.
+        lines_to_display: List[str] = []
+        with source_util.open_python_file(filename) as source_file:
+            source_lines = source_file.readlines()
+            lines_to_display.extend(source_lines[start_line:end_line])
+
+            # Our "end_line" is not necessarily the last line in the echo
+            # block. Iterate over the remaining lines in the source file
+            # until we find one that's indented less than the rest of the
+            # block. Note that this is *not* a perfect strategy, because
+            # de-denting is not guaranteed to signal "end of block". (A
+            # triple-quoted string might be dedented but still in the
+            # echo block, for example.)
+            if len(lines_to_display) > 0:
+                match = _SPACES_RE.match(lines_to_display[0])
+                initial_spaces = match.end() if match else 0
+                for line in source_lines[end_line:]:
+                    match = _SPACES_RE.match(line)
+                    indentation = match.end() if match else 0
+                    # The != 1 is because we want to allow '\n' between sections.
+                    if indentation != 1 and indentation < initial_spaces:
+                        break
+                    lines_to_display.append(line)
+
+        line_to_display = textwrap.dedent("".join(lines_to_display))
+        show_code(line_to_display, "python")
+
+    except FileNotFoundError as err:
+        show_warning("Unable to display code. %s" % err)
diff --git a/lib/streamlit/elements/doc_string.py b/lib/streamlit/elements/doc_string.py
index 9383162d6979..820f83857a55 100644
--- a/lib/streamlit/elements/doc_string.py
+++ b/lib/streamlit/elements/doc_string.py
@@ -26,6 +26,7 @@
 
 
 CONFUSING_STREAMLIT_MODULES = (
+    "streamlit.echo",
     "streamlit.delta_generator",
     "streamlit.legacy_caching.caching",
 )
diff --git a/lib/tests/streamlit/echo_test.py b/lib/tests/streamlit/echo_test.py
index d4cbf78a1c5a..eff98df0aa84 100644
--- a/lib/tests/streamlit/echo_test.py
+++ b/lib/tests/streamlit/echo_test.py
@@ -12,37 +12,40 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from parameterized import parameterized
+
 from tests import testutil
 import streamlit as st
 
 
 class EchoTest(testutil.DeltaGeneratorTestCase):
-    def test_echo(self):
-        for echo, echo_index, output_index in [
-            (lambda: st.echo(), 0, 1),
-            (lambda: st.echo("above"), 0, 1),
-            (lambda: st.echo("below"), 1, 0),
-        ]:
-
-            # The empty lines below are part of the test. Do not remove them.
-            with echo():
-                st.write("Hello")
+    @parameterized.expand(
+        [
+            ("code_location default", lambda: st.echo(), 0, 1),
+            ("code_location above", lambda: st.echo("above"), 0, 1),
+            ("code_location below", lambda: st.echo("below"), 1, 0),
+        ]
+    )
+    def test_echo(self, _, echo, echo_index, output_index):
+        # The empty lines below are part of the test. Do not remove them.
+        with echo():
+            st.write("Hello")
 
-                "hi"
+            "hi"
 
-                def foo(x):
-                    y = x + 10
+            def foo(x):
+                y = x + 10
 
-                    print(y)
+                print(y)
 
-                class MyClass(object):
-                    def do_x(self):
-                        pass
+            class MyClass(object):
+                def do_x(self):
+                    pass
 
-                    def do_y(self):
-                        pass
+                def do_y(self):
+                    pass
 
-            echo_str = """```python
+        echo_str = """```python
 st.write("Hello")
 
 "hi"
@@ -62,13 +65,13 @@ def do_y(self):
 
 ```"""
 
-            element = self.get_delta_from_queue(echo_index).new_element
-            self.assertEqual(echo_str, element.markdown.body)
+        element = self.get_delta_from_queue(echo_index).new_element
+        self.assertEqual(echo_str, element.markdown.body)
 
-            element = self.get_delta_from_queue(output_index).new_element
-            self.assertEqual("Hello", element.markdown.body)
+        element = self.get_delta_from_queue(output_index).new_element
+        self.assertEqual("Hello", element.markdown.body)
 
-            self.clear_queue()
+        self.clear_queue()
 
     def test_root_level_echo(self):
         import tests.streamlit.echo_test_data.root_level_echo

Some drive-by cleanup on st.echo as I was investigating a related issue

- Pulls `st.echo` out of `__init__.py` and into `echo.py`
- `st.echo()` gets more inline comments explaining what the code is doing
- `EchoTest.test_echo` is parameterized