pysrc2cpg: expand dynamic type hints based on imports (#2327)

* Add a pass to exapnd type hints

* Add test

---------

Co-authored-by: Fabian Yamaguchi <[email protected]>

console/src/main/scala/io/joern/console/cpgcreation/PythonSrcCpgGenerator.scala

@@ -2,7 +2,13 @@ package io.joern.console.cpgcreation
 
 import io.joern.console.FrontendConfig
 import io.shiftleft.codepropertygraph.Cpg
-import io.joern.pysrc2cpg.{ImportsPass, PythonNaiveCallLinker, PythonTypeHintCallLinker, PythonTypeRecovery}
+import io.joern.pysrc2cpg.{
+  DynamicTypeHintFullNamePass,
+  ImportsPass,
+  PythonNaiveCallLinker,
+  PythonTypeHintCallLinker,
+  PythonTypeRecovery
+}
 
 import java.nio.file.Path
 
@@ -25,6 +31,7 @@ case class PythonSrcCpgGenerator(config: FrontendConfig, rootPath: Path) extends
 
   override def applyPostProcessingPasses(cpg: Cpg): Cpg = {
     new ImportsPass(cpg).createAndApply()
+    new DynamicTypeHintFullNamePass(cpg).createAndApply()
     new PythonTypeRecovery(cpg).createAndApply()
     new PythonTypeHintCallLinker(cpg).createAndApply()
     new PythonNaiveCallLinker(cpg).createAndApply()

joern-cli/frontends/pysrc2cpg/src/main/scala/io/joern/pysrc2cpg/DynamicTypeHintFullNamePass.scala

@@ -0,0 +1,32 @@
+package io.joern.pysrc2cpg
+
+import io.shiftleft.codepropertygraph.Cpg
+import io.shiftleft.codepropertygraph.generated.PropertyNames
+import io.shiftleft.passes.CpgPass
+import overflowdb.BatchedUpdate
+import io.shiftleft.semanticcpg.language._
+
+/** The type hints we pick up via the parser are not full names. This pass fixes that by retrieving the import for each
+  * dynamic type hint and adjusting the dynamic type hint full name field accordingly.
+  */
+class DynamicTypeHintFullNamePass(cpg: Cpg) extends CpgPass(cpg) {
+  override def run(diffGraph: BatchedUpdate.DiffGraphBuilder): Unit = {
+    val fileToImports = cpg.imports.l
+      .flatMap { imp =>
+        imp.call.file.l.map { f => f.name -> imp }
+      }
+      .groupBy(_._1)
+      .view
+      .mapValues(_.map(_._2))
+
+    for {
+      methodReturn   <- cpg.methodReturn.filter(x => x.dynamicTypeHintFullName.nonEmpty)
+      typeHint       <- methodReturn.dynamicTypeHintFullName
+      file           <- methodReturn.file
+      imports        <- fileToImports.get(file.name)
+      importedEntity <- imports.importedAsExact(typeHint).importedEntity
+    } {
+      diffGraph.setNodeProperty(methodReturn, PropertyNames.DYNAMIC_TYPE_HINT_FULL_NAME, importedEntity)
+    }
+  }
+}

joern-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/PySrc2CpgFixture.scala

@@ -27,6 +27,7 @@ class PySrcTestCpg extends TestCpg with PythonFrontend {
   override def applyPasses(): Unit = {
     X2Cpg.applyDefaultOverlays(this)
     new ImportsPass(this).createAndApply()
+    new DynamicTypeHintFullNamePass(this).createAndApply()
     new PythonTypeRecovery(this).createAndApply()
     new PythonTypeHintCallLinker(this).createAndApply()
     new PythonNaiveCallLinker(this).createAndApply()

joern-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/passes/DynamicTypeHintFullNamePassTests.scala

@@ -0,0 +1,22 @@
+package io.joern.pysrc2cpg.passes
+
+import io.joern.pysrc2cpg.PySrc2CpgFixture
+import io.shiftleft.semanticcpg.language._
+
+class DynamicTypeHintFullNamePassTests extends PySrc2CpgFixture(withOssDataflow = false) {
+
+  "dynamic type hints" should {
+    lazy val cpg = code("""
+        |from foo.bar import Woo
+        |
+        |def m() -> Woo:
+        |   x
+        |
+        |""".stripMargin)
+
+    "take into accounts imports" in {
+      cpg.method("m").methodReturn.dynamicTypeHintFullName.l shouldBe List("foo.bar.Woo")
+    }
+  }
+
+}