From 13f82d4f4bf0ebe9af200ed5d5239b53fb0d7381 Mon Sep 17 00:00:00 2001 From: Leutrim Shala <83644358+lshala@users.noreply.github.com> Date: Wed, 25 Sep 2024 14:00:44 +0200 Subject: [PATCH] Support python type hints (#1701) --------- Co-authored-by: Maximilian Kaul --- .../python/PythonLanguageFrontend.kt | 34 +++++++++---------- .../cpg/frontends/python/StatementHandler.kt | 30 +++++++++------- .../statementHandler/StatementHandlerTest.kt | 31 +++++++++++------ .../src/test/resources/python/type_hints.py | 3 ++ 4 files changed, 58 insertions(+), 40 deletions(-) create mode 100644 cpg-language-python/src/test/resources/python/type_hints.py diff --git a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonLanguageFrontend.kt b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonLanguageFrontend.kt index 029981cbe1..73162462fa 100644 --- a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonLanguageFrontend.kt +++ b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonLanguageFrontend.kt @@ -144,24 +144,7 @@ class PythonLanguageFrontend(language: Language, ctx: Tr autoType() } is Python.AST.Name -> { - // We have some kind of name here; let's quickly check, if this is a primitive type - val id = type.id - if (id in language.primitiveTypeNames) { - return primitiveType(id) - } - - // Otherwise, this could already be a fully qualified type - val name = - if (language.namespaceDelimiter in id) { - // TODO: This might create problem with nested classes - parseName(id) - } else { - // otherwise, we can just simply take the unqualified name and the type - // resolver will take care of the rest - id - } - - objectType(name) + this.typeOf(type.id) } else -> { // The AST supplied us with some kind of type information, but we could not parse @@ -171,6 +154,21 @@ class PythonLanguageFrontend(language: Language, ctx: Tr } } + /** Resolves a [Type] based on its string identifier. */ + fun typeOf(typeId: String): Type { + // Check if the typeId contains a namespace delimiter for qualified types + val name = + if (language.namespaceDelimiter in typeId) { + // TODO: This might create problem with nested classes + parseName(typeId) + } else { + // Unqualified name, resolved by the type resolver + typeId + } + + return objectType(name) + } + /** * This functions extracts the source code from the input file given a location. This is a bit * tricky in Python, as indents are part of the syntax. We also don't want to include leading 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 9834971fe6..a6855e1df2 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 @@ -36,6 +36,7 @@ import de.fraunhofer.aisec.cpg.graph.declarations.* import de.fraunhofer.aisec.cpg.graph.statements.AssertStatement import de.fraunhofer.aisec.cpg.graph.statements.DeclarationStatement import de.fraunhofer.aisec.cpg.graph.statements.Statement +import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression import de.fraunhofer.aisec.cpg.graph.statements.expressions.Block import de.fraunhofer.aisec.cpg.graph.statements.expressions.Expression import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression @@ -196,18 +197,15 @@ class StatementHandler(frontend: PythonLanguageFrontend) : return frontend.expressionHandler.handle(node.value) } - private fun handleAnnAssign(node: Python.AST.AnnAssign): Statement { - // TODO: annotations + /** + * Translates a Python [`AnnAssign`](https://docs.python.org/3/library/ast.html#ast.AnnAssign) + * into an [AssignExpression]. + */ + private fun handleAnnAssign(node: Python.AST.AnnAssign): AssignExpression { val lhs = frontend.expressionHandler.handle(node.target) - return if (node.value != null) { - newAssignExpression( - lhs = listOf(lhs), - rhs = listOf(frontend.expressionHandler.handle(node.value!!)), // TODO !! - rawNode = node - ) - } else { - lhs - } + lhs.type = frontend.typeOf(node.annotation) + val rhs = node.value?.let { listOf(frontend.expressionHandler.handle(it)) } ?: emptyList() + return newAssignExpression(lhs = listOf(lhs), rhs = rhs, rawNode = node) } private fun handleIf(node: Python.AST.If): Statement { @@ -234,8 +232,16 @@ class StatementHandler(frontend: PythonLanguageFrontend) : return ret } - private fun handleAssign(node: Python.AST.Assign): Statement { + /** + * Translates a Python [`Assign`](https://docs.python.org/3/library/ast.html#ast.Assign) into an + * [AssignExpression]. + */ + private fun handleAssign(node: Python.AST.Assign): AssignExpression { val lhs = node.targets.map { frontend.expressionHandler.handle(it) } + node.type_comment?.let { typeComment -> + val tpe = frontend.typeOf(typeComment) + lhs.forEach { it.type = tpe } + } val rhs = frontend.expressionHandler.handle(node.value) if (rhs is List<*>) newAssignExpression( diff --git a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt index 68a45b3eef..c6f95f3bad 100644 --- a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt +++ b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/statementHandler/StatementHandlerTest.kt @@ -26,22 +26,18 @@ package de.fraunhofer.aisec.cpg.frontends.python.statementHandler import de.fraunhofer.aisec.cpg.TranslationResult -import de.fraunhofer.aisec.cpg.frontends.python.PythonLanguage +import de.fraunhofer.aisec.cpg.frontends.python.* import de.fraunhofer.aisec.cpg.graph.* -import de.fraunhofer.aisec.cpg.graph.statements.AssertStatement -import de.fraunhofer.aisec.cpg.graph.statements.expressions.Literal -import de.fraunhofer.aisec.cpg.test.analyze -import de.fraunhofer.aisec.cpg.test.analyzeAndGetFirstTU -import de.fraunhofer.aisec.cpg.test.assertResolvedType +import de.fraunhofer.aisec.cpg.graph.statements.* +import de.fraunhofer.aisec.cpg.graph.statements.expressions.* +import de.fraunhofer.aisec.cpg.test.* import java.nio.file.Path -import kotlin.test.Test -import kotlin.test.assertEquals -import kotlin.test.assertNotNull +import kotlin.test.* import org.junit.jupiter.api.BeforeAll import org.junit.jupiter.api.TestInstance @TestInstance(TestInstance.Lifecycle.PER_CLASS) -class StatementHandlerTest { +class StatementHandlerTest : BaseTest() { private lateinit var topLevel: Path private lateinit var result: TranslationResult @@ -123,4 +119,19 @@ class StatementHandlerTest { assertNotNull(message, "Assert statement should have a message") assertEquals("Test message", message.value, "The assert message is incorrect") } + + @Test + fun testTypeHints() { + analyzeFile("type_hints.py") + + // type comments + val a = result.refs["a"] + assertNotNull(a) + assertEquals(with(result) { assertResolvedType("int") }, a.type) + + // type annotation + val b = result.refs["b"] + assertNotNull(b) + assertEquals(with(result) { assertResolvedType("str") }, b.type) + } } diff --git a/cpg-language-python/src/test/resources/python/type_hints.py b/cpg-language-python/src/test/resources/python/type_hints.py new file mode 100644 index 0000000000..49032d1caa --- /dev/null +++ b/cpg-language-python/src/test/resources/python/type_hints.py @@ -0,0 +1,3 @@ +a = 1 #type: int + +b: str \ No newline at end of file