From 5dc6684b313e3cfc37a45131f7865dbe28daae72 Mon Sep 17 00:00:00 2001 From: "Yuichiro Tachibana (Tsuchiya)" Date: Thu, 17 Oct 2024 16:17:42 +0900 Subject: [PATCH] Fix codemod to handle function calls in assignments --- .../py/stlite-lib/stlite_lib/codemod.py | 3 +++ .../stlite_lib_tests/codemod_test.py | 26 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/packages/kernel/py/stlite-lib/stlite_lib/codemod.py b/packages/kernel/py/stlite-lib/stlite_lib/codemod.py index ebd956424..839336a8f 100644 --- a/packages/kernel/py/stlite-lib/stlite_lib/codemod.py +++ b/packages/kernel/py/stlite-lib/stlite_lib/codemod.py @@ -626,6 +626,9 @@ def visit(self, node: ast.AST) -> ast.AST: ) for target in node.targets: self._bind_expr(target, bound_to=bound_to) + + self.generic_visit(node) + return node assignment_occurs = isinstance( diff --git a/packages/kernel/py/stlite-lib/stlite_lib_tests/codemod_test.py b/packages/kernel/py/stlite-lib/stlite_lib_tests/codemod_test.py index 2515d91aa..ad95d5237 100644 --- a/packages/kernel/py/stlite-lib/stlite_lib_tests/codemod_test.py +++ b/packages/kernel/py/stlite-lib/stlite_lib_tests/codemod_test.py @@ -66,6 +66,32 @@ """ import streamlit as st +a = st.write_stream("Hello, world!") +""", + """ +import streamlit as st + +a = await st.write_stream("Hello, world!") +""", + id="call_in_assignment", + ), + pytest.param( + """ +from streamlit import write_stream + +write_stream("Hello, world!") +""", + """ +from streamlit import write_stream + +await write_stream("Hello, world!") +""", + id="from_streamlit_import_write_stream", + ), + pytest.param( + """ +import streamlit as st + if True: st.write_stream("Hello, world!") elif False: