python_to_mermaid/convert/
python_to_flowchart.rs

1use crate::flowchart::{self, Flowchart, FlowchartItem};
2use anyhow::Result;
3use ast::{StmtFunctionDef, Visitor as _};
4use itertools::Itertools as _;
5use rustpython_parser::{ast, Parse};
6
7pub fn enumerate_fn_defs(source: &str) -> Result<Vec<FnDef>> {
8    let ast = ast::ModModule::parse(source, ".")?;
9    let mut collector = FnDefCollector::new();
10    for stmt in ast.body {
11        collector.visit_stmt(stmt);
12    }
13
14    Ok(collector
15        .fn_defs
16        .into_iter()
17        .map(|(name, fn_def)| FnDef { name, fn_def })
18        .collect())
19}
20
21pub fn convert(fn_def: FnDef) -> Result<Flowchart> {
22    let mut fc = Flowchart::new_root(fn_def.name);
23    let mut gen = FlowchartGenerator::new(&mut fc);
24    gen.visit_stmts(fn_def.fn_def.body);
25    Ok(fc)
26}
27
28#[derive(Debug, Clone)]
29pub struct FnDef {
30    pub name: String,
31    pub fn_def: ast::StmtFunctionDef,
32}
33
34trait VisitorExt: ast::Visitor {
35    fn visit_stmts<I: IntoIterator<Item = ast::Stmt>>(&mut self, stmts: I) {
36        for stmt in stmts {
37            self.visit_stmt(stmt);
38        }
39    }
40}
41
42impl<V: ast::Visitor> VisitorExt for V {}
43
44struct FnDefCollector {
45    path: Vec<String>,
46    fn_defs: Vec<(String, ast::StmtFunctionDef)>,
47}
48
49impl FnDefCollector {
50    pub fn new() -> Self {
51        Self {
52            path: Vec::new(),
53            fn_defs: Vec::new(),
54        }
55    }
56}
57
58impl ast::Visitor for FnDefCollector {
59    fn visit_stmt_class_def(&mut self, node: ast::StmtClassDef) {
60        self.path.push(node.name.to_string());
61        self.generic_visit_stmt_class_def(node);
62        self.path.pop();
63    }
64
65    fn visit_stmt_function_def(&mut self, node: StmtFunctionDef) {
66        self.path.push(node.name.to_string());
67        self.fn_defs.push((self.path.join("."), node.clone()));
68        self.generic_visit_stmt_function_def(node);
69        self.path.pop();
70    }
71}
72
73struct FlowchartGenerator<'f> {
74    fc: &'f mut Flowchart,
75}
76
77impl<'f> FlowchartGenerator<'f> {
78    pub fn new(fc: &'f mut Flowchart) -> Self {
79        Self { fc }
80    }
81}
82
83impl<'f> ast::Visitor for FlowchartGenerator<'f> {
84    fn visit_stmt_return(&mut self, node: ast::StmtReturn) {
85        let label = format!(
86            "return: {}",
87            node.value.map_or("".to_string(), |v| v.to_string())
88        );
89
90        self.fc
91            .items
92            .push(FlowchartItem::Terminal(flowchart::Terminal::new(label)));
93    }
94
95    fn visit_stmt_delete(&mut self, node: ast::StmtDelete) {
96        let label = format!("delete: {}", node.targets.iter().format(", "));
97
98        self.fc
99            .items
100            .push(FlowchartItem::Step(flowchart::Step::new(label)));
101    }
102
103    fn visit_stmt_assign(&mut self, node: ast::StmtAssign) {
104        let label = format!("{} = {}", node.targets.iter().format(", "), node.value);
105
106        self.fc
107            .items
108            .push(FlowchartItem::Step(flowchart::Step::new(label)));
109    }
110
111    fn visit_stmt_aug_assign(&mut self, node: ast::StmtAugAssign) {
112        let op = op_to_string(&node.op);
113        let label = format!("{} {}= {}", node.target, op, node.value);
114
115        self.fc
116            .items
117            .push(FlowchartItem::Step(flowchart::Step::new(label)));
118    }
119
120    fn visit_stmt_ann_assign(&mut self, node: ast::StmtAnnAssign) {
121        let label = format!(
122            "{}: {}{}",
123            node.target,
124            node.annotation,
125            node.value.map_or_else(String::new, |v| format!(" = {}", v))
126        );
127
128        self.fc
129            .items
130            .push(FlowchartItem::Step(flowchart::Step::new(label)));
131    }
132
133    fn visit_stmt_for(&mut self, node: ast::StmtFor) {
134        let label = format!("for {} in {}", node.target, node.iter);
135        let mut fc = Flowchart::new_sub(label);
136        let mut gen = FlowchartGenerator::new(&mut fc);
137        gen.visit_stmts(node.body);
138        self.fc.items.push(FlowchartItem::SubFlowchart(fc));
139    }
140
141    fn visit_stmt_while(&mut self, node: ast::StmtWhile) {
142        let label = format!("while {}", node.test);
143        let mut fc = Flowchart::new_sub(label);
144        let mut gen = FlowchartGenerator::new(&mut fc);
145        gen.visit_stmts(node.body);
146        self.fc.items.push(FlowchartItem::SubFlowchart(fc));
147    }
148
149    fn visit_stmt_if(&mut self, node: ast::StmtIf) {
150        let then_items = {
151            let mut fc = Flowchart::new_unlabeled_sub();
152            let mut gen = FlowchartGenerator::new(&mut fc);
153            gen.visit_stmts(node.body);
154            fc.items
155        };
156
157        let else_items = {
158            let mut fc = Flowchart::new_unlabeled_sub();
159            let mut gen = FlowchartGenerator::new(&mut fc);
160            gen.visit_stmts(node.orelse);
161            fc.items
162        };
163
164        self.fc
165            .items
166            .push(FlowchartItem::Condition(flowchart::Condition {
167                condition: format!("if: {} ?", node.test),
168                then_items,
169                else_items,
170            }));
171    }
172
173    fn visit_stmt_with(&mut self, node: ast::StmtWith) {
174        let label = format!(
175            "with {}",
176            node.items
177                .iter()
178                .map(|item| &item.context_expr)
179                .format(", ")
180        );
181        let mut fc = Flowchart::new_sub(label);
182        let mut gen = FlowchartGenerator::new(&mut fc);
183        gen.visit_stmts(node.body);
184        self.fc.items.push(FlowchartItem::SubFlowchart(fc));
185    }
186
187    fn visit_stmt_match(&mut self, node: ast::StmtMatch) {
188        let label = format!("match {}", node.subject);
189        let mut fc = Flowchart::new_sub(label);
190
191        let mut curr = &mut fc.items;
192        for case in node.cases {
193            let items = {
194                let mut fc = Flowchart::new_unlabeled_sub();
195                let mut gen = FlowchartGenerator::new(&mut fc);
196                gen.visit_stmts(case.body);
197                fc.items
198            };
199
200            if matches!(
201                case.pattern,
202                ast::Pattern::MatchAs(ast::PatternMatchAs {
203                    pattern: None,
204                    name: None,
205                    ..
206                })
207            ) {
208                // Wildcard. Add items into current case.
209                curr.extend(items);
210            } else {
211                // Not wildcard.
212                let label = format!("case: {} ?", pattern_to_string(&case.pattern));
213                curr.push(FlowchartItem::Condition(flowchart::Condition::new(label)));
214                let FlowchartItem::Condition(cond) = curr.last_mut().unwrap() else {
215                    unreachable!("bug: last item is not a condition");
216                };
217                cond.then_items = items;
218                curr = &mut cond.else_items;
219            }
220        }
221
222        self.fc.items.push(FlowchartItem::SubFlowchart(fc));
223    }
224
225    fn visit_stmt_raise(&mut self, node: ast::StmtRaise) {
226        let label = format!(
227            "raise {}",
228            node.exc.map_or_else(String::new, |e| e.to_string())
229        );
230
231        self.fc
232            .items
233            .push(FlowchartItem::Terminal(flowchart::Terminal::new(label)));
234    }
235
236    fn visit_stmt_try(&mut self, node: ast::StmtTry) {
237        let mut fc = Flowchart::new_sub("try");
238        let mut gen = FlowchartGenerator::new(&mut fc);
239        gen.visit_stmts(node.body);
240        self.fc.items.push(FlowchartItem::SubFlowchart(fc));
241
242        let mut curr = &mut self.fc.items;
243        for handler in node.handlers {
244            let ast::ExceptHandler::ExceptHandler(handler) = handler;
245
246            let label = match (&handler.type_, &handler.name) {
247                (None, None) => "except ?".to_string(),
248                (Some(t), None) => format!("except {} ?", t),
249                (None, Some(n)) => format!("except as {} ?", n),
250                (Some(t), Some(n)) => format!("except {} as {} ?", t, n),
251            };
252
253            let items = {
254                let mut fc = Flowchart::new_unlabeled_sub();
255                let mut gen = FlowchartGenerator::new(&mut fc);
256                gen.visit_stmts(handler.body);
257                fc.items
258            };
259
260            curr.push(FlowchartItem::Condition(flowchart::Condition::new(label)));
261            let FlowchartItem::Condition(cond) = curr.last_mut().unwrap() else {
262                unreachable!("bug: last item is not a condition");
263            };
264
265            cond.then_items = items;
266            curr = &mut cond.else_items;
267        }
268    }
269
270    fn visit_stmt_expr(&mut self, node: ast::StmtExpr) {
271        let label = node.value.to_string().replace('"', "");
272
273        self.fc
274            .items
275            .push(FlowchartItem::Step(flowchart::Step::new(label)));
276    }
277
278    fn visit_stmt_break(&mut self, _node: ast::StmtBreak) {
279        let label = "break".to_string();
280        self.fc
281            .items
282            .push(FlowchartItem::Break(flowchart::Break::new(label)));
283    }
284
285    fn visit_stmt_continue(&mut self, _node: ast::StmtContinue) {
286        let label = "continue".to_string();
287        self.fc
288            .items
289            .push(FlowchartItem::Continue(flowchart::Continue::new(label)));
290    }
291}
292
293fn op_to_string(op: &ast::Operator) -> &str {
294    match op {
295        ast::Operator::Add => "+",
296        ast::Operator::Sub => "-",
297        ast::Operator::Mult => "*",
298        ast::Operator::MatMult => "@",
299        ast::Operator::Div => "/",
300        ast::Operator::Mod => "%",
301        ast::Operator::Pow => "**",
302        ast::Operator::LShift => "<<",
303        ast::Operator::RShift => ">>",
304        ast::Operator::BitOr => "|",
305        ast::Operator::BitXor => "^",
306        ast::Operator::BitAnd => "&",
307        ast::Operator::FloorDiv => "//",
308    }
309}
310
311fn pattern_to_string(pattern: &ast::Pattern) -> String {
312    match pattern {
313        ast::Pattern::MatchSingleton(ast::PatternMatchSingleton { value, .. }) => value.to_string(),
314        ast::Pattern::MatchValue(ast::PatternMatchValue { value, .. }) => value.to_string(),
315        ast::Pattern::MatchAs(ast::PatternMatchAs { pattern, name, .. }) => {
316            let pattern = pattern
317                .as_ref()
318                .map_or_else(|| "_".to_string(), |p| pattern_to_string(p));
319            if let Some(name) = name {
320                format!("{} as {}", pattern, name)
321            } else {
322                pattern
323            }
324        }
325        ast::Pattern::MatchSequence(ast::PatternMatchSequence { .. }) => "[...]".to_string(),
326        ast::Pattern::MatchMapping(ast::PatternMatchMapping { .. }) => "{...}".to_string(),
327        ast::Pattern::MatchClass(ast::PatternMatchClass { cls, .. }) => {
328            format!("{}()", cls)
329        }
330        ast::Pattern::MatchStar(ast::PatternMatchStar { name, .. }) => {
331            if let Some(name) = name {
332                format!("*{}", name)
333            } else {
334                "*".to_string()
335            }
336        }
337        ast::Pattern::MatchOr(ast::PatternMatchOr { patterns, .. }) => {
338            format!("({})", patterns.iter().map(pattern_to_string).format(" | "))
339        }
340    }
341}
python_to_mermaid/convert/python_to_flowchart.rs

python_to_mermaid/convert/
python_to_flowchart.rs
