Update flexdown parser to handle all Python statement types

Co-Authored-By: Alek Petuskey <[email protected]>

pcweb/flexdown.py

@@ -315,25 +315,24 @@ def render(self, env) -> rx.Component:
         elif exec_mode:
             return comp
         elif "box" in args:
-            # Check if the code contains a function definition
-            if code.strip().startswith("def "):
-                # Use exec for function definitions
+            # Use exec for any code block that can't be evaluated directly
+            if any(code.strip().startswith(stmt) for stmt in ["def ", "class ", "import ", "from "]):
                 temp_env = env.copy()
                 exec(code, temp_env, temp_env)
-                # Get the function name (first line after 'def ' until '(')
-                func_name = code.strip()[4:].split("(")[0].strip()
-                # Get the function from the environment and call it
-                comp = temp_env[func_name]()
+                # Get the last defined object from the environment
+                last_key = list(temp_env.keys())[-1]
+                # If it's a function or class, instantiate it
+                comp = temp_env[last_key]() if callable(temp_env[last_key]) else temp_env[last_key]
             else:
                 comp = eval(code, env, env)
             return rx.box(docdemobox(comp), margin_bottom="1em", id=comp_id)
         else:
-            # For non-box mode, still handle function definitions
-            if code.strip().startswith("def "):
+            # For non-box mode, handle all code types
+            if any(code.strip().startswith(stmt) for stmt in ["def ", "class ", "import ", "from "]):
                 temp_env = env.copy()
                 exec(code, temp_env, temp_env)
-                func_name = code.strip()[4:].split("(")[0].strip()
-                comp = temp_env[func_name]()
+                last_key = list(temp_env.keys())[-1]
+                comp = temp_env[last_key]() if callable(temp_env[last_key]) else temp_env[last_key]
             else:
                 comp = eval(code, env, env)