Repository: streamlit/streamlit

Issue:
Make sure magic works with all block statements in Python
Magic should work in "if" blocks, "while" blocks, "for", "try", "with", "async", etc...

https://docs.python.org/3/reference/compound_stmts.html


Discussion:

Resolving patch:
diff --git a/examples/magic.py b/examples/magic.py
new file mode 100644
--- /dev/null
+++ b/examples/magic.py
@@ -0,0 +1,128 @@
+import asyncio
+import contextlib
+
+import streamlit as st
+
+async_loop = asyncio.new_event_loop()
+
+
+magic = None
+
+
+def prep_magic(name):
+    global magic
+    magic = "_%s_" % name
+    st.markdown("**%s** expected:" % name)
+
+
+prep_magic("no block")
+magic
+
+prep_magic("IF block")
+if True:
+    magic
+
+prep_magic("ELIF block")
+if False:
+    pass
+elif True:
+    magic
+
+prep_magic("ELSE block")
+if False:
+    pass
+else:
+    magic
+
+
+prep_magic("FOR block")
+for ii in range(1):
+    magic
+
+prep_magic("WHILE block")
+ii = 0
+while ii < 1:
+    magic
+    ii += 1
+
+prep_magic("WITH block")
+
+
+@contextlib.contextmanager
+def context_mgr():
+    try:
+        yield
+    finally:
+        pass
+
+
+with context_mgr():
+    magic
+
+prep_magic("TRY block")
+try:
+    magic
+except:
+    raise
+
+prep_magic("EXCEPT block")
+try:
+    raise RuntimeError("shenanigans!")
+except RuntimeError:
+    magic
+
+prep_magic("FINALLY block")
+try:
+    pass
+finally:
+    magic
+
+prep_magic("FUNCTION block")
+
+
+def func():
+    magic
+
+
+func()
+
+prep_magic("ASYNC FUNCTION block")
+
+
+async def async_func():
+    magic
+
+
+async_loop.run_until_complete(async_func())
+
+prep_magic("ASYNC FOR block")
+
+
+async def async_for():
+    async def async_iter():
+        yield
+
+    async for _ in async_iter():
+        magic
+
+
+async_loop.run_until_complete(async_for())
+
+prep_magic("ASYNC WITH block")
+
+
+async def async_with():
+    @contextlib.asynccontextmanager
+    async def async_context_mgr():
+        try:
+            yield
+        finally:
+            pass
+
+    async with async_context_mgr():
+        magic
+
+
+async_loop.run_until_complete(async_with())
+
+st.info("done")
diff --git a/lib/streamlit/magic.py b/lib/streamlit/magic.py
--- a/lib/streamlit/magic.py
+++ b/lib/streamlit/magic.py
@@ -53,6 +53,10 @@ def _modify_ast_subtree(tree, body_attr="body", is_root=False):
             node_type is ast.FunctionDef
             or node_type is ast.With
             or node_type is ast.For
+            or node_type is ast.While
+            or node_type is ast.AsyncFunctionDef
+            or node_type is ast.AsyncWith
+            or node_type is ast.AsyncFor
         ):
             _modify_ast_subtree(node)
 
@@ -69,7 +73,7 @@ def _modify_ast_subtree(tree, body_attr="body", is_root=False):
             _modify_ast_subtree(node)
             _modify_ast_subtree(node, "orelse")
 
-        # Convert expression nodes to st.write
+        # Convert standalone expression nodes to st.write
         elif node_type is ast.Expr:
             value = _get_st_write_from_expr(node, i, parent_type=type(tree))
             if value is not None:
@@ -142,6 +146,10 @@ def _get_st_write_from_expr(node, i, parent_type):
     ):
         return None
 
+    # Don't change yield nodes
+    if type(node.value) is ast.Yield or type(node.value) is ast.YieldFrom:
+        return None
+
     # If tuple, call st.write on the 0th element (rather than the
     # whole tuple). This allows us to add a comma at the end of a statement
     # to turn it into an expression that should be st-written. Ex:


Test patch:
diff --git a/lib/tests/streamlit/magic_test.py b/lib/tests/streamlit/magic_test.py
--- a/lib/tests/streamlit/magic_test.py
+++ b/lib/tests/streamlit/magic_test.py
@@ -43,9 +43,9 @@ def _testCode(self, code, expected_count):
             if type(node) is ast.Call and "streamlit" in ast.dump(node.func):
                 count += 1
         self.assertEqual(
-            count,
             expected_count,
-            ("There must be at least {} streamlit nodes," "but found {}").format(
+            count,
+            ("There must be exactly {} streamlit nodes, but found {}").format(
                 expected_count, count
             ),
         )
@@ -122,3 +122,54 @@ def test_with_statement(self):
     a
 """
         self._testCode(CODE_WITH_STATEMENT, 1)
+
+    def test_while_statement(self):
+        """Test that 'yield' expressions do not get magicked"""
+        CODE_WHILE_STATEMENT = """
+a = 10
+while True:
+    a
+"""
+        self._testCode(CODE_WHILE_STATEMENT, 1)
+
+    def test_yield_statement(self):
+        """Test that 'yield' expressions do not get magicked"""
+        CODE_YIELD_STATEMENT = """
+def yield_func():
+    yield
+"""
+        self._testCode(CODE_YIELD_STATEMENT, 0)
+
+    def test_yield_from_statement(self):
+        """Test that 'yield from' expressions do not get magicked"""
+        CODE_YIELD_FROM_STATEMENT = """
+def yield_func():
+    yield from None
+"""
+        self._testCode(CODE_YIELD_FROM_STATEMENT, 0)
+
+    def test_async_function_statement(self):
+        """Test async function definitions"""
+        CODE_ASYNC_FUNCTION = """
+async def myfunc(a):
+    a
+"""
+        self._testCode(CODE_ASYNC_FUNCTION, 1)
+
+    def test_async_with_statement(self):
+        """Test 'async with' statements"""
+        CODE_ASYNC_WITH = """
+async def myfunc(a):
+    async with None:
+        a
+"""
+        self._testCode(CODE_ASYNC_WITH, 1)
+
+    def test_async_for_statement(self):
+        """Test 'async for' statements"""
+        CODE_ASYNC_FOR = """
+async def myfunc(a):
+    async for _ in None:
+        a
+"""
+        self._testCode(CODE_ASYNC_FOR, 1)