Support for python packages (#1779)

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

@@ -277,13 +277,22 @@ class PythonLanguageFrontend(language: Language<PythonLanguageFrontend>, ctx: Tr
             modulePaths.fold(null) { previous: NamespaceDeclaration?, path ->
                 var fqn = previous?.name.fqn(path)
 
-                if (path != "__init__") {
+                // The __init__ module is very special in Python. The symbols that are declared by
+                // __init__.py are available directly under the path of the package (not module) it
+                // lies in. For example, if the contents of the file foo/bar/__init__.py are
+                // available in the module foo.bar (under the assumption that both foo and bar are
+                // packages). We therefore do not want to create an additional __init__ namespace.
+                // However, in reality, the symbols are actually available in foo.bar as well as in
+                // foo.bar.__init__, although the latter is practically not used, and therefore we
+                // do not support it because major workarounds would be needed.
+                if (path == "__init__") {
+                    previous
+                } else {
                     val nsd = newNamespaceDeclaration(fqn, rawNode = pythonASTModule)
+                    nsd.path = relative?.parent?.pathString + "/" + module
                     scopeManager.addDeclaration(nsd)
                     scopeManager.enterScope(nsd)
                     nsd
-                } else {
-                    previous
                 }
             }
 

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

@@ -36,6 +36,12 @@ import de.fraunhofer.aisec.cpg.graph.declarations.*
 import de.fraunhofer.aisec.cpg.graph.scopes.BlockScope
 import de.fraunhofer.aisec.cpg.graph.scopes.NameScope
 import de.fraunhofer.aisec.cpg.graph.statements.*
+import de.fraunhofer.aisec.cpg.graph.statements.AssertStatement
+import de.fraunhofer.aisec.cpg.graph.statements.CatchClause
+import de.fraunhofer.aisec.cpg.graph.statements.DeclarationStatement
+import de.fraunhofer.aisec.cpg.graph.statements.ForEachStatement
+import de.fraunhofer.aisec.cpg.graph.statements.Statement
+import de.fraunhofer.aisec.cpg.graph.statements.TryStatement
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
 import de.fraunhofer.aisec.cpg.graph.types.FunctionType
 import de.fraunhofer.aisec.cpg.helpers.Util
@@ -419,13 +425,31 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
     private fun handleImportFrom(node: Python.AST.ImportFrom): 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"
-            )
+        var module = parseName(node.module ?: "")
+
+        if (level != null && level > 0L) {
+            // Because the __init__ module is omitted from our current namespace, we need to check
+            // for its existence and add __init__, otherwise the relative path would be off by one
+            // level.
+            var parent =
+                if (isInitModule()) {
+                    frontend.scopeManager.currentNamespace.fqn("__init__")
+                } else {
+                    frontend.scopeManager.currentNamespace
+                }
+
+            // If the level is specified, we need to relative the module path. We basically need to
+            // move upwards in the parent namespace in the amount of dots
+            for (i in 0 until level) {
+                parent = parent?.parent
+                if (parent == null) {
+                    break
+                }
+            }
+
+            module = parent.fqn(module.localName)
         }
 
-        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
@@ -453,6 +477,13 @@ class StatementHandler(frontend: PythonLanguageFrontend) :
         return declStmt
     }
 
+    /** Small utility function to check, whether we are inside an __init__ module. */
+    private fun isInitModule(): Boolean =
+        (frontend.scopeManager.firstScopeIsInstanceOrNull<NameScope>()?.astNode
+                as? NamespaceDeclaration)
+            ?.path
+            ?.endsWith("__init__") == true
+
     private fun handleWhile(node: Python.AST.While): Statement {
         val ret = newWhileStatement(rawNode = node)
         ret.condition = frontend.expressionHandler.handle(node.test)

cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt

@@ -38,11 +38,10 @@ import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
 import de.fraunhofer.aisec.cpg.graph.types.InitializerTypePropagation
 import de.fraunhofer.aisec.cpg.helpers.SubgraphWalker
-import de.fraunhofer.aisec.cpg.passes.configuration.DependsOn
 import de.fraunhofer.aisec.cpg.passes.configuration.ExecuteBefore
 import de.fraunhofer.aisec.cpg.passes.configuration.RequiredFrontend
 
-@DependsOn(TypeResolver::class)
+@ExecuteBefore(ImportResolver::class)
 @ExecuteBefore(SymbolResolver::class)
 @RequiredFrontend(PythonLanguageFrontend::class)
 class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx) {

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

@@ -25,6 +25,7 @@
  */
 package de.fraunhofer.aisec.cpg.frontends.python
 
+import de.fraunhofer.aisec.cpg.InferenceConfiguration
 import de.fraunhofer.aisec.cpg.analysis.ValueEvaluator
 import de.fraunhofer.aisec.cpg.graph.*
 import de.fraunhofer.aisec.cpg.graph.Annotation
@@ -1486,6 +1487,61 @@ class PythonFrontendTest : BaseTest() {
         assertEquals("e = \"é\"", unicodeFunc.body?.code)
     }
 
+    @Test
+    fun testPackageResolution() {
+        val topLevel = Path.of("src", "test", "resources", "python", "packages")
+        var result =
+            analyze(listOf(topLevel.resolve("foobar")).map { it.toFile() }, topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+                it.useParallelFrontends(false)
+                it.failOnError(false)
+                it.inferenceConfiguration(
+                    InferenceConfiguration.builder().inferFunctions(false).build()
+                )
+            }
+        assertNotNull(result)
+
+        var expected =
+            setOf(
+                "foobar",
+                "foobar.__main__",
+                "foobar.module1",
+                "foobar.config",
+                "foobar.implementation",
+                "foobar.implementation.internal_bar",
+                "foobar.implementation.internal_foo"
+            )
+        assertEquals(expected, result.namespaces.map { it.name.toString() }.distinct().toSet())
+
+        var bar = result.functions["bar"]
+        assertNotNull(bar)
+        assertFullName("foobar.implementation.internal_bar.bar", bar)
+
+        var foo = result.functions["foo"]
+        assertNotNull(foo)
+        assertFullName("foobar.implementation.internal_foo.foo", foo)
+
+        var barCall = result.calls["bar"]
+        assertNotNull(barCall)
+        assertInvokes(barCall, bar)
+
+        var fooCalls = result.calls("foo")
+        assertEquals(2, fooCalls.size)
+        fooCalls.forEach { assertInvokes(it, foo) }
+
+        val refBarString = result.refs("bar_string")
+        refBarString.forEach {
+            assertNotNull(it)
+            assertNotNull(it.refersTo)
+        }
+
+        val refFooString = result.refs("foo_string")
+        refFooString.forEach {
+            assertNotNull(it)
+            assertNotNull(it.refersTo)
+        }
+    }
+
     class PythonValueEvaluator : ValueEvaluator() {
         override fun computeBinaryOpEffect(
             lhsValue: Any?,

cpg-language-python/src/test/resources/python/packages/foobar/__main__.py

@@ -0,0 +1,12 @@
+# module foobar.implementation exports "foo" from foobar.implementation.internal_foo,
+# so we can directly import it
+from foobar.implementation import foo
+
+# bar is not exported in the package, so we need to import the specific module to
+# access it
+from foobar.implementation.internal_bar import bar
+
+bar()
+foo()
+
+print("Main!")

cpg-language-python/src/test/resources/python/packages/foobar/config/__init__.py

@@ -0,0 +1,2 @@
+bar_string = "bar"
+foo_string = "foo"

cpg-language-python/src/test/resources/python/packages/foobar/implementation/__init__.py

@@ -0,0 +1,4 @@
+# this also "exports" foo in this module
+from .internal_foo import foo
+
+foo()

cpg-language-python/src/test/resources/python/packages/foobar/implementation/internal_bar.py

@@ -0,0 +1,5 @@
+import foobar.config
+
+def bar():
+    print(foobar.config.bar_string)
+    pass

cpg-language-python/src/test/resources/python/packages/foobar/implementation/internal_foo.py

@@ -0,0 +1,5 @@
+from ..config import foo_string
+
+def foo():
+    print(foo_string)
+    pass