Skip to content

Commit bf75d7f

Browse files
fabsx00Fabian Yamaguchi
andauthored
pysrc2cpg: propagation of return types (#2328)
* stash * Propagate dynamic type hint * stash * Fix works now * Remove `println`s * scalafmt --------- Co-authored-by: Fabian Yamaguchi <[email protected]>
1 parent 94d69c7 commit bf75d7f

File tree

2 files changed

+25
-3
lines changed

2 files changed

+25
-3
lines changed

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -502,9 +502,7 @@ class TypeRecoveryPassTests extends PySrc2CpgFixture(withOssDataflow = false) {
502502

503503
"recover its full name successfully" in {
504504
val List(methodFullName) = cpg.call("query").methodFullName.l
505-
methodFullName shouldBe Seq("data", "db_session.py:<module>.create_session.<returnValue>.query").mkString(
506-
File.separator
507-
)
505+
methodFullName shouldBe "sqlalchemy.orm.Session.query"
508506
}
509507
}
510508

joern-cli/frontends/x2cpg/src/main/scala/io/joern/x2cpg/passes/frontend/XTypeRecovery.scala

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ abstract class XTypeRecovery[CompilationUnitType <: AstNode](cpg: Cpg, iteration
6060
compilationUnit
6161
.map(unit => generateRecoveryForCompilationUnitTask(unit, builder).fork())
6262
.foreach(_.get())
63+
6364
} finally {
6465
globalTable.clear()
6566
addedNodes.clear()
@@ -131,6 +132,26 @@ abstract class RecoverForXCompilationUnit[CompilationUnitType <: AstNode](
131132
protected def members: Traversal[Member] =
132133
cu.ast.isMember
133134

135+
protected def visitCall(call: Call): Unit = {
136+
symbolTable.get(call).foreach { methodFullName =>
137+
val index = methodFullName.indexOf("<returnValue>")
138+
if (index != -1) {
139+
val methodToLookup = methodFullName.substring(0, index - 1)
140+
val remainder = methodFullName.substring(index + "<returnValue>".length)
141+
val methods = cpg.method.fullNameExact(methodToLookup).l
142+
methods match {
143+
case List(method) =>
144+
val dynamicTypeHints = method.methodReturn.dynamicTypeHintFullName.toSet
145+
val newTypes = dynamicTypeHints.map(x => x + remainder)
146+
symbolTable.put(call, newTypes)
147+
case List() =>
148+
case _ =>
149+
logger.warn(s"More than a single function matches method full name: $methodToLookup")
150+
}
151+
}
152+
}
153+
}
154+
134155
override def compute(): Unit = try {
135156
prepopulateSymbolTable()
136157
// Set known aliases that point to imports for local and external methods/modules
@@ -141,6 +162,9 @@ abstract class RecoverForXCompilationUnit[CompilationUnitType <: AstNode](
141162
members.foreach(visitMembers)
142163
// Populate local symbol table with assignments
143164
assignments.foreach(visitAssignments)
165+
// Propagate return values
166+
cpg.method.ast.isCall.foreach(visitCall)
167+
144168
// Persist findings
145169
setTypeInformation()
146170
} finally {

0 commit comments

Comments
 (0)