diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 88b8023..6e0b5fb 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -30,4 +30,4 @@
"forwardPorts": [
8501
]
-}
\ No newline at end of file
+}
diff --git a/.github/workflows/pylint.yml b/.github/workflows/pylint.yml
new file mode 100644
index 0000000..15c0712
--- /dev/null
+++ b/.github/workflows/pylint.yml
@@ -0,0 +1,71 @@
+name: Pylint
+
+on: [pull_request]
+
+jobs:
+ pylint:
+ runs-on: ubuntu-latest
+ permissions:
+ contents: read
+ pull-requests: write
+ steps:
+ - uses: actions/checkout@v4
+
+ - name: Set up Python
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.10"
+
+ - name: Install dependencies
+ run: |
+ python -m pip install --upgrade pip
+ pip install poetry
+ poetry install
+
+ - name: Run Pylint
+ id: pylint
+ continue-on-error: true
+ run: |
+ # Run pylint on source, app.py, and tests
+ # Redirect output to file, also capture stderr
+ poetry run pylint src/serpent app.py tests > pylint_report.txt 2>&1
+
+ - name: Comment on PR
+ uses: actions/github-script@v7
+ with:
+ script: |
+ const fs = require('fs');
+ try {
+ const report = fs.readFileSync('pylint_report.txt', 'utf8');
+ const MAX_LENGTH = 60000; // GitHub comment limit is ~65536
+ let bodyContent = report;
+
+ if (report.length > MAX_LENGTH) {
+ bodyContent = report.substring(0, MAX_LENGTH) + "\n\n... (Output truncated due to length)";
+ }
+
+ if (!bodyContent.trim()) {
+ bodyContent = "No output captured from Pylint.";
+ }
+
+ // Sanitize bodyContent to avoid breaking the template literal
+ bodyContent = bodyContent.replace(/`/g, '\\`').replace(/\$\{/g, '\\${');
+
+ const outcome = '${{ steps.pylint.outcome }}';
+ const icon = outcome === 'success' ? '✅' : '⚠️';
+ const summary = outcome === 'success' ? 'Pylint passed' : 'Pylint found issues';
+
+ await github.rest.issues.createComment({
+ issue_number: context.issue.number,
+ owner: context.repo.owner,
+ repo: context.repo.repo,
+ body: `### ${icon} ${summary}\n\n\nView Pylint Report
\n\n\`\`\`\n${bodyContent}\n\`\`\`\n `
+ });
+
+ // If we want to fail the workflow if pylint failed:
+ if (outcome === 'failure') {
+ core.setFailed('Pylint found issues.');
+ }
+ } catch (error) {
+ core.setFailed(error.message);
+ }
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index ab7a8d6..91067c6 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -11,10 +11,10 @@ jobs:
runs-on: ubuntu-latest
steps:
- - uses: actions/checkout@v3
+ - uses: actions/checkout@v6
- name: Set up Python
- uses: actions/setup-python@v4
+ uses: actions/setup-python@v6
with:
python-version: "3.10"
diff --git a/.gitignore b/.gitignore
index 1701059..ab685a5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@ __pycache__/
.idea/
SERPENT.egg-info/
build/
+.pytest_cache/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1b339c6..5adac6a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,6 +6,10 @@ repos:
- id: end-of-file-fixer
- id: check-yaml
- id: check-added-large-files
+ - id: check-toml
+ - id: check-json
+ - id: pretty-format-json
+ args: [--autofix]
- repo: https://github.com/psf/black
rev: 25.9.0
diff --git a/.pylintrc b/.pylintrc
deleted file mode 100644
index e908ab2..0000000
--- a/.pylintrc
+++ /dev/null
@@ -1,5 +0,0 @@
-# Source - https://stackoverflow.com/q
-# Posted by Łukasz Rogalski, modified by community. See post 'Timeline' for change history
-# Retrieved 2026-01-10, License - CC BY-SA 4.0
-
-pylint: max-line-length=120
diff --git a/app.py b/app.py
new file mode 100644
index 0000000..8dc7a2d
--- /dev/null
+++ b/app.py
@@ -0,0 +1,263 @@
+"""
+Launch the Streamlit web application for converting Python functions into flowcharts.
+"""
+
+import ast
+import logging
+import re
+from typing import Any
+import shutil
+from pathlib import Path
+
+import streamlit as st
+from PIL import Image
+from streamlit_extras.badges import badge
+from streamlit_extras.stylable_container import stylable_container
+
+from serpent.core import generate_graphviz_flowchart
+from serpent.resources import EXAMPLES, THEMES
+
+assets_dir = Path(__file__).parent / "assets"
+try:
+ LOGO = Image.open(assets_dir / "serpent_logo_transparent.png")
+ ICON = (
+ Image.open(assets_dir / "serpent_logo_compact.png")
+ if (assets_dir / "serpent_logo_compact.png").exists()
+ else LOGO
+ )
+except (FileNotFoundError, OSError):
+ LOGO = None
+ ICON = None
+
+# Page Configuration
+st.set_page_config(
+ page_title="SERPENT",
+ page_icon=ICON or "🐍",
+ layout="wide",
+ initial_sidebar_state="expanded",
+)
+
+# Custom CSS for cleaner look
+st.markdown(
+ """
+
+""",
+ unsafe_allow_html=True,
+)
+
+
+def _render_sidebar() -> tuple[str, str, str]:
+ """Render the sidebar and return settings."""
+ with st.sidebar:
+ if LOGO:
+ st.image(LOGO, width="stretch")
+ else:
+ st.title("SERPENT 🐍")
+
+ st.write("Turn your Python functions into clear flowcharts.")
+ st.divider()
+
+ def update_example():
+ """Callback to update code input when example changes."""
+ ex = st.session_state.example_selector
+ if ex != "(Custom)":
+ st.session_state.code_area_widget = EXAMPLES[ex]
+ st.session_state.code_input = EXAMPLES[ex]
+
+ st.subheader("⚙️ Settings")
+
+ selected_example = st.selectbox(
+ "Load Example",
+ ["(Custom)"] + list(EXAMPLES.keys()),
+ index=0,
+ help="Select an example to see how it works.",
+ key="example_selector",
+ on_change=update_example,
+ )
+
+ st.caption("Appearance")
+ rankdir = st.selectbox(
+ "Orientation",
+ options=["TB", "LR"],
+ format_func=lambda x: "Top-Down" if x == "TB" else "Left-Right",
+ )
+
+ theme_name = st.selectbox("Theme", options=list(THEMES.keys()))
+
+ st.divider()
+
+ with st.expander("About & Help"):
+ st.markdown(
+ """
+ **How to use:**
+ 1. Paste your Python function.
+ 2. The flowchart updates automatically.
+ 3. Download the result.
+
+ **Tips:**
+ - Works best with single functions.
+ - Supports `if/else`, `loops`, `break/continue`.
+ """
+ )
+
+ st.caption("Created by Asif Sayyed")
+ badge(
+ type="github",
+ name="Asifdotexe/SERPENT",
+ url="https://github.com/Asifdotexe/SERPENT",
+ )
+ return selected_example, rankdir, theme_name
+
+
+def _render_input_area(selected_example: str) -> tuple[str, str]:
+ """Render the code input area."""
+ st.subheader("📝 Input Code")
+
+ if selected_example != "(Custom)":
+ default_code = EXAMPLES[selected_example]
+ else:
+ # Keep previous input if possible, else empty
+ default_code = "def my_func():\n pass"
+
+ # Initial state
+ if "code_input" not in st.session_state:
+ st.session_state.code_input = default_code
+
+ def _sync_code_input():
+ """Sync widget state to session state code_input."""
+ st.session_state.code_input = st.session_state.code_area_widget
+
+ # We handle updates via callback now, so we remove the imperative check
+ code = st.text_area(
+ "Python Code",
+ value=st.session_state.code_input,
+ height=400,
+ label_visibility="collapsed",
+ key="code_area_widget",
+ on_change=_sync_code_input,
+ )
+
+ chart_title = st.text_input("Chart Title", placeholder="Enter a title (optional)")
+ return code, chart_title
+
+
+def _render_output_area(
+ code: str, chart_title: str, rankdir: str, theme_name: str
+) -> None:
+ """Render the flowchart output area."""
+ valid_graph = None
+ st.subheader("🖼️ Flowchart")
+
+ with stylable_container(
+ key="output_container",
+ css_styles="""
+ {
+ border: 1px solid rgba(49, 51, 63, 0.2);
+ border-radius: 0.5rem;
+ padding: 1rem;
+ min-height: 480px;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ background-color: rgba(255, 255, 255, 0.05);
+ }
+ """,
+ ):
+ if not code.strip():
+ st.info("Waiting for code input...")
+ else:
+ try:
+ # Validate
+ ast.parse(code)
+
+ # Generate
+ final_title = chart_title or "Flowchart"
+ selected_theme = THEMES[theme_name]
+
+ graph = generate_graphviz_flowchart(
+ code,
+ title=final_title,
+ rankdir=rankdir,
+ style_config=selected_theme,
+ )
+
+ st.graphviz_chart(graph, width="stretch")
+
+ # Store for download buttons below
+ valid_graph = graph
+
+ except SyntaxError as e:
+ st.error(f"Syntax Error: {e}")
+ valid_graph = None
+ except Exception as e:
+ st.error(f"An error occurred: {e}")
+ logging.exception("Graph generation failed")
+ valid_graph = None
+ return valid_graph
+
+
+def _render_download_buttons(valid_graph: Any, chart_title: str) -> None:
+ """Render download buttons if graph is valid."""
+ if not valid_graph:
+ return
+
+ st.divider()
+ d_col1, d_col2 = st.columns([1, 1])
+
+ # Sanitize filename
+ safe_title = re.sub(
+ r"[^a-z0-9_\-]", "", (chart_title or "flowchart").lower().replace(" ", "_")
+ )
+
+ with d_col1:
+ if shutil.which("dot"):
+ try:
+ png_bytes = valid_graph.pipe(format="png")
+ st.download_button(
+ "📥 Download PNG",
+ data=png_bytes,
+ file_name=f"{safe_title}.png",
+ mime="image/png",
+ width="stretch",
+ )
+ except Exception:
+ st.warning("Could not generate PNG (Check Graphviz installation).")
+
+ with d_col2:
+ st.download_button(
+ "📄 Download DOT",
+ data=valid_graph.source,
+ file_name=f"{safe_title}.dot",
+ mime="text/vnd.graphviz",
+ width="stretch",
+ )
+
+
+def main() -> None:
+ """Main application entry point."""
+ selected_example, rankdir, theme_name = _render_sidebar()
+
+ col_header, _ = st.columns([3, 1])
+ with col_header:
+ st.title("Flowchart Generator")
+
+ input_col, output_col = st.columns(2)
+
+ with input_col:
+ code, chart_title = _render_input_area(selected_example)
+
+ with output_col:
+ valid_graph = _render_output_area(code, chart_title, rankdir, theme_name)
+
+ # Download Area (Full width below columns)
+ _render_download_buttons(valid_graph, chart_title)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pyproject.toml b/pyproject.toml
index e7e2af5..bc0a0f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ authors = ["Asif Sayyed "]
readme = "README.md"
license = "MIT"
keywords = ["flowchart", "code visualization", "streamlit", "graphviz"]
-packages = [{include = "serpent"}]
+packages = [{ include = "serpent", from = "src" }]
[tool.poetry.dependencies]
python = "^3.10"
@@ -22,12 +22,25 @@ isort = "^7.0.0"
pylint = "^4.0.4"
pre-commit = "^4.5.1"
-[tool.poetry.scripts]
-serpent = "serpent.app:main"
[tool.pylint.messages_control]
-disable = ["redefined-outer-name", "invalid-name", "assignment-from-no-return", "too-many-branches",
- "unused-argument", "unnecessary-pass", "fixme", "too-few-public-methods"]
+disable = [
+ "redefined-outer-name",
+ "assignment-from-no-return",
+ "too-many-branches",
+ "unused-argument",
+ "unnecessary-pass",
+ "fixme",
+ "too-few-public-methods",
+ "broad-exception-caught",
+]
+
+[tool.pylint.design]
+max-args = 7
+max-positional-arguments = 7
+
+[tool.pylint.format]
+max-line-length = 120
[build-system]
requires = ["poetry-core"]
diff --git a/serpent/app.py b/serpent/app.py
deleted file mode 100644
index ab0da75..0000000
--- a/serpent/app.py
+++ /dev/null
@@ -1,233 +0,0 @@
-"""
-Launch the Streamlit web application for converting Python functions into flowcharts using Graphviz.
-
-Purpose:
-- Clean side-by-side layout with custom styling.
-- Uses modern Streamlit layout features for clarity.
-"""
-
-import re
-import ast
-import shutil
-import logging
-import textwrap
-from pathlib import Path
-
-import streamlit as st
-from PIL import Image
-from streamlit_extras.avatar import avatar
-from streamlit_extras.badges import badge
-from streamlit_extras.stylable_container import stylable_container
-
-from serpent.core import generate_graphviz_flowchart
-
-assets_dir = Path(__file__).parent.parent / "assets"
-try:
- logo = Image.open(assets_dir / "serpent_logo_transparent.png")
- banner = Image.open(assets_dir / "serpent_banner_transparent.png")
-except FileNotFoundError:
- # Fallback to no images if not found
- # (safer for package distribution if assets aren't included yet)
- logo = None
- banner = None
-
-# Page Configuration
-st.set_page_config(
- page_title="SERPENT",
- page_icon=logo,
- layout="wide",
- initial_sidebar_state="collapsed",
-)
-
-# Just hiding the boring default Streamlit menu and footer.
-# Also making buttons full width because big buttons are clickable buttons.
-st.markdown(
- """
-
-""",
- unsafe_allow_html=True,
-)
-
-# A little sample to show user what to do.
-placeholder_code = textwrap.dedent(
- """\
- def should_i_code_today(coffee_level: int, deadline_approaching: bool) -> str:
- if deadline_approaching:
- return "Code like your life depends on it!"
- elif coffee_level > 5:
- return "Let's write some beautiful, elegant code."
- else:
- return "Go get more coffee, then we'll talk."
-"""
-)
-
-
-def main() -> None:
- """
- Launch the Streamlit web application.
-
- This is the main entry point. It paints the UI, takes the user input,
- calls the generator, and serves the hot, fresh flowchart.
- """
- if banner:
- st.image(banner, use_container_width=True)
- else:
- # Fallback title
- st.title("SERPENT")
-
- st.caption(
- "Turn your Python functions into clear, standard flowcharts in a few clicks."
- )
-
- # Using avatar to show who's the boss (author).
- caption_msg = (
- "Hey, turn your code into a flowchart and turn knowledge sharing into a breeze."
- )
- avatar(
- [
- {
- "url": "https://avatars.githubusercontent.com/u/115421661?v=4",
- "size": 40,
- "title": "Asif Sayyed",
- "caption": caption_msg,
- "key": "author_avatar",
- }
- ]
- )
-
- with st.expander("How to use this tool?"):
- st.markdown(
- """
- 1. **Paste your code**: Drop a valid Python function into the text area.
- 2. **Add a title**: Give your flowchart a descriptive title.
- 3. **Generate**: Click the button to see your flowchart appear side-by-side.
- """
- )
- st.code(placeholder_code, language="python")
-
- st.divider()
-
- # Side-by-side layout (make 2 columns)
- input_col, output_col = st.columns(2)
-
- # Input Column
- with (
- input_col,
- stylable_container(
- key="input_container",
- css_styles="""
- {
- border: 1px solid rgba(49, 51, 63, 0.2);
- border-radius: 0.5rem;
- padding: calc(1em - 1px);
- }
- """,
- ),
- ):
- st.subheader("Input")
- code = st.text_area(
- "Your Python Code",
- height=350,
- placeholder="def my_function(): ...",
- help="""Paste a valid Python function here,
- The script will ignore comments and docstrings.""",
- )
- chart_title = st.text_input(
- "Flowchart Title",
- value="Some fancy title for your flowchart?",
- help="Enter a title to be displayed at the top of your flowchart.",
- )
- generate = st.button(
- "Generate Flowchart",
- help="Click here to generate the flowchart from your code.",
- type="primary",
- use_container_width=True,
- )
-
- # Output Column
- with (
- output_col,
- stylable_container(
- key="output_container",
- css_styles="""
- {
- border: 1px solid rgba(49, 51, 63, 0.2);
- border-radius: 0.5rem;
- padding: calc(1em - 1px);
- }
- """,
- ),
- ):
- st.subheader("Output")
-
- if not generate:
- st.info("Your generated flowchart will appear here.")
- elif not code.strip():
- st.warning(
- "Please paste some Python code first. Can't make juice without oranges!"
- )
- else:
- try:
- # Validating input first. Safety first!
- ast.parse(code)
- graph = generate_graphviz_flowchart(code, title=chart_title)
- st.success("✅ Flowchart generated! Looking good.")
- st.graphviz_chart(graph.source)
-
- # Sanitize chart_title for filename
- # Collapse whitespace, remove unsafe chars, lower case, truncate
- safe_title = re.sub(r"[^a-z0-9_\-]", "", chart_title.lower().replace(" ", "_"))
- safe_title = re.sub(r"_+", "_", safe_title).strip("_")
- safe_title = safe_title[:50] or "flowchart"
-
- # Download Buttons
- # If 'dot' is installed, we can give a PNG. If not, only DOT file.
- if shutil.which("dot"):
- png_bytes = graph.pipe(format="png")
- st.download_button(
- label="📥 Download as PNG",
- data=png_bytes,
- file_name=f"{safe_title}.png",
- mime="image/png",
- help="Download the flowchart as a PNG image.",
- use_container_width=True,
- type="primary",
- )
- else:
- st.warning(
- "PNG export not available. Install `Graphviz` system-wide to enable it."
- )
- st.download_button(
- label="⬇️ Download DOT source",
- data=graph.source,
- file_name=f"{safe_title}.dot",
- mime="text/vnd.graphviz",
- help="Download the Graphviz source file (.dot) to render it locally.",
- )
- st.caption("💡 Tip: You can view `.dot` files in VSCode or online.")
-
- except SyntaxError as e:
- st.error(
- f"Syntax Error: Your Python code is invalid.\n\n**Details:** {e}"
- )
- except Exception:
- logging.exception("Unexpected error in main UI loop")
- st.error(
- "An unexpected error occurred. Please try again or contact support."
- )
-
- st.divider()
- st.caption("Like the result? Starring the repository helps a lot!")
- badge(
- type="github",
- name="Asifdotexe/SERPENT",
- url="https://github.com/Asifdotexe/SERPENT",
- )
-
-
-if __name__ == "__main__":
- main()
diff --git a/serpent/core.py b/serpent/core.py
deleted file mode 100644
index da0be6f..0000000
--- a/serpent/core.py
+++ /dev/null
@@ -1,392 +0,0 @@
-"""
-Core logic for parsing Python code with AST and drawing a flowchart.
-
-It leverages Python's Abstract Syntax Tree (AST) to traverse code structure and `graphviz`
-to generate visual flowcharts.
-"""
-
-import ast
-from typing import Any, Optional, Union
-
-from graphviz import Digraph
-
-
-class PythonFlowchartGV(ast.NodeVisitor):
- """
- A custom AST NodeVisitor that generates a Graphviz Digraph
- representing the control flow of Python code.
-
- This class walks through the Python code structure (AST) and converts it into a visual graph.
- Instead of boring old stack, we use a smarter `last_nodes` list to keep track of connections.
- """
-
- def __init__(self) -> None:
- """
- Initialize the Flowchart Visitor.
-
- We need to setup the blank canvas (Digraph) where we will draw our shapes.
- Also need some counters and lists to remember where we came from,
- so we know where to go next.
- """
- self.graph: Digraph = Digraph(format="png")
- self.counter: int = 0
-
- # Why last_nodes?
- # Imagine you are walking in a park. You need to know where you are standing right now
- # to know where you can walk to next. This list keeps track of the ID of the nodes
- # that are waiting to be connected to the next instruction.
- # Initially empty because we haven't started walking yet!
- self.last_nodes: list[Union[str, tuple[str, Optional[str]]]] = []
-
- # Why loop_stack?
- # Loops are tricky, bhai.
- # Sometimes you want to `break` out (exit) or `continue` (go back start).
- # We need to remember which loop we are currently inside so we know where to jump to.
- # It's like inception - loop inside loop inside loop.
- self.loop_stack: list[dict[str, Any]] = []
-
- def new_node(
- self,
- label: str,
- shape: str = "box",
- connect_from: Optional[list[Union[str, tuple[str, Optional[str]]]]] = None,
- edge_label: str = "",
- ) -> str:
- """
- Create a new node in the flowchart and connect it to previous nodes.
-
- This is the main painter! It takes a command, makes a shape, and draws lines (edges)
- from the previous steps to this new step. It handles all the coloring and labeling too.
-
- :param label: The text to display inside the box/diamond/circle.
- :param shape: The geometric shape of the node (box, diamond, oval, etc.).
- :param connect_from: Specific list of nodes to connect FROM.
- If None, uses the last visited nodes.
- :param edge_label: Text to write on the arrow connecting to this new node
- (e.g., "True", "False").
- :return: The unique ID of the newly created node (like "n1", "n2").
- """
- # Global override check
- # Checking if some previous logic left a note saying
- # "Hey, the next edge needs this label!"
- override_label = getattr(self, "next_edge_label", None)
- if override_label:
- edge_label = override_label
- # Used it, now forget it.
- self.next_edge_label = None
-
- # FIXME: Make these enumerations
- # Making it look pretty with pastel colors.
- # Life is too short for boring black and white charts, na?
- color_map = {
- "box": "lightyellow",
- "diamond": "lightblue",
- "oval": "lightgreen",
- "circle": "thistle",
- "parallelogram": "lightcyan",
- }
- fillcolor = color_map.get(shape, "white")
-
- node_id = f"n{self.counter}"
- self.counter += 1
-
- self.graph.node(
- node_id, label=label, shape=shape, style="filled", fillcolor=fillcolor
- )
-
- # Deciding who is the parent of this new node.
- # Default is `self.last_nodes` (the immediate previous steps).
- sources = connect_from if connect_from is not None else self.last_nodes
-
- for source in sources:
- src_id = source
- lbl = edge_label
-
- # If the source is complex (tuple), it might carry its own specific label instruction.
- # Example: An `If` node sends "True" to one guy and "False" to another.
- if isinstance(source, tuple):
- src_id = source[0]
- if source[1]: # If specific label exists, it wins!
- lbl = source[1]
-
- self.graph.edge(src_id, node_id, label=lbl)
-
- # Update current state: This new node is now the "last node".
- self.last_nodes = [node_id]
- return node_id
-
- def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
- """
- Handle function definition.
-
- :param node: The AST node for the function.
- """
- name = node.name
- start_node = self.new_node(f"Function: {name}", shape="oval")
-
- # Function body starts here.
- self.last_nodes = [start_node]
-
- # Go inside the function and visit every line of code.
- for stmt in node.body:
- self.visit(stmt)
-
- # Function over! We don't explicitly add "End" node, flow just stops.
-
- def visit_If(self, node: ast.If) -> None:
- """
- Handle `if` statements.
-
- Life is full of choices, and so is code. We create a diamond shape for the decision.
- Then we split the path: one way for YES (True), one way for NO (False).
- Finally, we have to bring them back together (merge) because the code eventually moves on.
-
- :param node: The AST node for the if statement.
- """
- condition_label = f"If: {ast.unparse(node.test)}"
- decision_node = self.new_node(condition_label, shape="diamond")
-
- decision_node_id = decision_node
-
- # Condition is "True"
- # We prepare to visit the lines inside the `if` block.
- # We cheat a bit and set a global flag
- # so the first node inside gets connected with "True" label.
- self.last_nodes = [decision_node_id]
- true_end_nodes = []
-
- if node.body:
- self.last_nodes = [(decision_node_id, "True")]
- for stmt in node.body:
- self.visit(stmt)
- true_end_nodes = self.last_nodes
- else:
- # Empty body? Just pass through.
- true_end_nodes = [(decision_node_id, "True")]
-
- # Condition is "False"
- false_end_nodes = []
- if node.orelse:
- self.last_nodes = [(decision_node_id, "False")]
- for stmt in node.orelse:
- self.visit(stmt)
- false_end_nodes = self.last_nodes
- else:
- # If there is no else,
- # the False path goes straight from decision to the merge point.
- # We explicitly label this edge as "False".
- false_end_nodes = [(decision_node_id, "False")]
-
- # Decision merge (Joining of the splits)
- # We collect all the endpoints from True path and False path.
- # The next line of code after this big IF will connect from ALL of these endpoints.
- start_merge_nodes = []
- for n in true_end_nodes:
- if isinstance(n, tuple):
- start_merge_nodes.append(n)
- else:
- start_merge_nodes.append((n, None))
-
- for n in false_end_nodes:
- if isinstance(n, tuple):
- start_merge_nodes.append(n)
- else:
- start_merge_nodes.append((n, None))
-
- self.last_nodes = start_merge_nodes
- # Cleanup
- self.next_edge_label = None
-
- def visit_For(self, node: ast.For) -> None:
- """Handle `for` loops using common loop logic.
-
- :param node: The loop node.
- """
- return self._handle_loop(
- node, f"For: {ast.unparse(node.target)} in {ast.unparse(node.iter)}"
- )
-
- def visit_While(self, node: ast.While) -> None:
- """Handle `while` loops using common loop logic.
-
- :param node: The loop node.
- """
- return self._handle_loop(node, f"While: {ast.unparse(node.test)}")
-
- def _handle_loop(self, node: Union[ast.For, ast.While], label: str) -> None:
- """
- Common logic for processing loops.
-
- Loops are circles of life.
- 1. Enter the circle (Condition).
- 2. Do the work (Body).
- 3. Go back to start (Back-edge).
- 4. Or leave if done (Exit).
- Also handles those rebellious `break` and `continue` statements.
-
- :param node: The loop node.
- :param label: Label for the condition node.
- """
- # The Gatekeeper (Condition Node)
- # Using diamond because it is a decision point (True/False).
- condition_node = self.new_node(label, shape="diamond")
-
- # We push a new context to the stack
- # so `break` and `continue` know who their daddy is (current loop).
- self.loop_stack.append(
- {"break": [], "continue": [], "start_node": condition_node}
- )
-
- # Doing the work (Body - True Path)
- self.last_nodes = [(condition_node, "True")]
-
- for stmt in node.body:
- self.visit(stmt)
-
- # The Return Logic (Back-edge)
- # Everyone who reached the end of the body must go back to the condition to check again.
- for n in self.last_nodes:
- src = n[0] if isinstance(n, tuple) else n
- lbl = n[1] if isinstance(n, tuple) else None
- if lbl:
- self.graph.edge(src, condition_node, label=lbl)
- else:
- self.graph.edge(src, condition_node)
-
- # Handling `continue` (Shortcuts)
- # Any `continue` we found inside just jumps straight back to condition.
- loop_ctx = self.loop_stack.pop()
- for cont_node in loop_ctx["continue"]:
- self.graph.edge(cont_node, condition_node)
-
- # The Exit Strategy
- # The next code connects from:
- # - The condition (when it becomes False).
- # - Any `break` statements (they escape the loop).
- exit_nodes = [(condition_node, "False")]
- for break_node in loop_ctx["break"]:
- exit_nodes.append((break_node, None))
-
- self.last_nodes = exit_nodes
- self.next_edge_label = None
-
- def visit_Break(self, _node: ast.Break) -> None:
- """
- Handle `break` statement.
-
- Emergency exit! We cut the current flow and register this node in the
- `break` list of the parent loop. It will be re-connected later to the exit.
- """
- if self.loop_stack:
- break_node = self.new_node("break", shape="box")
- self.loop_stack[-1]["break"].append(break_node)
- # Dead end here, flow transfers to loop exit.
- self.last_nodes = []
- else:
- self.new_node("break (orphaned)", shape="box")
-
- def visit_Continue(self, _node: ast.Continue) -> None:
- """
- Handle `continue` statement.
-
- Skip the rest, go back to start! We register this in `continue` list
- and cut the local flow.
- """
- if self.loop_stack:
- cont_node = self.new_node("continue", shape="box")
- self.loop_stack[-1]["continue"].append(cont_node)
- self.last_nodes = []
- else:
- self.new_node("continue (orphaned)", shape="box")
-
- def visit_Return(self, node: ast.Return) -> None:
- """
- Handle `return` statement.
-
- Game over for this function. We return the value and stop the flow here.
- """
- val = ast.unparse(node.value) if node.value else "None"
- self.new_node(f"Return: {val}", shape="box")
- self.last_nodes = []
-
- def visit_Expr(self, node: ast.Expr) -> None:
- """Handle expression statements (ignoring docstrings).
-
- :param node: The expression node to handle.
- """
- if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
- # Skip docstrings, they are for reading not flowcharting.
- return
- self.new_node(ast.unparse(node).strip(), shape="box")
-
- def visit_Assign(self, node: ast.Assign) -> None:
- """Handle variable assignment.
-
- :param node: The assignment node to handle.
- """
- self.new_node(ast.unparse(node).strip(), shape="box")
-
- def visit_AugAssign(self, node: ast.AugAssign) -> None:
- """Handle augmented assignment (+=, -= etc).
-
- :param node: The augmented assignment node to handle.
- """
- self.new_node(ast.unparse(node).strip(), shape="box")
-
- def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
- """Handle annotated assignment (x: int = 5).
-
- :param node: The annotated assignment node to handle.
- """
- self.new_node(ast.unparse(node).strip(), shape="box")
-
- def visit_Pass(self, node: ast.Pass) -> None:
- """
- Handle `pass`.
-
- Do nothing. Chill. The flow continues as if nothing happened.
-
- :param node: The pass node to handle.
- """
- pass
-
- def generic_visit(self, node: ast.AST) -> None:
- """Fallback for any other node types.
-
- :param node: The node to handle.
- """
- if isinstance(node, ast.stmt):
- self.new_node(f"{type(node).__name__}", shape="box")
- else:
- super().generic_visit(node)
-
-
-def generate_graphviz_flowchart(code_str: str, title: str = "Flowchart") -> Digraph:
- """
- Parse Python source code and generate a Graphviz Digraph.
-
- This is the wrapper function that the outside world calls.
- It catches syntax errors so the app doesn't crash on bad code.
-
- :param code_str: The raw Python code string to parse.
- :param title: The title to display on the flowchart.
- :returns: The generated Graphviz object ready to rendering.
- """
- try:
- tree = ast.parse(code_str)
- except SyntaxError:
- # If the user types garbage, we show a nice error box instead of crashing.
- graph = Digraph()
- graph.node(
- "error",
- label="Syntax Error: Cannot parse code",
- shape="box",
- style="filled",
- fillcolor="lightpink",
- )
- return graph
-
- fc = PythonFlowchartGV()
- fc.visit(tree)
- fc.graph.attr(label=title, labelloc="t", fontsize="20")
- return fc.graph
diff --git a/serpent/__init__.py b/src/serpent/__init__.py
similarity index 100%
rename from serpent/__init__.py
rename to src/serpent/__init__.py
diff --git a/src/serpent/core.py b/src/serpent/core.py
new file mode 100644
index 0000000..5ba9809
--- /dev/null
+++ b/src/serpent/core.py
@@ -0,0 +1,323 @@
+"""
+Core logic for parsing Python code with AST and drawing a flowchart.
+"""
+
+import ast
+from typing import Any, Optional, Union
+
+from graphviz import Digraph
+
+from serpent.resources import THEMES
+
+
+class PythonFlowchartGV(ast.NodeVisitor):
+ """
+ A custom AST NodeVisitor that generates a Graphviz Digraph
+ representing the control flow of Python code.
+ """
+
+ def __init__(
+ self, rankdir: str = "TB", style_config: Optional[dict[str, str]] = None
+ ) -> None:
+ """
+ Initialize the Flowchart Visitor.
+ """
+ self.graph: Digraph = Digraph(format="png")
+ self.graph.attr(rankdir=rankdir) # Set orientation (TB or LR)
+ self.counter: int = 0
+ self.last_nodes: list[Union[str, tuple[str, Optional[str]]]] = []
+ self.loop_stack: list[dict[str, Any]] = []
+ self.next_edge_label: Optional[str] = None
+
+ # Default colors if not provided
+ self.style_config = style_config or THEMES["Classic (Pastel)"]
+
+ def new_node(
+ self,
+ label: str,
+ shape: str = "box",
+ connect_from: Optional[list[Union[str, tuple[str, Optional[str]]]]] = None,
+ edge_label: str = "",
+ node_type: Optional[str] = None,
+ ) -> str:
+ """
+ Create a new node in the flowchart and connect it to previous nodes.
+ """
+ override_label = getattr(self, "next_edge_label", None)
+ if override_label:
+ edge_label = override_label
+ self.next_edge_label = None
+
+ fillcolor = self.style_config.get(
+ node_type, self.style_config.get(shape, "white")
+ )
+
+ node_id = f"n{self.counter}"
+ self.counter += 1
+
+ self.graph.node(
+ node_id, label=label, shape=shape, style="filled", fillcolor=fillcolor
+ )
+
+ sources = connect_from if connect_from is not None else self.last_nodes
+
+ for source in sources:
+ src_id = source
+ lbl = edge_label
+
+ if isinstance(source, tuple):
+ src_id = source[0]
+ if source[1]:
+ lbl = source[1]
+
+ self.graph.edge(src_id, node_id, label=lbl)
+
+ self.last_nodes = [node_id]
+ return node_id
+
+ def visit(self, node: ast.AST) -> Any:
+ """Override visit to dispatch to snake_case methods."""
+ name = node.__class__.__name__
+ snake_name = "".join(
+ ["_" + c.lower() if c.isupper() else c for c in name]
+ ).lstrip("_")
+
+ method_name = "visit_" + snake_name
+ visitor = getattr(self, method_name, self.generic_visit)
+ return visitor(node)
+
+ def visit_function_def(self, node: ast.FunctionDef) -> None:
+ """Handle function definition."""
+ name = node.name
+ start_node = self.new_node(f"Function: {name}", shape="oval")
+ self.last_nodes = [start_node]
+ for stmt in node.body:
+ self.visit(stmt)
+
+ def visit_if(self, node: ast.If) -> None:
+ """Handle `if` statements."""
+ condition_label = f"If: {ast.unparse(node.test)}"
+ decision_node = self.new_node(condition_label, shape="diamond")
+ decision_node_id = decision_node
+
+ self.last_nodes = [decision_node_id]
+ true_end_nodes = []
+
+ if node.body:
+ self.last_nodes = [(decision_node_id, "True")]
+ for stmt in node.body:
+ self.visit(stmt)
+ true_end_nodes = self.last_nodes
+ else:
+ true_end_nodes = [(decision_node_id, "True")]
+
+ false_end_nodes = []
+ if node.orelse:
+ self.last_nodes = [(decision_node_id, "False")]
+ for stmt in node.orelse:
+ self.visit(stmt)
+ false_end_nodes = self.last_nodes
+ else:
+ false_end_nodes = [(decision_node_id, "False")]
+
+ start_merge_nodes = []
+ for n in true_end_nodes:
+ if isinstance(n, tuple):
+ start_merge_nodes.append(n)
+ else:
+ start_merge_nodes.append((n, None))
+
+ for n in false_end_nodes:
+ if isinstance(n, tuple):
+ start_merge_nodes.append(n)
+ else:
+ start_merge_nodes.append((n, None))
+
+ self.last_nodes = start_merge_nodes
+ self.next_edge_label = None
+
+ def visit_for(self, node: ast.For) -> None:
+ """Handle `for` loops."""
+ return self._handle_loop(
+ node, f"For: {ast.unparse(node.target)} in {ast.unparse(node.iter)}"
+ )
+
+ def visit_while(self, node: ast.While) -> None:
+ """Handle `while` loops."""
+ return self._handle_loop(node, f"While: {ast.unparse(node.test)}")
+
+ def _handle_loop(self, node: Union[ast.For, ast.While], label: str) -> None:
+ """Common logic for processing loops."""
+ condition_node = self.new_node(label, shape="diamond")
+
+ self.loop_stack.append(
+ {"break": [], "continue": [], "start_node": condition_node}
+ )
+
+ self.last_nodes = [(condition_node, "True")]
+
+ for stmt in node.body:
+ self.visit(stmt)
+
+ for n in self.last_nodes:
+ src = n[0] if isinstance(n, tuple) else n
+ lbl = n[1] if isinstance(n, tuple) else None
+ if lbl:
+ self.graph.edge(src, condition_node, label=lbl)
+ else:
+ self.graph.edge(src, condition_node)
+
+ loop_ctx = self.loop_stack.pop()
+ for cont_node in loop_ctx["continue"]:
+ self.graph.edge(cont_node, condition_node)
+
+ exit_nodes = [(condition_node, "False")]
+ for break_node in loop_ctx["break"]:
+ exit_nodes.append((break_node, None))
+
+ self.last_nodes = exit_nodes
+ self.next_edge_label = None
+
+ def visit_break(self, _node: ast.Break) -> None:
+ """Handle `break` statement."""
+
+ if self.loop_stack:
+ break_node = self.new_node("break", shape="box", node_type="break")
+ self.loop_stack[-1]["break"].append(break_node)
+ self.last_nodes = []
+ else:
+ self.new_node("break (orphaned)", shape="box", node_type="break")
+
+ def visit_continue(self, _node: ast.Continue) -> None:
+ """Handle `continue` statement."""
+ if self.loop_stack:
+ cont_node = self.new_node("continue", shape="box", node_type="continue")
+ self.loop_stack[-1]["continue"].append(cont_node)
+ self.last_nodes = []
+ else:
+ self.new_node("continue (orphaned)", shape="box", node_type="continue")
+
+ def visit_return(self, node: ast.Return) -> None:
+ """Handle `return` statement."""
+ val = ast.unparse(node.value) if node.value else "None"
+ self.new_node(f"Return: {val}", shape="box")
+ self.last_nodes = []
+
+ def visit_expr(self, node: ast.Expr) -> None:
+ """Handle expression statements."""
+ if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
+ return
+ self.new_node(ast.unparse(node).strip(), shape="box")
+
+ def visit_assign(self, node: ast.Assign) -> None:
+ """Handle variable assignment."""
+ self.new_node(ast.unparse(node).strip(), shape="box")
+
+ def visit_aug_assign(self, node: ast.AugAssign) -> None:
+ """Handle augmented assignment."""
+ self.new_node(ast.unparse(node).strip(), shape="box")
+
+ def visit_ann_assign(self, node: ast.AnnAssign) -> None:
+ """Handle annotated assignment."""
+ self.new_node(ast.unparse(node).strip(), shape="box")
+
+ def visit_try(self, node: ast.Try) -> None:
+ """
+ Handle try...except...finally blocks.
+
+ Try is a risky business.
+ 1. We attempt the risky code (Try body).
+ 2. If it blows up, we catch it (Except handlers).
+ 3. If it works, we might do something else (Else).
+ 4. No matter what, we clean up (Finally).
+ """
+ try_node = self.new_node("Try", shape="diamond")
+
+ # 1. Try Body
+ self.last_nodes = [(try_node, "Attempt")]
+ for stmt in node.body:
+ self.visit(stmt)
+
+ success_nodes = self.last_nodes
+ all_end_nodes = []
+
+ # 2. Exception Handlers
+ for handler in node.handlers:
+ exc_label = "Exception"
+ if handler.type:
+ exc_label = f"Exc: {ast.unparse(handler.type)}"
+
+ # Conceptually, exceptions branch off the "Try" attempt
+ self.last_nodes = [(try_node, exc_label)]
+
+ for stmt in handler.body:
+ self.visit(stmt)
+
+ all_end_nodes.extend(self.last_nodes)
+
+ # 3. Else Block (executed if no exception)
+ if node.orelse:
+ self.last_nodes = success_nodes
+ for stmt in node.orelse:
+ self.visit(stmt)
+ all_end_nodes.extend(self.last_nodes)
+ else:
+ all_end_nodes.extend(success_nodes)
+
+ if node.finalbody:
+ self.last_nodes = all_end_nodes
+ for stmt in node.finalbody:
+ self.visit(stmt)
+ # Flow continues from end of finally
+ else:
+ self.last_nodes = all_end_nodes
+
+ self.next_edge_label = None
+
+ def visit_raise(self, node: ast.Raise) -> None:
+ """Handle `raise` statement."""
+ if node.exc:
+ val = ast.unparse(node.exc)
+ self.new_node(f"Raise: {val}", shape="box")
+ else:
+ self.new_node("Raise", shape="box")
+ self.last_nodes = []
+
+ def visit_pass(self, node: ast.Pass) -> None:
+ """Handle `pass`."""
+ pass
+
+ def generic_visit(self, node: ast.AST) -> None:
+ """Fallback for any other node types."""
+ if isinstance(node, ast.stmt):
+ self.new_node(f"{type(node).__name__}", shape="box")
+ else:
+ super().generic_visit(node)
+
+
+def generate_graphviz_flowchart(
+ code_str: str,
+ title: str = "Flowchart",
+ rankdir: str = "TB",
+ style_config: Optional[dict[str, str]] = None,
+) -> Digraph:
+ """
+ Parse Python source code and generate a Graphviz Digraph.
+ """
+ try:
+ tree = ast.parse(code_str)
+ except SyntaxError:
+ graph = Digraph()
+ graph.node(
+ "error",
+ label="Syntax Error: Cannot parse code",
+ shape="box",
+ style="filled",
+ fillcolor="lightpink",
+ )
+ return graph
+
+ fc = PythonFlowchartGV(rankdir=rankdir, style_config=style_config)
+ fc.visit(tree)
+ fc.graph.attr(label=title, labelloc="t", fontsize="20")
+ return fc.graph
diff --git a/src/serpent/resources.py b/src/serpent/resources.py
new file mode 100644
index 0000000..19af29c
--- /dev/null
+++ b/src/serpent/resources.py
@@ -0,0 +1,154 @@
+"""
+Resources for the SERPENT application.
+"""
+
+import textwrap
+from typing import Dict
+
+# --- Themes ---
+THEMES: Dict[str, Dict[str, str]] = {
+ "Classic (Pastel)": {
+ "box": "lightyellow",
+ "diamond": "lightblue",
+ "oval": "lightgreen",
+ "circle": "thistle",
+ "parallelogram": "lightcyan",
+ "break": "mistyrose",
+ "continue": "lightgray",
+ },
+ "Clean White": {
+ "box": "white",
+ "diamond": "white",
+ "oval": "white",
+ "circle": "white",
+ "parallelogram": "white",
+ "break": "white",
+ "continue": "white",
+ },
+ "Dark Mode": {
+ "box": "#444444",
+ "diamond": "#555555",
+ "oval": "#222222",
+ "circle": "#666666",
+ "parallelogram": "#333333",
+ "break": "#883333",
+ "continue": "#333388",
+ },
+ "Blueberry": {
+ "box": "#e3f2fd",
+ "diamond": "#bbdefb",
+ "oval": "#90caf9",
+ "circle": "#64b5f6",
+ "parallelogram": "#42a5f5",
+ "break": "#ffcdd2",
+ "continue": "#e1bee7",
+ },
+}
+
+# --- Examples ---
+EXAMPLES = {
+ "ATM Machine (If/Else)": textwrap.dedent(
+ """\
+ def atm_withdrawal(balance: float, request: float, is_authenticated: bool) -> str:
+ if not is_authenticated:
+ return "Authentication failed."
+
+ if request <= 0:
+ print("Invalid amount.")
+ result = "Error: Amount must be positive."
+ elif request > balance:
+ print("Insufficient funds.")
+ result = "Error: Not enough money."
+ else:
+ balance -= request
+ print(f"Dispensing ${request}...")
+ result = f"Success. New balance: ${balance}"
+
+ return result
+ """
+ ),
+ "Smart Light (Loop & Condition)": textwrap.dedent(
+ """\
+ def smart_lighting_system(sensor_readings: list[float], threshold: float):
+ for reading in sensor_readings:
+ if reading < 0:
+ print("Sensor Error: Negative light level.")
+ continue # Skip invalid reading
+
+ if reading > threshold:
+ print(f"Bright ({reading} lux): Turning lights OFF.")
+ break # Sufficient light found, stop checking
+ else:
+ print(f"Dim ({reading} lux): Keep lights ON.")
+
+ print("Lighting check complete.")
+ """
+ ),
+ "Server Connection (While Loop)": textwrap.dedent(
+ """\
+ def connect_to_server(max_retries: int):
+ attempt = 0
+ connected = False
+
+ while attempt < max_retries and not connected:
+ print(f"Connecting... Attempt {attempt + 1}")
+ # Simulate connection logic
+ if attempt == 2: # Pretend success on 3rd try
+ connected = True
+ else:
+ attempt += 1
+
+ if connected:
+ return "Connection Established"
+ else:
+ return "Connection Failed Service Unavailable"
+ """
+ ),
+ "File Safer (Try/Except/Finally)": textwrap.dedent(
+ """\
+ def safe_file_reader(filepath: str) -> str:
+ file_handle = None
+ try:
+ print(f"Opening {filepath}...")
+ # Simulate opening file (would naturally raise OSError)
+ if not filepath:
+ raise ValueError("Empty filepath")
+ file_handle = open(filepath, 'r')
+ data = file_handle.read()
+ return data
+ except FileNotFoundError:
+ return "Error: File not found."
+ except ValueError as e:
+ return f"Error: Invalid input - {e}"
+ finally:
+ if file_handle:
+ print("Closing file handle...")
+ file_handle.close()
+ print("Cleanup complete.")
+ """
+ ),
+ "Order Processing (Nested)": textwrap.dedent(
+ """\
+ def process_orders(orders: list[dict]):
+ for order in orders:
+ status = order.get("status")
+
+ if status == "cancelled":
+ continue
+
+ if status == "pending":
+ amount = order.get("amount", 0)
+ if amount > 1000:
+ print("Flagging for manual review (High Value)")
+ elif amount < 0:
+ print("Error: Invalid Order")
+ break # Stop critical error
+ else:
+ print("Auto-approving order")
+ else:
+ print(f"Skipping order with status: {status}")
+
+ return "Batch Complete"
+ """
+ ),
+}
diff --git a/tests/test_app_smoke.py b/tests/test_app_smoke.py
index bf56ad9..09d6203 100644
--- a/tests/test_app_smoke.py
+++ b/tests/test_app_smoke.py
@@ -1,15 +1,26 @@
"""
Smoke tests for the Streamlit application.
"""
+
+from pathlib import Path
+
from streamlit.testing.v1 import AppTest
-from serpent import app
+
+
+# Add project root to path to import app.py
+root_dir = Path(__file__).parent.parent
+
def test_app_startup():
"""
Smoke test to verify the app starts up without errors.
"""
- at = AppTest.from_file(app.__file__).run()
-
- # Check if the caption is correct (title is conditional on banner presence)
- assert at.caption[0].value == "Turn your Python functions into clear, standard flowcharts in a few clicks."
+ at = AppTest.from_file(str(root_dir / "app.py")).run()
+
+ # Check if the app runs without exception
assert not at.exception
+
+ # Check if the title is correct
+ # The title might be an image or text depending on assets, but "SERPENT" should be somewhere
+ # logic in app.py: st.title("SERPENT 🐍") if no logo
+ pass
diff --git a/tests/test_break_continue_colors.py b/tests/test_break_continue_colors.py
new file mode 100644
index 0000000..8ed0756
--- /dev/null
+++ b/tests/test_break_continue_colors.py
@@ -0,0 +1,47 @@
+"""
+Test file for break and continue colors.
+"""
+
+import textwrap
+from serpent.core import generate_graphviz_flowchart
+from serpent.resources import THEMES
+
+
+def test_break_continue_colors():
+ """
+ Verify that break and continue nodes get their specific colors
+ from the theme configuration.
+ """
+ code = textwrap.dedent(
+ """
+ while True:
+ if condition:
+ break
+ else:
+ continue
+ """
+ )
+
+ # Use Classic theme which has specific colors for break/continue
+ theme = THEMES["Classic (Pastel)"]
+ break_color = theme["break"] # mistyrose
+ continue_color = theme["continue"] # lightgray
+
+ graph = generate_graphviz_flowchart(code, style_config=theme)
+ dot_source = graph.source
+
+ print(f"Checking for break color: {break_color}")
+ assert (
+ f'fillcolor="{break_color}"' in dot_source
+ or f"fillcolor={break_color}" in dot_source
+ ), f"Break node missing color {break_color} in dot_source"
+
+ print(f"Checking for continue color: {continue_color}")
+ assert (
+ f'fillcolor="{continue_color}"' in dot_source
+ or f"fillcolor={continue_color}" in dot_source
+ ), f"Continue node missing color {continue_color} in dot_source"
+
+
+if __name__ == "__main__":
+ test_break_continue_colors()
diff --git a/tests/test_try_except.py b/tests/test_try_except.py
new file mode 100644
index 0000000..d213a32
--- /dev/null
+++ b/tests/test_try_except.py
@@ -0,0 +1,32 @@
+"""
+Test file for try/except blocks.
+"""
+
+from serpent.core import generate_graphviz_flowchart
+
+
+def test_try_except_structure():
+ """Test that try/except blocks are correctly structured."""
+ code = """
+try:
+ process()
+except ValueError:
+ handle_error()
+else:
+ success()
+finally:
+ cleanup()
+"""
+ graph = generate_graphviz_flowchart(code)
+ source = graph.source
+
+ # Check for nodes
+ assert "label=Try" in source
+ assert "label=Attempt" in source
+ assert "Exc: ValueError" in source
+
+ # Check content
+ assert "process" in source
+ assert "handle_error" in source
+ assert "success" in source
+ assert "cleanup" in source