From 1f204b77d741017154391f9dc7b1143be4319808 Mon Sep 17 00:00:00 2001
From: Alexander Kuechler <alexander.kuechler@aisec.fraunhofer.de>
Date: Fri, 18 Oct 2024 12:27:07 +0200
Subject: [PATCH] Support async with

---
 .../de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt   |  4 ++--
 .../aisec/cpg/frontends/python/StatementHandler.kt | 14 +++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt
index ffa70363df..6e323eb7c5 100644
--- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt
+++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt
@@ -284,8 +284,8 @@ fun <T : Node> T.codeAndLocationFrom(other: Node): T {
  * expression handler.
  */
 context(CodeAndLocationProvider<AstNode>)
-fun <T : Node, AstNode> T.codeAndLocationFromOtherRawNode(rawNode: AstNode): T {
-    setCodeAndLocation(this@CodeAndLocationProvider, rawNode)
+fun <T : Node, AstNode> T.codeAndLocationFromOtherRawNode(rawNode: AstNode?): T {
+    rawNode?.let { setCodeAndLocation(this@CodeAndLocationProvider, it) }
     return this
 }
 
diff --git a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt
index 27e9b3cec4..b9845db25a 100644
--- a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt
+++ b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt
@@ -72,13 +72,13 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
             is Python.AST.Assert -> handleAssert(node)
             is Python.AST.Try -> handleTryStatement(node)
             is Python.AST.Delete -> handleDelete(node)
-            is Python.AST.With -> handleWithStatement(node)
+            is Python.AST.With,
+            is Python.AST.AsyncWith -> handleWithStatement(node)
             is Python.AST.Global -> handleGlobal(node)
             is Python.AST.Nonlocal -> handleNonLocal(node)
             is Python.AST.Match,
             is Python.AST.Raise,
-            is Python.AST.TryStar,
-            is Python.AST.AsyncWith ->
+            is Python.AST.TryStar ->
                 newProblemExpression(
                     "The statement of class ${node.javaClass} is not supported yet",
                     rawNode = node
@@ -122,7 +122,7 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
      *     manager.__exit__(None, None, None)
      * ```
      */
-    private fun handleWithStatement(node: Python.AST.With): Block {
+    private fun handleWithStatement(node: Python.AST.NormalOrAsyncWith): Block {
         /**
          * Prepares the `manager = ContextManager()` and returns the random name for the "manager"
          * as well as the assignment.
@@ -236,8 +236,8 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
                 tmpValName
             )
         }
-
-        val result = newBlock().codeAndLocationFromOtherRawNode(node).implicit()
+        val result =
+            newBlock().codeAndLocationFromOtherRawNode(node as? Python.AST.BaseStmt).implicit()
 
         // If there are multiple elements in node.items, we have to nest the try statements.
         // We start with a generic block for the outer context manager.
@@ -283,7 +283,7 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
                                         )
                                     }
                                 }
-                                .codeAndLocationFromOtherRawNode(node)
+                                .codeAndLocationFromOtherRawNode(node as? Python.AST.BaseStmt)
                                 .implicit()
                         // Add the catch block
                         this.catchClauses.add(