[pysrc2cpg] Set new MODULE modifier on Python module methods (#3826)

joernio · Nov 15, 2023 · cb2a6ef · cb2a6ef
1 parent dea6ab2
commit cb2a6ef
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git a/build.sbt b/build.sbt
@@ -2,7 +2,7 @@ name                     := "joern"
 ThisBuild / organization := "io.joern"
 ThisBuild / scalaVersion := "3.3.1"
 
-val cpgVersion = "1.4.25"
+val cpgVersion = "1.4.29"
 
 lazy val joerncli          = Projects.joerncli
 lazy val querydb           = Projects.querydb

diff --git a/...-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/LambdaTests.scala b/...-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/LambdaTests.scala
@@ -702,7 +702,7 @@ class LambdaTests extends JavaSrcCode2CpgFixture {
       cpg.all.collectAll[ClosureBinding].l match {
         case myValue :: Nil =>
           myValue.closureOriginalName.head shouldBe "myValue"
-          myValue._localViaRefOut.name shouldBe "myValue"
+          myValue._localViaRefOut.get.name shouldBe "myValue"
           myValue._captureIn.collectFirst { case x: MethodRef =>
             x.methodFullName
           }.head shouldBe "Foo.lambda$0:<unresolvedSignature>(1)"

diff --git a/joern-cli/frontends/pysrc2cpg/src/main/scala/io/joern/pysrc2cpg/PythonAstVisitor.scala b/joern-cli/frontends/pysrc2cpg/src/main/scala/io/joern/pysrc2cpg/PythonAstVisitor.scala
@@ -96,6 +96,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
         "<module>",
         methodFullName,
         Some("<module>"),
+        ModifierTypes.VIRTUAL :: ModifierTypes.MODULE :: Nil,
         parameterProvider = () => MethodParameters.empty(),
         bodyProvider = () =>
           createBuiltinIdentifiers(memOpCalculator.names)
@@ -332,6 +333,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
         methodName,
         methodFullName,
         scopeName,
+        ModifierTypes.VIRTUAL :: Nil,
         parameterProvider,
         bodyProvider,
         returns,
@@ -351,6 +353,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
     name: String,
     fullName: String,
     scopeName: Option[String],
+    modifiers: List[String],
     parameterProvider: () => MethodParameters,
     bodyProvider: () => Iterable[nodes.NewNode],
     returns: Option[ast.iexpr],
@@ -367,8 +370,12 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
 
     contextStack.pushMethod(scopeName, methodNode, blockNode, methodRefNode)
 
-    val virtualModifierNode = nodeBuilder.modifierNode(ModifierTypes.VIRTUAL)
-    edgeBuilder.astEdge(virtualModifierNode, methodNode, 0)
+    var order = 0
+    for (modifier <- modifiers) {
+      val modifierNode = nodeBuilder.modifierNode(modifier)
+      edgeBuilder.astEdge(modifierNode, methodNode, order)
+      order += 1
+    }
 
     val methodParameter = parameterProvider()
     val parameterOrder  = new AutoIncIndex(methodParameter.posStartIndex)
@@ -601,6 +608,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
       adapterMethodName,
       adapterMethodFullName,
       Some(adaptedMethodName),
+      ModifierTypes.VIRTUAL :: Nil,
       parameterProvider = () => {
         MethodParameters(
           0,
@@ -692,6 +700,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
       methodName,
       methodFullName,
       Some(methodName),
+      ModifierTypes.VIRTUAL :: Nil,
       parameterProvider = () => {
         MethodParameters(1, convert(parametersWithoutSelf, 1))
       },
@@ -742,6 +751,7 @@ class PythonAstVisitor(relFileName: String, protected val nodeToCode: NodeToCode
       newMethodName,
       newMethodStubFullName,
       Some(newMethodName),
+      ModifierTypes.VIRTUAL :: Nil,
       parameterProvider = () => {
         MethodParameters(
           0,

diff --git a/...n-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/cpg/FunctionDefCpgTests.scala b/...n-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/cpg/FunctionDefCpgTests.scala
@@ -1,7 +1,8 @@
 package io.joern.pysrc2cpg.cpg
 
 import io.joern.pysrc2cpg.{Constants, Py2CpgTestContext}
-import io.shiftleft.semanticcpg.language._
+import io.shiftleft.codepropertygraph.generated.ModifierTypes
+import io.shiftleft.semanticcpg.language.*
 import org.scalatest.freespec.AnyFreeSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -196,4 +197,15 @@ class FunctionDefCpgTests extends AnyFreeSpec with Matchers {
     }
   }
 
+  "module function" - {
+    lazy val cpg = Py2CpgTestContext.buildCpg("""
+        |""".stripMargin)
+    "test existence of MODULE modifier on module method node" in {
+      cpg.method
+        .name("<module>")
+        .modifier
+        .modifierType(ModifierTypes.MODULE)
+        .nonEmpty shouldBe true
+    }
+  }
 }