Skip to content

Commit

Permalink
Handling imports in Python (#1555)
Browse files Browse the repository at this point in the history
  • Loading branch information
@oxisto
oxisto authored Jun 27, 2024
1 parent b448b1e commit 777dfe2
Show file tree
Hide file tree
Showing 10 changed files with 242 additions and 105 deletions.
68 changes: 40 additions & 28 deletions ...uage-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@
package de.fraunhofer.aisec.cpg.frontends.python

import de.fraunhofer.aisec.cpg.graph.*
import de.fraunhofer.aisec.cpg.graph.declarations.ImportDeclaration
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.MemberExpression
import de.fraunhofer.aisec.cpg.graph.statements.expressions.ProblemExpression
import jep.python.PyObject

Expand Down Expand Up @@ -279,7 +281,21 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
}

private fun handleAttribute(node: Python.ASTAttribute): Expression {
return newMemberExpression(name = node.attr, base = handle(node.value), rawNode = node)
var base = handle(node.value)

// We do a quick check, if this refers to an import. This is faster than doing
// this in a pass and most likely valid, since we are under the assumption that
// our current file is (more or less) complete, but we might miss some
// additional dependencies
var ref =
if (isImport(base.name)) {
// Yes, it's an import, so we need to construct a reference with an FQN
newReference(base.name.fqn(node.attr), rawNode = node)
} else {
newMemberExpression(name = node.attr, base = base, rawNode = node)
}

return ref
}

