Fix Python Assignments

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/AssignExpression.kt

@@ -138,7 +138,7 @@ class AssignExpression :
 
     /** Finds the value (of [rhs]) that is assigned to the particular [lhs] expression. */
     fun findValue(lhsExpression: HasType): Expression? {
-        return if (lhs.size > 1) {
+        return if (lhs.size == 1) {
             rhs.singleOrNull()
         } else {
             // Basically, we need to find out which index on the lhs this variable belongs to and

cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt

@@ -30,6 +30,7 @@ import de.fraunhofer.aisec.cpg.graph.declarations.ImportDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.MethodDeclaration
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
+import de.fraunhofer.aisec.cpg.graph.types.TupleType
 import jep.python.PyObject
 
 class ExpressionHandler(frontend: PythonLanguageFrontend) :
@@ -340,7 +341,7 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
             lst += handle(e)
         }
         val ile = newInitializerListExpression(rawNode = node)
-        ile.type = frontend.objectType("tuple")
+        ile.type = TupleType(types = lst.map { it.type })
         ile.initializers = lst
         return ile
     }

cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt

@@ -634,7 +634,31 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
      * [AssignExpression].
      */
     private fun handleAssign(node: Python.AST.Assign): AssignExpression {
-        val lhs = node.targets.map { frontend.expressionHandler.handle(it) }
+        val lhsCandidate = node.targets.map { frontend.expressionHandler.handle(it) }
+
+        /*
+         * We have to unpack the lhs of the assign expression, because this might be a multi assign.
+         * In this case, the handler returns an [InitializerListExpression] which does not fit very
+         * well for our lhs. We thus unpack the [InitializerListExpression] and directly use the
+         * [Reference]s stored inside.
+         */
+        val lhs =
+            lhsCandidate.flatMap {
+                when (it) {
+                    is Reference -> listOf(it)
+                    is InitializerListExpression -> it.initializers
+                    else ->
+                        listOf(
+                            newProblemExpression(
+                                problem =
+                                    "Expected a `Reference` or an `InitializerListExpression`.",
+                                type = ProblemNode.ProblemType.TRANSLATION,
+                                rawNode = node
+                            )
+                        )
+                }
+            }
+
         node.type_comment?.let { typeComment ->
             val tpe = frontend.typeOf(typeComment)
             lhs.forEach { it.type = tpe }

cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandlerTest.kt

@@ -347,4 +347,67 @@ class ExpressionHandlerTest {
         assertContains(nestedBoolOpDifferentOp2.lhs.nextEOG, nestedBoolOpDifferentOp2Rhs.lhs)
         assertContains(nestedBoolOpDifferentOp2Rhs.lhs.nextEOG, nestedBoolOpDifferentOp2Rhs.rhs)
     }
+
+    @Test
+    fun testAssignmentsMultiReturn() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val result =
+            analyze(listOf(topLevel.resolve("assignments.py").toFile()), topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(result)
+
+        // `a, b = 1, 2`
+        val refA = result.refs["a"]
+        assertIs<Reference>(refA)
+        val refB = result.refs["b"]
+        assertIs<Reference>(refB)
+
+        val assignment = result.assigns.firstOrNull()
+        assertIs<AssignExpression>(assignment)
+        assertEquals(2, assignment.assignments.size)
+
+        // extract both assignments
+        val assignmentA = assignment.assignments[0]
+        assertIs<Assignment>(assignmentA)
+        val assignmentB = assignment.assignments[1]
+        assertIs<Assignment>(assignmentB)
+
+        // check that the assignments assign the correct value and assign to the correct reference
+        assertLiteralValue(1, assignmentA.value)
+        assertLiteralValue(2, assignmentB.value)
+        assertEquals(refA, assignmentA.target)
+        assertEquals(refB, assignmentB.target)
+
+        /*
+        ```python
+        def foo():
+            return (1, 2)
+        fooA, fooB = foo()
+        ```
+         */
+        val refFooA = result.refs["fooA"]
+        assertIs<Reference>(refFooA)
+        val refFooB = result.refs["fooB"]
+        assertIs<Reference>(refFooB)
+
+        val assignmentFoo = result.assigns.getOrNull(1)
+        assertIs<AssignExpression>(assignmentFoo)
+        assertEquals(2, assignmentFoo.assignments.size)
+
+        // extract both assignments
+        val assignmentFooA = assignmentFoo.assignments[0]
+        assertIs<Assignment>(assignmentFooA)
+        val assignmentFooB = assignmentFoo.assignments[1]
+        assertIs<Assignment>(assignmentFooB)
+
+        // check that the assignments assign the correct value and assign to the correct reference
+        assertLiteralValue(3, assignmentFooA.value)
+        assertLiteralValue(4, assignmentFooB.value)
+        assertEquals(refFooA, assignmentFooA.target)
+        assertEquals(refFooB, assignmentFooB.target)
+
+        // check that there is only one call to `foo` modeled
+        assertEquals(1, result.calls.filter { it.invokes == result.functions["foo"] }.size)
+    }
 }

cpg-language-python/src/test/resources/python/assignments.py

@@ -0,0 +1,19 @@
+a, b = (1, 2)
+
+def foo():
+    return (3, 4)
+
+fooA, fooB = foo()
+
+
+# more python code to be used for more tests...
+a = 42
+a, b = [21, 42]
+c = d = 42
+e = 42, 42
+f, g = '42' # yes, this also unpacks strings!
+
+class Foo:
+    foo = 42
+
+a, newVar, Foo.foo, Foo.bar, *rest = 1, 2, 3, 4, 5, 6