Initial support for yield

Author: tmr232

src/control-flow/cfg-defs.ts

@@ -3,7 +3,7 @@ import type Parser from "web-tree-sitter";
 
 export type NodeType =
   | "YIELD"
-  | "RAISE"
+  | "THROW"
   | "MARKER_COMMENT"
   | "LOOP_HEAD"
   | "LOOP_EXIT"
@@ -154,6 +154,10 @@ export function mergeNodeAttrs(
   if (from.cluster !== into.cluster) {
     return null;
   }
+  const noMergeTypes: NodeType[] = ["YIELD", "THROW"];
+  if (noMergeTypes.includes(from.type) || noMergeTypes.includes(into.type)) {
+    return null;
+  }
   return {
     type: from.type,
     code: `${from.code}\n${into.code}`,

src/control-flow/cfg-python.ts

@@ -178,16 +178,16 @@ export class CFGBuilder {
     }
   }
   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 hasYield = this.matchExistsIn(syntax, "yield", `(yield) @yield`);
+    if (hasYield) {
+      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);
+    const raiseNode = this.addNode("THROW", raiseSyntax.text);
     return { entry: raiseNode, exit: null };
   }
   private processReturnStatement(returnSyntax: Parser.SyntaxNode): BasicBlock {
@@ -506,6 +506,17 @@ export class CFGBuilder {
     return match;
   }
 
+  private matchExistsIn(
+    syntax: Parser.SyntaxNode,
+    mainName: string,
+    queryString: string,
+  ): boolean {
+    const language = syntax.tree.getLanguage();
+    const query = language.query(queryString);
+    const matches = query.matches(syntax);
+    return matches.length > 0;
+  }
+
   private getSyntax(
     match: Parser.QueryMatch,
     name: string,

src/control-flow/render.ts

@@ -86,11 +86,11 @@ function buildHierarchy(cfg: CFG): Hierarchy {
   // for (const cluster of clusterNodes.keys()) {
   //   console.log(cluster.id, getParents(cluster).map(c => c.id))
   // }
-  // showHierarchy(hierarchy);
+  // _showHierarchy(hierarchy);
   return hierarchy;
 }
 
-function showHierarchy(hierarchy: Hierarchy) {
+function _showHierarchy(hierarchy: Hierarchy) {
   const stack: [Hierarchy, number][] = [[hierarchy, 0]];
   const spaces = "                                           ";
   for (
@@ -116,7 +116,6 @@ function getParents(cluster: Cluster) {
   return parents.toReversed();
 }
 
-
 function renderHierarchy(
   cfg: CFG,
   hierarchy: Hierarchy,
@@ -309,7 +308,7 @@ function renderEdge(
 
 function renderNode(graph: CFGGraph, node: string, verbose: boolean): string {
   const dotAttrs: DotAttributes = {};
-  dotAttrs.style = "filled"
+  dotAttrs.style = "filled";
   dotAttrs.label = "";
   const nodeAttrs = graph.getNodeAttributes(node);
   if (verbose) {
@@ -333,13 +332,14 @@ function renderNode(graph: CFGGraph, node: string, verbose: boolean): string {
     minHeight = 0.5;
   }
   switch (nodeAttrs.type) {
-    case "RAISE":
+    case "THROW":
       dotAttrs.shape = "triangle";
       dotAttrs.fillcolor = "#fdd";
       break;
     case "YIELD":
-      dotAttrs.shallow = "hexagon";
-      dotAttrs.fillcolor = "aqua";
+      dotAttrs.shape = "hexagon";
+      dotAttrs.orientation = 90;
+      dotAttrs.fillcolor = "deepskyblue";
       break;
   }
 

src/test/commentTestSamples/sample.py

@@ -161,8 +161,7 @@ def WithNestedCluster():
 def raise_exception():
     raise
 
-# render: true,
-# exits: 999
+# render: true
 def raise_again():
     try:
         raise x
@@ -180,7 +179,8 @@ def yield_value():
         yield x
 
 
-# exits: 0,
+# exits: 2,
+# nodes: 4,
 # render: true
 def try_except():
     try:
@@ -189,7 +189,8 @@ def try_except():
         g()
     return
 
-# exits: 0,
+# exits: 1,
+# nodes: 5,
 # render: true
 def try_except_finally():
     try:
@@ -200,7 +201,8 @@ def try_except_finally():
         h()
     return
 
-# exits: 0,
+# exits: 2,
+# nodes: 9,
 # render: true
 def try_many_except_finally():
     try:
@@ -215,15 +217,15 @@ def try_many_except_finally():
     return
 
 
-# exits: 0,
+# exits: 1,
+# nodes: 6,
 # render: true
 def try_except_else_finally():
     try: pass
     except: pass
     else: pass
     finally: pass
 
-# exits: 0,
 # render: true
 def massive_try_except_else_finally():
     try: 
@@ -249,7 +251,8 @@ def massive_try_except_else_finally():
             for y in a:
                 pass
 
-# exits: 5,
+# exits: 2,
+# nodes: 8,
 # render: true
 def try_finally():
     try:
@@ -259,7 +262,8 @@ def try_finally():
     finally:
         pass
 
-# exits: 5,
+# exits: 1,
+# nodes: 5,
 # render: true
 def try_with_finally():
     try: