From 68b3d575633222a2af66ea0235e3204ba9885714 Mon Sep 17 00:00:00 2001
From: Taus <tausbn@github.com>
Date: Tue, 28 Apr 2026 14:59:11 +0000
Subject: [PATCH] Cleanup, printCFG

Co-authored-by: yoff <yoff@github.com>
---
 python/ql/lib/printCfgNew.ql                  | 45 +++++++++++++++++++
 .../controlflow/internal/AstNodeImpl.qll      | 43 +++++++++++++++++-
 2 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 python/ql/lib/printCfgNew.ql

diff --git a/python/ql/lib/printCfgNew.ql b/python/ql/lib/printCfgNew.ql
new file mode 100644
index 00000000000..7c098cbf8f6
--- /dev/null
+++ b/python/ql/lib/printCfgNew.ql
@@ -0,0 +1,45 @@
+/**
+ * @name Print CFG (New)
+ * @description Produces a representation of a file's Control Flow Graph
+ *              using the new shared control flow library.
+ *              This query is used by the VS Code extension.
+ * @id python/print-cfg
+ * @kind graph
+ * @tags ide-contextual-queries/print-cfg
+ */
+
+private import python as Py
+import semmle.python.controlflow.internal.AstNodeImpl
+
+external string selectedSourceFile();
+
+private predicate selectedSourceFileAlias = selectedSourceFile/0;
+
+external int selectedSourceLine();
+
+private predicate selectedSourceLineAlias = selectedSourceLine/0;
+
+external int selectedSourceColumn();
+
+private predicate selectedSourceColumnAlias = selectedSourceColumn/0;
+
+module ViewCfgQueryInput implements ControlFlow::ViewCfgQueryInputSig<Py::File> {
+  predicate selectedSourceFile = selectedSourceFileAlias/0;
+
+  predicate selectedSourceLine = selectedSourceLineAlias/0;
+
+  predicate selectedSourceColumn = selectedSourceColumnAlias/0;
+
+  predicate cfgScopeSpan(
+    AstSigImpl::Callable callable, Py::File file, int startLine, int startColumn, int endLine,
+    int endColumn
+  ) {
+    exists(Py::Scope scope |
+      scope = callable.asScope() and
+      file = scope.getLocation().getFile() and
+      scope.getLocation().hasLocationInfo(_, startLine, startColumn, endLine, endColumn)
+    )
+  }
+}
+
+import ControlFlow::ViewCfgQuery<Py::File, ViewCfgQueryInput>
diff --git a/python/ql/lib/semmle/python/controlflow/internal/AstNodeImpl.qll b/python/ql/lib/semmle/python/controlflow/internal/AstNodeImpl.qll
index c5e2d010688..15ec5dbfa73 100644
--- a/python/ql/lib/semmle/python/controlflow/internal/AstNodeImpl.qll
+++ b/python/ql/lib/semmle/python/controlflow/internal/AstNodeImpl.qll
@@ -11,6 +11,7 @@
 private import python as Py
 private import codeql.controlflow.ControlFlowGraph
 private import codeql.controlflow.SuccessorType
+private import codeql.util.Void
 
 private module Ast {
   /** The newtype representing AST nodes for the shared CFG library. */
@@ -717,6 +718,8 @@ module AstSigImpl implements AstSig<Py::Location> {
       index = 0 and result = w.getTest()
       or
       index = 1 and result = w.getBody()
+      or
+      index = 2 and result = w.getOrelse()
     )
     or
     // ForStmt (mapped as ForeachStmt): collection (0), variable (1), body (2)
@@ -1046,12 +1049,17 @@ module AstSigImpl implements AstSig<Py::Location> {
     Expr getExpr() { result = this.getValue() }
   }
 
-  /** A `raise` statement (mapped to `ThrowStmt`). */
-  class ThrowStmt extends Stmt, Ast::RaiseNode {
+  /** A `raise` statement (mapped to `Throw`). */
+  class Throw extends Stmt, Ast::RaiseNode {
     /** Gets the expression being raised. */
     Expr getExpr() { result = this.getException() }
   }
 
+  /** A `goto` statement. Python has no goto. */
+  class GotoStmt extends Stmt {
+    GotoStmt() { none() }
+  }
+
   // ===== Try/except =====
   /** A `try` statement. */
   class TryStmt extends Stmt {
@@ -1171,6 +1179,26 @@ module AstSigImpl implements AstSig<Py::Location> {
     NullCoalescingExpr() { none() }
   }
 
+  /** An assignment expression. Python has no assignment expressions in the BinaryExpr sense. */
+  class Assignment extends BinaryExpr {
+    Assignment() { none() }
+  }
+
+  /** A simple assignment expression. */
+  class AssignExpr extends Assignment { }
+
+  /** A compound assignment expression. */
+  class CompoundAssignment extends Assignment { }
+
+  /** A short-circuiting logical AND compound assignment. Python has no `&&=` operator. */
+  class AssignLogicalAndExpr extends CompoundAssignment { }
+
+  /** A short-circuiting logical OR compound assignment. Python has no `||=` operator. */
+  class AssignLogicalOrExpr extends CompoundAssignment { }
+
+  /** A short-circuiting null-coalescing compound assignment. Python has no `??=` operator. */
+  class AssignNullCoalescingExpr extends CompoundAssignment { }
+
   /** A unary expression. Exists for the `not` subclass. */
   class UnaryExpr extends Expr {
     UnaryExpr() { this instanceof Ast::NotExprNode }
@@ -1186,6 +1214,15 @@ module AstSigImpl implements AstSig<Py::Location> {
     /** Gets the boolean value of this literal. */
     boolean getValue() { result = this.getBoolValue() }
   }
+
+  /** A pattern match expression. Python has no `instanceof`-style pattern match expr. */
+  class PatternMatchExpr extends Expr {
+    PatternMatchExpr() { none() }
+
+    Expr getExpr() { none() }
+
+    AstNode getPattern() { none() }
+  }
 }
 
 private module Cfg0 = Make0<Py::Location, AstSigImpl>;
@@ -1209,6 +1246,8 @@ private module Input implements InputSig1, InputSig2 {
     string toString() { result = "label" }
   }
 
+  class CallableBodyPartContext = Void;
+
   predicate inConditionalContext(AstSigImpl::AstNode n, ConditionKind kind) {
     kind.isBoolean() and
     n = any(Ast::AssertNode a).getTest()