diff --git a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt index f035f95df8..232c1b6900 100644 --- a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt +++ b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt @@ -26,6 +26,7 @@ package de.fraunhofer.aisec.cpg.frontends.python import de.fraunhofer.aisec.cpg.graph.* +import de.fraunhofer.aisec.cpg.graph.declarations.MethodDeclaration import de.fraunhofer.aisec.cpg.graph.statements.expressions.Expression import de.fraunhofer.aisec.cpg.graph.statements.expressions.ProblemExpression @@ -181,6 +182,24 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) : } private fun handleName(node: PythonAST.Name): Expression { - return newDeclaredReferenceExpression(name = node.id, rawNode = node) + val r = newDeclaredReferenceExpression(name = node.id, rawNode = node) + + /* + * TODO: this is not nice... :( + * + * Take a little shortcut and set refersTo, in case this is a method receiver. This allows us to play more + * nicely with member (call) expressions on the current class, since then their base type is known. + */ + val currentFunction = frontend.scopeManager.currentFunction + if (currentFunction is MethodDeclaration) { + val recv = currentFunction.receiver + recv.let { + if (node.id == it?.name?.localName) { + r.refersTo = it + r.type = it.type + } + } + } + return r } } 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 4ec6be4cf1..0e219ead27 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 @@ -214,6 +214,7 @@ class StatementHandler(frontend: PythonLanguageFrontend) : implicitInitializerAllowed = false, rawNode = recvPythonNode ) + frontend.scopeManager.addDeclaration(recvNode) when (result) { is ConstructorDeclaration -> result.receiver = recvNode is MethodDeclaration -> result.receiver = recvNode diff --git a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt index 05404ca6f8..34db4f3830 100644 --- a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt +++ b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt @@ -27,19 +27,20 @@ package de.fraunhofer.aisec.cpg.passes import de.fraunhofer.aisec.cpg.TranslationContext import de.fraunhofer.aisec.cpg.frontends.python.PythonLanguageFrontend -import de.fraunhofer.aisec.cpg.graph.Component -import de.fraunhofer.aisec.cpg.graph.Node -import de.fraunhofer.aisec.cpg.graph.newFieldDeclaration -import de.fraunhofer.aisec.cpg.graph.newVariableDeclaration +import de.fraunhofer.aisec.cpg.graph.* +import de.fraunhofer.aisec.cpg.graph.declarations.FieldDeclaration +import de.fraunhofer.aisec.cpg.graph.declarations.MethodDeclaration +import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression import de.fraunhofer.aisec.cpg.graph.statements.expressions.DeclaredReferenceExpression +import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression import de.fraunhofer.aisec.cpg.helpers.SubgraphWalker import de.fraunhofer.aisec.cpg.passes.order.ExecuteFirst import de.fraunhofer.aisec.cpg.passes.order.RequiredFrontend @ExecuteFirst @RequiredFrontend(PythonLanguageFrontend::class) -class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) { +class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), NamespaceProvider { override fun cleanup() { // nothing to do } @@ -57,45 +58,93 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) { * This function checks for each [AssignExpression] whether there is already a matching variable * or not. New variables can be one of: * - [FieldDeclaration] if we are currently in a record - * - [VariableDeclaratrion] otherwise + * - [VariableDeclaration] otherwise * * TODO: loops */ - private fun handle(assignExpression: Node?) { - if (assignExpression !is AssignExpression) { - return + private fun handle(node: Node?) { + when (node) { + is AssignExpression -> handleAssignExpression(node) + is DeclaredReferenceExpression -> handleDeclaredReferenceExpression(node) + else -> {} } + } - for (target in assignExpression.lhs) { - (target as? DeclaredReferenceExpression)?.let { - val resolved = scopeManager.resolveReference(it) - - if (resolved == null) { - val decl = - if (scopeManager.isInRecord && !scopeManager.isInFunction) { - val field = newFieldDeclaration(it.name, code = it.code) - field.location = it.location + /* + * Return null when not creating a new decl + */ + private fun handleDeclaredReferenceExpression( + node: DeclaredReferenceExpression + ): VariableDeclaration? { + val resolved = scopeManager.resolveReference(node) + if (resolved == null) { + val decl = + if (scopeManager.isInRecord) { + if (scopeManager.isInFunction) { + if ( + node is MemberExpression && + node.base.name == + (scopeManager.currentFunction as? MethodDeclaration) + ?.receiver + ?.name + ) { + val field = newFieldDeclaration(node.name, code = node.code) + field.location = node.location scopeManager.currentRecord?.addField(field) // TODO why do we need this? field } else { - val v = newVariableDeclaration(it.name, code = it.code) - v.location = it.location + val v = newVariableDeclaration(node.name, code = node.code) + v.location = node.location + scopeManager.currentFunction?.addDeclaration(v) v } + } else { + val field = newFieldDeclaration(node.name, code = node.code) + field.location = node.location + scopeManager.currentRecord?.addField(field) // TODO why do we need this? + field + } + } else { if (scopeManager.isInFunction) { - scopeManager.currentFunction?.addDeclaration( - decl - ) // TODO why do we need this? + val v = newVariableDeclaration(node.name, code = node.code) + v.location = node.location + scopeManager.currentFunction + ?.body + ?.addDeclaration(v) // TODO why do we need this? + v + } else { + val v = newVariableDeclaration(node.name, code = node.code) + v.location = node.location + v } - decl.isImplicit = true + } + + decl.isImplicit = true + + decl.scope = scopeManager.currentScope // TODO why do we need this? + scopeManager.addDeclaration(decl) + return decl + } else { + return null + } + } + + private fun handleAssignExpression(assignExpression: AssignExpression) { + for (target in assignExpression.lhs) { + (target as? DeclaredReferenceExpression)?.let { + val resolved = scopeManager.resolveReference(it) + if (resolved == null) { + val decl = handleDeclaredReferenceExpression(it) assignExpression.findValue(it)?.let { value -> - decl.type = value.type + decl?.type = value.type } // TODO why do we need this (testCtor test case for example)? - assignExpression.declarations += decl - decl.scope = scopeManager.currentScope // TODO why do we need this? - scopeManager.addDeclaration(decl) + + decl?.let { d -> assignExpression.declarations += d } } } } } + + override val namespace: Name? + get() = scopeManager.currentNamespace } diff --git a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt index a811832d55..11feeaec59 100644 --- a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt +++ b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt @@ -365,10 +365,10 @@ class PythonFrontendTest : BaseTest() { assertLocalName("z", fieldZ) assertLocalName("baz", fieldBaz) - assertNull(fieldX.initializer) - assertNotNull(fieldY.initializer) - assertNull(fieldZ.initializer) - assertNotNull(fieldBaz.initializer) + assertNull(fieldX.firstAssignment) + assertNotNull(fieldY.firstAssignment) + assertNull(fieldZ.firstAssignment) + assertNotNull(fieldBaz.firstAssignment) val methBar = recordFoo.methods[0] assertNotNull(methBar) @@ -378,13 +378,12 @@ class PythonFrontendTest : BaseTest() { assertNotNull(barZ) assertEquals(fieldZ, barZ.refersTo) - val barBaz = - (methBar.body as? CompoundStatement)?.statements?.get(1) as? DeclarationStatement + val barBaz = (methBar.body as? CompoundStatement)?.statements?.get(1) as? AssignExpression assertNotNull(barBaz) val barBazInner = barBaz.declarations[0] as? FieldDeclaration assertNotNull(barBazInner) assertLocalName("baz", barBazInner) - assertNotNull(barBazInner.initializer) + assertNotNull(barBazInner.firstAssignment) } @Test @@ -404,7 +403,7 @@ class PythonFrontendTest : BaseTest() { assertNotNull(recordFoo) assertLocalName("Foo", recordFoo) - assertEquals(1, recordFoo.fields.size) + assertEquals(2, recordFoo.fields.size) val somevar = recordFoo.fields[0] assertNotNull(somevar) assertLocalName("somevar", somevar)