private fun handleConstant(node: Python.ASTConstant): Expression {
Expand Down Expand Up @@ -327,34 +343,25 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
* TODO: cast, memberexpression, magic
*/
private fun handleCall(node: Python.ASTCall): Expression {
var callee = frontend.expressionHandler.handle(node.func)

val ret =
when (node.func) {
is Python.ASTAttribute -> {
newMemberCallExpression(
frontend.expressionHandler.handle(node.func),
rawNode = node
)
}
else -> {
val func = handle(node.func)

// try to resolve -> [ConstructExpression]
val currentScope = frontend.scopeManager.currentScope
val record =
currentScope?.let { frontend.scopeManager.getRecordForName(func.name) }

if (record != null) {
// construct expression
val constructExpr =
newConstructExpression(
(node.func as? Python.ASTName)?.id,
rawNode = node
)
constructExpr.type = record.toType()
constructExpr
} else {
newCallExpression(func, rawNode = node)
}
if (callee is MemberExpression) {
newMemberCallExpression(callee, rawNode = node)
} else {
// try to resolve -> [ConstructExpression]
val currentScope = frontend.scopeManager.currentScope
val record =
currentScope?.let { frontend.scopeManager.getRecordForName(callee.name) }

if (record != null) {
// construct expression
val constructExpr =
newConstructExpression((node.func as? Python.ASTName)?.id, rawNode = node)
constructExpr.type = record.toType()
constructExpr
} else {
newCallExpression(callee, rawNode = node)
}
}

Expand All @@ -369,6 +376,11 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
return ret
}

private fun isImport(name: Name): Boolean {
val decl = frontend.scopeManager.findSymbols(name).filterIsInstance<ImportDeclaration>()
return decl.isNotEmpty()
}

private fun handleName(node: Python.ASTName): Expression {
val r = newReference(name = node.id, rawNode = node)

Expand Down
2 changes: 1 addition & 1 deletion ...python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonLanguageFrontend.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -387,7 +387,7 @@ fun fromPython(pyObject: Any?): Python.AST {
"ast.Is" -> Python.ASTIs(pyObject)
"ast.IsNot" -> Python.ASTIsNot(pyObject)
"ast.In" -> Python.ASTIn(pyObject)
"ast.NotInt" -> Python.ASTNotIn(pyObject)
"ast.NotIn" -> Python.ASTNotIn(pyObject)

// `"ast.expr_context`
"ast.Load" -> Python.ASTLoad(pyObject)
Expand Down
142 changes: 70 additions & 72 deletions ...guage-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/StatementHandler.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
package de.fraunhofer.aisec.cpg.frontends.python

import de.fraunhofer.aisec.cpg.graph.*
import de.fraunhofer.aisec.cpg.graph.Annotation
import de.fraunhofer.aisec.cpg.graph.declarations.*
import de.fraunhofer.aisec.cpg.graph.statements.DeclarationStatement
import de.fraunhofer.aisec.cpg.graph.statements.Statement
Expand Down Expand Up @@ -76,14 +77,57 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
private fun handleImport(node: Python.ASTImport): Statement {
val declStmt = newDeclarationStatement(rawNode = node)
for (imp in node.names) {
val v =
if (imp.asname != null) {
newVariableDeclaration(imp.asname, rawNode = imp) // TODO refers to original????
val alias = imp.asname
val decl =
if (alias != null) {
newImportDeclaration(
parseName(imp.name),
false,
parseName(alias),
rawNode = imp
)
} else {
newVariableDeclaration(imp.name, rawNode = imp)
newImportDeclaration(parseName(imp.name), false, rawNode = imp)
}
frontend.scopeManager.addDeclaration(v)
declStmt.addDeclaration(v)
frontend.scopeManager.addDeclaration(decl)
declStmt.addToPropertyEdgeDeclaration(decl)
}
return declStmt
}

private fun handleImportFrom(node: Python.ASTImportFrom): Statement {
val declStmt = newDeclarationStatement(rawNode = node)
val level = node.level
if (level == null || level > 0) {
return newProblemExpression(
"not supporting relative paths in from (...) import syntax yet"
)
}

val module = parseName(node.module ?: "")
for (imp in node.names) {
// We need to differentiate between a wildcard import and an individual symbol.
// Wildcards luckily do not have aliases
val decl =
if (imp.name == "*") {
// In the wildcard case, our "import" is the module name and we set "wildcard"
// to true
newImportDeclaration(module, true, rawNode = imp)
} else {
// If we import an individual symbol, we need to FQN the symbol with our module
// name and import that. We also need to take care of any alias
val name = module.fqn(imp.name)
val alias = imp.asname
if (alias != null) {
newImportDeclaration(name, false, parseName(alias), rawNode = imp)
} else {
newImportDeclaration(name, false, rawNode = imp)
}
}

// Finally, add our declaration to the scope and the declaration statement
frontend.scopeManager.addDeclaration(decl)
declStmt.addToPropertyEdgeDeclaration(decl)
}
return declStmt
}
Expand Down Expand Up @@ -187,22 +231,6 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
)
}

private fun handleImportFrom(node: Python.ASTImportFrom): Statement {
val declStmt = newDeclarationStatement(rawNode = node)
for (stmt in node.names) {
val name =
if (stmt.asname != null) {
stmt.asname
} else {
stmt.name
}
val decl = newVariableDeclaration(name = name, rawNode = node)
frontend.scopeManager.addDeclaration(decl)
declStmt.addDeclaration(decl)
}
return declStmt
}

private fun handleClassDef(stmt: Python.ASTClassDef): Statement {
val cls = newRecordDeclaration(stmt.name, "class", rawNode = stmt)
stmt.bases.map { cls.superClasses.add(frontend.typeOf(it)) }
Expand Down Expand Up @@ -456,64 +484,33 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
}
}

private fun handleAnnotations(
node: Python.ASTAsyncFunctionDef
): Collection<de.fraunhofer.aisec.cpg.graph.Annotation> {
val annotations = mutableListOf<de.fraunhofer.aisec.cpg.graph.Annotation>()
for (decorator in node.decorator_list) {
if (decorator !is Python.ASTCall) {
TODO()
}

val decFuncParsed = frontend.expressionHandler.handle(decorator.func)
if (decFuncParsed !is MemberExpression) {
TODO()
}

val annotation =
newAnnotation(
name =
Name(
localName = decFuncParsed.name.localName,
parent = decFuncParsed.base.name
),
rawNode = node
)
for (arg in decorator.args) {
val argParsed = frontend.expressionHandler.handle(arg)
annotation.members +=
newAnnotationMember(
"annotationArg" + decorator.args.indexOf(arg), // TODO
argParsed,
rawNode = arg
)
}
for (keyword in decorator.keywords) {
annotation.members +=
newAnnotationMember(
name = keyword.arg,
value = frontend.expressionHandler.handle(keyword.value),
rawNode = keyword
)
}
private fun handleAnnotations(node: Python.ASTAsyncFunctionDef): Collection<Annotation> {
return handleDeclaratorList(node, node.decorator_list)
}

annotations += annotation
}
return annotations
private fun handleAnnotations(node: Python.ASTFunctionDef): Collection<Annotation> {
return handleDeclaratorList(node, node.decorator_list)
}

private fun handleAnnotations(
node: Python.ASTFunctionDef
): Collection<de.fraunhofer.aisec.cpg.graph.Annotation> {
val annotations = mutableListOf<de.fraunhofer.aisec.cpg.graph.Annotation>()
for (decorator in node.decorator_list) {
fun handleDeclaratorList(
node: Python.AST,
decoratorList: List<Python.ASTBASEexpr>
): List<Annotation> {
val annotations = mutableListOf<Annotation>()
for (decorator in decoratorList) {
if (decorator !is Python.ASTCall) {
TODO()
log.warn(
"Decorator (${decorator::class}) is not ASTCall, cannot handle this (yet)."
)
continue
}

val decFuncParsed = frontend.expressionHandler.handle(decorator.func)
if (decFuncParsed !is MemberExpression) {
TODO()
log.warn(
"parsed function expression (${decFuncParsed::class}) is not a member expression, cannot handle this (yet)."
)
continue
}

val annotation =
Expand Down Expand Up @@ -545,6 +542,7 @@ class StatementHandler(frontend: PythonLanguageFrontend) :

annotations += annotation
}

return annotations
}

Expand Down
6 changes: 3 additions & 3 deletions ...nguage-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,7 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) {
*/
private fun handle(node: Node?) {
when (node) {
// TODO is doubled. Whatever this means
is AssignExpression -> handleAssignExpression(node)
is Reference -> handleReference(node)
is ForEachStatement -> handleForEach(node)
else -> {
// Nothing to do for all other types of nodes
Expand All @@ -87,6 +85,9 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) {
if (node.resolutionHelper is CallExpression) {
return null
}

// TODO(oxisto): Actually the logic here is far more complex in reality, taking into account
// local and global variables, but for now we just resolve it using the scope manager
val resolved = scopeManager.resolveReference(node)

// Nothing to create
Expand Down Expand Up @@ -167,7 +168,6 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) {
handled.let { node.addDeclaration(it) }
}
}
else -> TODO()
}
}
}
Loading

0 comments on commit 777dfe2

Please sign in to comment.