ToTeMM/code-complexity-visualizer

1import os
2import re
3import streamlit as st
4from radon.visitors import ComplexityVisitor
5from radon.metrics import mi_visit
6import google.generativeai as genai
7import ast
8
9# 1) Page configuration
10st.set_page_config(page_title="Code Complexity Visualizer", page_icon="🔬", layout="centered")
11
12# 2) API Key Management
13GEMINI_API_KEY = None
14try:
15	# Prefer Streamlit secrets when available
16	if "GEMINI_API_KEY" in st.secrets:
17		GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
18except Exception:
19	# Secrets not available (e.g., local execution)
20	pass
21
22# Fallback to environment variable if not found in secrets
23if not GEMINI_API_KEY:
24	GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
25
26if GEMINI_API_KEY:
27	genai.configure(api_key=GEMINI_API_KEY)
28else:
29	st.warning("No Gemini API key found. Add it to Streamlit secrets or set environment variable 'GEMINI_API_KEY'. AI refactoring will be disabled.")
30
31# 3) Core Analysis Function
32
33def analyze_code_complexity(code: str):
34	"""Return (total_cyclomatic_complexity, maintainability_index).
35	If radon fails to parse, return (None, None).
36	"""
37	try:
38		visitor = ComplexityVisitor.from_code(code)
39		total_complexity = sum(block.complexity for block in visitor.functions + visitor.methods)
40		# Maintainability Index (multi=True returns per block MI and aggregated values)
41		mi_scores = mi_visit(code, multi=True)
42		# mi_visit returns list of tuples when multi=True; we use global score if available
43		# Fallback: compute average of returned MI values if list of numbers
44		maintainability_index = None
45		if isinstance(mi_scores, (list, tuple)) and len(mi_scores) > 0:
46			# radon returns list of tuples (filename, mi) or list of MI values depending on version
47			# Normalize to a float MI: try to extract numeric values then average
48			values = []
49			for item in mi_scores:
50				if isinstance(item, (int, float)):
51					values.append(float(item))
52				elif isinstance(item, (list, tuple)) and len(item) >= 2 and isinstance(item[1], (int, float)):
53					values.append(float(item[1]))
54			if values:
55				maintainability_index = sum(values) / len(values)
56		if maintainability_index is None:
57			# As a last resort, compute single value with multi=False
58			maintainability_index = float(mi_visit(code, multi=False))
59		return total_complexity, maintainability_index
60	except Exception:
61		return None, None
62
63# Helper: extract per-function slices using AST
64
65def extract_functions(code: str):
66	"""Return list of dicts: {name, start, end, code} for top-level and nested functions."""
67	results = []
68	try:
69		tree = ast.parse(code)
70		lines = code.splitlines()
71		for node in ast.walk(tree):
72			if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
73				name = node.name
74				start = getattr(node, 'lineno', None)
75				end = getattr(node, 'end_lineno', None)
76				if start is not None and end is not None and 1 <= start <= end <= len(lines):
77					snippet = "\n".join(lines[start-1:end])
78					results.append({
79						"name": name,
80						"start": start,
81						"end": end,
82						"code": snippet,
83					})
84	except Exception:
85		return []
86	return results
87
88# 4) Core LLM Refactoring Function
89
90def _is_api_configured() -> bool:
91	return bool(GEMINI_API_KEY)
92
93def _strip_markdown_fences(text: str) -> str:
94	# Remove triple backtick code fences, optionally with language
95	text = re.sub(r"^\s*```[a-zA-Z0-9_+-]*\s*\n", "", text)
96	text = re.sub(r"\n\s*```\s*$", "", text)
97	# Also remove stray backticks
98	text = text.replace("```python", "").replace("```", "").strip()
99	return text
100
101def get_llm_refactoring_suggestion(code: str, complexity_score: float) -> str:
102	"""Return refactored code suggestion from Gemini, or error text if unavailable."""
103	if not _is_api_configured():
104		return "[Error] Gemini API key not configured."
105	try:
106		model = genai.GenerativeModel('gemini-1.5-flash')
107		prompt = (
108			"You are an expert Python developer. Your goal is to reduce cyclomatic complexity "
109			"and improve maintainability while preserving behavior and I/O.\n\n"
110			f"Given this Python code (cyclomatic complexity score: {complexity_score}), refactor it.\n"
111			"Strict requirements:\n"
112			"- Preserve functionality and public API.\n"
113			"- Reduce branching and deeply nested logic; split into small functions where helpful.\n"
114			"- Prefer clear names and straight-line logic; avoid cleverness.\n"
115			"- Keep imports minimal and standard.\n"
116			"- Only output the complete refactored Python code in a SINGLE code block.\n"
117			"- Do NOT include any explanation before or after the code.\n\n"
118			f"Original code:\n<CODE>\n{code}\n</CODE>\n\n"
119			"Output: a single Python code block with the refactoring."
120		)
121		response = model.generate_content(prompt)
122		text = getattr(response, 'text', '') or ''
123		if not text:
124			# Some SDK versions nest parts differently
125			candidates = getattr(response, 'candidates', None)
126			if candidates:
127				text = candidates[0].content.parts[0].text if candidates[0].content.parts else ''
128		suggestion = _strip_markdown_fences(text)
129		return suggestion if suggestion.strip() else "[Error] Empty response from model."
130	except Exception as e:
131		return f"[Error] Failed to generate suggestion: {e}"
132
133# 5) Streamlit User Interface
134st.title("🔬 Code Complexity Visualizer")
135st.markdown(
136	"Analyze Python code for Cyclomatic Complexity and Maintainability Index. "
137	"If code is overly complex, an AI refactoring suggestion will be generated using Google Gemini."
138)
139
140user_code = st.text_area(
141	"Paste your Python function or module here:",
142	height=280,
143	placeholder="""def example(a, b):
144	if a > 0:
145		if b > 0:
146			return a + b
147		else:
148			return a - b
149	else:
150		return b - a""",
151)
152
153analyze_clicked = st.button("Analyze Code")
154
155if analyze_clicked:
156	if not user_code or not user_code.strip():
157		st.error("Please paste some Python code before analyzing.")
158	else:
159		total_complexity, mi_score = analyze_code_complexity(user_code)
160		if total_complexity is None or mi_score is None:
161			st.error("Failed to analyze code. Ensure the code is valid Python.")
162		else:
163			st.header("Analysis Results")
164			col1, col2 = st.columns(2)
165
166			with col1:
167				st.metric(label="Cyclomatic Complexity (Σ of functions/methods)", value=int(total_complexity))
168				if total_complexity < 6:
169					st.success("Lower is better. 1-5 Good")
170				elif 6 <= total_complexity <= 10:
171					st.warning("Moderate complexity. 6-10 Moderate")
172				else:
173					st.error("High complexity. 11+ High")
174
175			with col2:
176				st.metric(label="Maintainability Index (0-100)", value=round(float(mi_score), 2))
177				if mi_score > 19:
178					st.success("Higher is better. 20-100 High")
179				elif 10 <= mi_score <= 19:
180					st.warning("Medium maintainability. 10-19 Medium")
181				else:
182					st.error("Low maintainability. 0-9 Low")
183
184			# Per-function refactor workflow
185			functions = extract_functions(user_code)
186			if functions:
187				st.subheader("Per-function Refactor")
188				# Compute per-function metrics
189				items = []
190				for fn in functions:
191					cc, mi = analyze_code_complexity(fn["code"])
192					if cc is not None and mi is not None:
193						items.append({"label": f"{fn['name']} (L{fn['start']}-{fn['end']}) – CC {int(cc)}, MI {round(float(mi),2)}", "fn": fn, "cc": cc, "mi": mi})
194				if items:
195					labels = [it["label"] for it in items]
196					choice = st.selectbox("Select a function to refactor", labels, index=0)
197					selected = next((it for it in items if it["label"] == choice), None)
198					can_refactor = _is_api_configured()
199					refactor_clicked = st.button("Refactor selected function", disabled=not can_refactor)
200					if not can_refactor:
201						st.info("Provide GEMINI_API_KEY to enable AI refactoring.")
202					if refactor_clicked and selected:
203						with st.spinner("Generating function refactor..."):
204							suggestion = get_llm_refactoring_suggestion(selected["fn"]["code"], selected["cc"])
205						# Compute before/after metrics
206						before_cc, before_mi = selected["cc"], selected["mi"]
207						after_cc, after_mi = analyze_code_complexity(suggestion)
208						col_a, col_b = st.columns(2)
209						with col_a:
210							st.subheader("Original Function")
211							st.code(selected["fn"]["code"], language="python")
212							st.metric("CC (before)", int(before_cc))
213							st.metric("MI (before)", round(float(before_mi), 2))
214						with col_b:
215							st.subheader("Suggested Refactoring")
216							st.code(suggestion, language="python")
217							if after_cc is not None and after_mi is not None:
218								st.metric("CC (after)", int(after_cc))
219								st.metric("MI (after)", round(float(after_mi), 2))
220								if int(after_cc) < int(before_cc) and float(after_mi) >= float(before_mi):
221									st.success("Improved complexity and maintainability.")
222								elif int(after_cc) < int(before_cc):
223									st.success("Improved complexity.")
224								elif float(after_mi) > float(before_mi):
225									st.success("Improved maintainability.")
226								else:
227									st.warning("No improvement detected.")
228				else:
229					st.warning("No functions detected or metrics unavailable for selection.")
230
231			# AI Refactoring for high complexity (whole snippet)
232			if total_complexity > 10:
233				st.header("🤖 AI Refactoring Suggestion")
234				with st.spinner("Generating suggestion..."):
235					suggestion = get_llm_refactoring_suggestion(user_code, total_complexity)
236				col_a, col_b = st.columns(2)
237				with col_a:
238					st.subheader("Original Code")
239					st.code(user_code, language="python")
240				with col_b:
241					st.subheader("Suggested Refactoring")
242					st.code(suggestion, language="python")
243