Update the codemod to await st.write()

whitphx · whitphx · commit 5172a43f2ad9 · 2025-05-17T16:57:25.000-04:00
diff --git a/packages/kernel/py/stlite-lib/stlite_lib/codemod.py b/packages/kernel/py/stlite-lib/stlite_lib/codemod.py
@@ -911,6 +911,7 @@ def patch(code: str | ast.Module, script_path: str) -> ast.Module:
     wildcard_import_monitor_targets = set(
         [
             WildcardImportTarget(module="time", attr="sleep"),
+            WildcardImportTarget(module="streamlit", attr="write"),
             WildcardImportTarget(module="streamlit", attr="write_stream"),
             WildcardImportTarget(module="asyncio", attr="run"),
         ]
@@ -924,6 +925,7 @@ def patch(code: str | ast.Module, script_path: str) -> ast.Module:
     func_call_handler = FuncCallTransformHandler(
         {
             FunctionCall(name="time.sleep"): TransformRuleAction.TIME_SLEEP,
+            FunctionCall(name="streamlit.write"): TransformRuleAction.AWAIT_CALL,
             FunctionCall(name="streamlit.write_stream"): TransformRuleAction.AWAIT_CALL,
             AttrFunctionCall(
                 obj=ReturnValue(called_function="streamlit.navigation"),
diff --git a/streamlit b/streamlit
@@ -1 +1 @@
-Subproject commit 4dfc2e692d3126dec566e46babdd360d2d4c1bdb
+Subproject commit 61b7eb578df81310c421b8ebc8985424dcf4b79d

Original file line number	Diff line number	Diff line change
`@@ -911,6 +911,7 @@ def patch(code: str \| ast.Module, script_path: str) -> ast.Module:`
`911`	`911`	`wildcard_import_monitor_targets = set(`
`912`	`912`	`[`
`913`	`913`	`WildcardImportTarget(module="time", attr="sleep"),`
	`914`	`+ WildcardImportTarget(module="streamlit", attr="write"),`
`914`	`915`	`WildcardImportTarget(module="streamlit", attr="write_stream"),`
`915`	`916`	`WildcardImportTarget(module="asyncio", attr="run"),`
`916`	`917`	`]`
`@@ -924,6 +925,7 @@ def patch(code: str \| ast.Module, script_path: str) -> ast.Module:`
`924`	`925`	`func_call_handler = FuncCallTransformHandler(`
`925`	`926`	`{`
`926`	`927`	`FunctionCall(name="time.sleep"): TransformRuleAction.TIME_SLEEP,`
	`928`	`+ FunctionCall(name="streamlit.write"): TransformRuleAction.AWAIT_CALL,`
`927`	`929`	`FunctionCall(name="streamlit.write_stream"): TransformRuleAction.AWAIT_CALL,`
`928`	`930`	`AttrFunctionCall(`
`929`	`931`	`obj=ReturnValue(called_function="streamlit.navigation"),`