Basic support for raise

Author: tmr232

src/control-flow/cfg-defs.ts

@@ -2,6 +2,8 @@ import { MultiDirectedGraph } from "graphology";
 import type Parser from "web-tree-sitter";
 
 export type NodeType =
+  | "YIELD"
+  | "RAISE"
   | "MARKER_COMMENT"
   | "LOOP_HEAD"
   | "LOOP_EXIT"

src/control-flow/cfg-python.ts

@@ -171,12 +171,25 @@ export class CFGBuilder {
         return this.processWithStatement(node);
       case "try_statement":
         return this.processTryStatement(node);
-      default: {
-        const newNode = this.addNode("STATEMENT", node.text);
-        return { entry: newNode, exit: newNode };
-      }
+      case "raise_statement":
+        return this.processRaiseStatement(node);
+      default:
+        return this.defaultProcessStatement(node);
     }
   }
+  private defaultProcessStatement(syntax: Parser.SyntaxNode): BasicBlock {
+    // const match = this.matchQuery(syntax, "yield", `(yield) @yield`);
+    // if (match.captures.length > 0) {
+    //   const yieldNode = this.addNode("YIELD", syntax.text);
+    //   return { entry: yieldNode, exit: yieldNode };
+    // }
+    const newNode = this.addNode("STATEMENT", syntax.text);
+    return { entry: newNode, exit: newNode };
+  }
+  private processRaiseStatement(raiseSyntax: Parser.SyntaxNode): BasicBlock {
+    const raiseNode = this.addNode("RAISE", raiseSyntax.text);
+    return { entry: raiseNode, exit: null };
+  }
   private processReturnStatement(returnSyntax: Parser.SyntaxNode): BasicBlock {
     const returnNode = this.addNode("RETURN", returnSyntax.text);
     return { entry: returnNode, exit: null, returns: [returnNode] };

src/control-flow/render.ts

@@ -116,35 +116,6 @@ function getParents(cluster: Cluster) {
   return parents.toReversed();
 }
 
-function renderNode(graph: CFGGraph, node: string, verbose: boolean): string {
-  let label = "";
-  if (verbose) {
-    label = `${node} ${graph.getNodeAttributes(node).type} ${graph.getNodeAttributes(node).code}`;
-
-    const clusterAttrs = graph.getNodeAttribute(node, "cluster");
-    label = `${clusterAttrs?.id} ${clusterAttrs?.type}\n${label}`;
-  }
-  let shape = "box";
-  let fillColor = "lightgray";
-  let minHeight = 0.2;
-  if (graph.degree(node) === 0) {
-    minHeight = 0.5;
-  } else if (graph.inDegree(node) === 0) {
-    shape = "invhouse";
-    fillColor = "#48AB30";
-    minHeight = 0.5;
-  } else if (graph.outDegree(node) === 0) {
-    shape = "house";
-    fillColor = "#AB3030";
-    minHeight = 0.5;
-  }
-
-  const height = Math.max(
-    graph.getNodeAttribute(node, "lines") * 0.3,
-    minHeight,
-  );
-  return `    ${node} [label="${label}" shape="${shape}" fillcolor="${fillColor}" style="filled" height=${height}];\n`;
-}
 
 function renderHierarchy(
   cfg: CFG,
@@ -335,3 +306,46 @@ function renderEdge(
   }
   return `    ${source} -> ${target} [${formatStyle(dotAttrs)}];\n`;
 }
+
+function renderNode(graph: CFGGraph, node: string, verbose: boolean): string {
+  const dotAttrs: DotAttributes = {};
+  dotAttrs.style = "filled"
+  dotAttrs.label = "";
+  const nodeAttrs = graph.getNodeAttributes(node);
+  if (verbose) {
+    dotAttrs.label = `${node} ${nodeAttrs.type} ${graph.getNodeAttributes(node).code}`;
+
+    const clusterAttrs = graph.getNodeAttribute(node, "cluster");
+    dotAttrs.label = `${clusterAttrs?.id} ${clusterAttrs?.type}\n${dotAttrs.label}`;
+  }
+  dotAttrs.shape = "box";
+  dotAttrs.fillcolor = "lightgray";
+  let minHeight = 0.2;
+  if (graph.degree(node) === 0) {
+    dotAttrs.minHeight = 0.5;
+  } else if (graph.inDegree(node) === 0) {
+    dotAttrs.shape = "invhouse";
+    dotAttrs.fillcolor = "#48AB30";
+    minHeight = 0.5;
+  } else if (graph.outDegree(node) === 0) {
+    dotAttrs.shape = "house";
+    dotAttrs.fillcolor = "#AB3030";
+    minHeight = 0.5;
+  }
+  switch (nodeAttrs.type) {
+    case "RAISE":
+      dotAttrs.shape = "triangle";
+      dotAttrs.fillcolor = "#fdd";
+      break;
+    case "YIELD":
+      dotAttrs.shallow = "hexagon";
+      dotAttrs.fillcolor = "aqua";
+      break;
+  }
+
+  dotAttrs.height = Math.max(
+    graph.getNodeAttribute(node, "lines") * 0.3,
+    minHeight,
+  );
+  return `    ${node} [${formatStyle(dotAttrs)}];\n`;
+}

src/test/commentTestSamples/sample.py

@@ -156,6 +156,29 @@ def WithNestedCluster():
             return
 
 
+# render: true,
+# exits: 999
+def raise_exception():
+    raise
+
+# render: true,
+# exits: 999
+def raise_again():
+    try:
+        raise x
+    except: 
+        if x:
+            pass
+    else: 
+        if x:
+            pass
+
+# render: true,
+# exits: 3
+def yield_value():
+    for x in y:
+        yield x
+    
 
 # exits: 0,
 # render: true