SONARPY-2295 Support relative imports in type inference V2 (#2122)

SonarSource · Nov 1, 2024 · e6ec935 · e6ec935
1 parent 4833ba3
commit e6ec935
Show file tree

Hide file tree

Showing 10 changed files with 105 additions and 25 deletions.
diff --git a/python-frontend/src/main/java/org/sonar/plugins/python/api/PythonVisitorContext.java b/python-frontend/src/main/java/org/sonar/plugins/python/api/PythonVisitorContext.java
@@ -43,15 +43,15 @@ public class PythonVisitorContext extends PythonInputFileContext {
   private List<PreciseIssue> issues = new ArrayList<>();
   private final TypeChecker typeChecker;
 
-  public PythonVisitorContext(FileInput rootTree, PythonFile pythonFile, @Nullable File workingDirectory, @Nullable String packageName) {
+  public PythonVisitorContext(FileInput rootTree, PythonFile pythonFile, @Nullable File workingDirectory, String packageName) {
     super(pythonFile, workingDirectory, CacheContextImpl.dummyCache());
     this.rootTree = rootTree;
     this.parsingException = null;
-    SymbolTableBuilder symbolTableBuilder = packageName != null ? new SymbolTableBuilder(packageName, pythonFile) : new SymbolTableBuilder(pythonFile);
+    SymbolTableBuilder symbolTableBuilder = new SymbolTableBuilder(packageName, pythonFile);
     symbolTableBuilder.visitFileInput(rootTree);
     var symbolTable = new SymbolTableBuilderV2(rootTree).build();
     var projectLevelTypeTable = new ProjectLevelTypeTable(ProjectLevelSymbolTable.empty());
-    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable).inferTypes(rootTree);
+    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable, packageName).inferTypes(rootTree);
     this.typeChecker = new TypeChecker(projectLevelTypeTable);
   }
 
@@ -65,7 +65,7 @@ public PythonVisitorContext(FileInput rootTree, PythonFile pythonFile, @Nullable
     var symbolTable = new SymbolTableBuilderV2(rootTree)
       .build();
     var projectLevelTypeTable = new ProjectLevelTypeTable(projectLevelSymbolTable);
-    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable).inferTypes(rootTree);
+    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable, packageName).inferTypes(rootTree);
     this.typeChecker = new TypeChecker(projectLevelTypeTable);
   }
 
@@ -78,7 +78,7 @@ public PythonVisitorContext(FileInput rootTree, PythonFile pythonFile, @Nullable
     var symbolTable = new SymbolTableBuilderV2(rootTree)
       .build();
     var projectLevelTypeTable = new ProjectLevelTypeTable(projectLevelSymbolTable);
-    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable).inferTypes(rootTree);
+    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable, packageName).inferTypes(rootTree);
     this.typeChecker = new TypeChecker(projectLevelTypeTable);
   }
 

diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/TypeInferenceV2.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/TypeInferenceV2.java
@@ -36,6 +36,7 @@
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.StatementList;
 import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.python.semantic.SymbolUtils;
 import org.sonar.python.semantic.v2.types.FlowSensitiveTypeInference;
 import org.sonar.python.semantic.v2.types.Propagation;
 import org.sonar.python.semantic.v2.types.PropagationVisitor;
@@ -49,15 +50,17 @@ public class TypeInferenceV2 {
   private final TypeTable projectLevelTypeTable;
   private final SymbolTable symbolTable;
   private final PythonFile pythonFile;
+  private final String fullyQualifiedModuleName;
 
-  public TypeInferenceV2(TypeTable projectLevelTypeTable, PythonFile pythonFile, SymbolTable symbolTable) {
+  public TypeInferenceV2(TypeTable projectLevelTypeTable, PythonFile pythonFile, SymbolTable symbolTable, String packageName) {
     this.projectLevelTypeTable = projectLevelTypeTable;
     this.symbolTable = symbolTable;
     this.pythonFile = pythonFile;
+    this.fullyQualifiedModuleName = SymbolUtils.fullyQualifiedModuleName(packageName, pythonFile.fileName());
   }
 
   public Map<SymbolV2, Set<PythonType>> inferTypes(FileInput fileInput) {
-    TrivialTypeInferenceVisitor trivialTypeInferenceVisitor = new TrivialTypeInferenceVisitor(projectLevelTypeTable, pythonFile);
+    TrivialTypeInferenceVisitor trivialTypeInferenceVisitor = new TrivialTypeInferenceVisitor(projectLevelTypeTable, pythonFile, fullyQualifiedModuleName);
     fileInput.accept(trivialTypeInferenceVisitor);
 
     var typesBySymbol = inferTypesAndMemberAccessSymbols(fileInput);

diff --git a/...rontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java b/...rontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java
@@ -26,6 +26,7 @@
 import java.util.Deque;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import org.sonar.plugins.python.api.PythonFile;
 import org.sonar.plugins.python.api.tree.AliasedName;
@@ -96,13 +97,15 @@ public class TrivialTypeInferenceVisitor extends BaseTreeVisitor {
 
   private final TypeTable projectLevelTypeTable;
   private final String fileId;
+  private final String fullyQualifiedModuleName;
 
   private final Deque<PythonType> typeStack = new ArrayDeque<>();
 
-  public TrivialTypeInferenceVisitor(TypeTable projectLevelTypeTable, PythonFile pythonFile) {
+  public TrivialTypeInferenceVisitor(TypeTable projectLevelTypeTable, PythonFile pythonFile, String fullyQualifiedModuleName) {
     this.projectLevelTypeTable = projectLevelTypeTable;
     Path path = pathOf(pythonFile);
     this.fileId = path != null ? path.toString() : pythonFile.toString();
+    this.fullyQualifiedModuleName = fullyQualifiedModuleName;
   }
 
 
@@ -381,10 +384,17 @@ private static void generateNames(PythonType resolvedType, List<Name> names, Lis
 
   @Override
   public void visitImportFrom(ImportFrom importFrom) {
-    Optional.of(importFrom)
-      .map(ImportFrom::module)
+    List<String> fromModuleFqn = Optional.ofNullable(importFrom.module())
       .map(TrivialTypeInferenceVisitor::dottedNameToPartFqn)
-      .ifPresent(fqn -> setTypeToImportFromStatement(importFrom, fqn));
+      .orElse(new ArrayList<>());
+    List<Token> dotPrefixTokens = importFrom.dottedPrefixForModule();
+    if (!dotPrefixTokens.isEmpty()) {
+      // Relative import: we start from the current module FQN and go up as many levels as there are dots in the import statement
+      List<String> moduleFqnElements = List.of(fullyQualifiedModuleName.split("\\."));
+      int sizeLimit = Math.max(0, moduleFqnElements.size() - dotPrefixTokens.size());
+      fromModuleFqn = Stream.concat(moduleFqnElements.stream().limit(sizeLimit), fromModuleFqn.stream()).toList();
+    }
+    setTypeToImportFromStatement(importFrom, fromModuleFqn);
   }
 
   private static List<String> dottedNameToPartFqn(DottedName dottedName) {
@@ -413,7 +423,7 @@ private void setTypeToImportFromStatement(ImportFrom importFrom, List<String> fq
 
   private static UnknownType.UnresolvedImportType createUnresolvedImportType(List<String> moduleFqnList, Name name) {
     String fromModuleFqn = String.join(".", moduleFqnList);
-    String fqn = fromModuleFqn + "." + name.name();
+    String fqn = fromModuleFqn.isEmpty() ? name.name() : String.join(".", fromModuleFqn, name.name());
     return new UnknownType.UnresolvedImportType(fqn);
   }
 

diff --git a/python-frontend/src/test/java/org/sonar/python/PythonTestUtils.java b/python-frontend/src/test/java/org/sonar/python/PythonTestUtils.java
@@ -59,7 +59,7 @@ public static String appendNewLine(String s) {
   }
 
   public static FileInput parse(String... lines) {
-    return parse(new SymbolTableBuilder(pythonFile("")), lines);
+    return parse(new SymbolTableBuilder(pythonFile("mod")), lines);
   }
 
   public static FileInput parse(SymbolTableBuilder symbolTableBuilder, String... lines) {

diff --git a/python-frontend/src/test/java/org/sonar/python/SubscriptionVisitorTest.java b/python-frontend/src/test/java/org/sonar/python/SubscriptionVisitorTest.java
@@ -62,7 +62,7 @@ public void initialize(Context context) {
     };
 
     FileInput fileInput = PythonTestUtils.parse("'.*'");
-    PythonVisitorContext context = new PythonVisitorContext(fileInput, PythonTestUtils.pythonFile("file"), null, null);
+    PythonVisitorContext context = new PythonVisitorContext(fileInput, PythonTestUtils.pythonFile("file"), null, "");
     SubscriptionVisitor.analyze(Collections.singleton(check), context);
   }
 
@@ -100,7 +100,7 @@ public void initialize(Context context) {
     };
 
     FileInput fileInput = PythonTestUtils.parse("class A:\n  def foo(self): ...");
-    PythonVisitorContext context = new PythonVisitorContext(fileInput, PythonTestUtils.pythonFile("file"), null, null);
+    PythonVisitorContext context = new PythonVisitorContext(fileInput, PythonTestUtils.pythonFile("file"), null, "");
     SubscriptionVisitor.analyze(Collections.singleton(check), context);
   }
 }
diff --git a/python-frontend/src/test/java/org/sonar/python/semantic/v2/ProjectLevelTypeTableTest.java b/python-frontend/src/test/java/org/sonar/python/semantic/v2/ProjectLevelTypeTableTest.java
@@ -321,4 +321,35 @@ def foo(): ...
     var typeWrapper = (LazyTypeWrapper) fooType.decorators().get(0);
     assertThat(typeWrapper.hasImportPath("lib2.lib.lib_decorator")).isTrue();
   }
+
+  @Test
+  void relativeImports() {
+    var projectLevelSymbolTable = new ProjectLevelSymbolTable();
+    FileInput initTree = parseWithoutSymbols("");
+    PythonFile initFile = pythonFile("__init__.py");
+    projectLevelSymbolTable.addModule(initTree, "my_package", initFile);
+
+    var libTree = parseWithoutSymbols(
+      """
+      def foo(): ...
+      """
+    );
+    projectLevelSymbolTable.addModule(libTree, "my_package", pythonFile("lib.py"));
+
+    var projectLevelTypeTable = new ProjectLevelTypeTable(projectLevelSymbolTable);
+    var mainFile = pythonFile("main.py");
+    var fileInput = parseAndInferTypes(projectLevelTypeTable, mainFile, """
+      from .lib import foo
+      from . import lib
+      foo
+      lib
+      """
+    );
+    PythonType fooType = ((ExpressionStatement) fileInput.statements().statements().get(2)).expressions().get(0).typeV2();
+    assertThat(fooType).isInstanceOf(FunctionType.class);
+    assertThat(fooType.name()).isEqualTo("foo");
+    PythonType libType = ((ExpressionStatement) fileInput.statements().statements().get(3)).expressions().get(0).typeV2();
+    assertThat(libType).isInstanceOf(ModuleType.class);
+    assertThat(libType.name()).isEqualTo("lib");
+  }
 }
diff --git a/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java b/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java
@@ -97,7 +97,7 @@
 
 public class TypeInferenceV2Test {
 
-  static PythonFile pythonFile = PythonTestUtils.pythonFile("");
+  static PythonFile pythonFile = PythonTestUtils.pythonFile("mod");
 
   @Test
   void testTypeshedImports() {
@@ -191,6 +191,42 @@ def function():
     assertThat(argType).isInstanceOfSatisfying(UnresolvedImportType.class, a -> assertThat(a.importPath()).isEqualTo("a.b"));
   }
 
+  @Test
+  void testRelativeImport() {
+    FileInput fileInput = inferTypes("""
+      from . import module
+      module
+      """);
+
+    PythonType pythonType = ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(pythonType).isInstanceOf(UnresolvedImportType.class);
+    assertThat(((UnresolvedImportType) pythonType).importPath()).isEqualTo("my_package.module");
+
+    fileInput = inferTypes("""
+      from .. import module
+      module
+      """);
+    pythonType = ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(pythonType).isInstanceOf(UnresolvedImportType.class);
+    assertThat(((UnresolvedImportType) pythonType).importPath()).isEqualTo("module");
+
+    fileInput = inferTypes("""
+      from .hello import module
+      module
+      """);
+    pythonType = ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(pythonType).isInstanceOf(UnresolvedImportType.class);
+    assertThat(((UnresolvedImportType) pythonType).importPath()).isEqualTo("my_package.hello.module");
+
+    fileInput = inferTypes("""
+      from ..second_hello import module
+      module
+      """);
+    pythonType = ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(pythonType).isInstanceOf(UnresolvedImportType.class);
+    assertThat(((UnresolvedImportType) pythonType).importPath()).isEqualTo("second_hello.module");
+  }
+
   @Test
   void testProjectLevelSymbolTableImports() {
     var classSymbol = new ClassSymbolImpl("C", "something.known.C");
@@ -2987,7 +3023,7 @@ def foo(): ...
   private static Map<SymbolV2, Set<PythonType>> inferTypesBySymbol(String lines) {
     FileInput root = parse(lines);
     var symbolTable = new SymbolTableBuilderV2(root).build();
-    var typeInferenceV2 = new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable);
+    var typeInferenceV2 = new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "");
     return typeInferenceV2.inferTypes(root);
   }
 
@@ -3000,7 +3036,7 @@ private static FileInput inferTypes(String lines, ProjectLevelTypeTable projectL
 
     var symbolTable = new SymbolTableBuilderV2(root)
       .build();
-    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable).inferTypes(root);
+    new TypeInferenceV2(projectLevelTypeTable, pythonFile, symbolTable, "my_package").inferTypes(root);
     return root;
   }
 

diff --git a/python-frontend/src/test/java/org/sonar/python/types/v2/ClassTypeTest.java b/python-frontend/src/test/java/org/sonar/python/types/v2/ClassTypeTest.java
@@ -173,7 +173,7 @@ void multiple_bindings() {
       "C = \"hello\"");
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
 
     ClassDef classDef = (ClassDef) fileInput.statements().statements().get(0);
     PythonType pythonType = classDef.name().typeV2();
@@ -197,7 +197,7 @@ void multiple_bindings_2() {
     );
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
 
     ClassDef classDef = (ClassDef) fileInput.statements().statements().get(1);
     PythonType pythonType = classDef.name().typeV2();
@@ -564,7 +564,7 @@ void type_annotations_scope() {
     );
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
     ClassDef firstDef = (ClassDef) fileInput.statements().statements().get(0);
     ClassDef innerClass = (ClassDef) firstDef.body().statements().get(0);
     FunctionDef functionDef = (FunctionDef) firstDef.body().statements().get(1);
@@ -613,7 +613,7 @@ public static List<ClassType> classTypes(String... code) {
     FileInput fileInput = parseWithoutSymbols(code);
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
     return PythonTestUtils.getAllDescendant(fileInput, t -> t.is(Tree.Kind.CLASSDEF))
       .stream()
       .map(ClassDef.class::cast)

diff --git a/python-frontend/src/test/java/org/sonar/python/types/v2/FunctionTypeTest.java b/python-frontend/src/test/java/org/sonar/python/types/v2/FunctionTypeTest.java
@@ -175,7 +175,7 @@ void owner() {
     );
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
 
     ClassDef classDef = (ClassDef) PythonTestUtils.getAllDescendant(fileInput, t -> t.is(Tree.Kind.CLASSDEF)).get(0);
     FunctionDef functionDef = (FunctionDef) PythonTestUtils.getAllDescendant(fileInput, t -> t.is(Tree.Kind.FUNCDEF)).get(0);
@@ -201,7 +201,7 @@ public static FunctionType functionType(String... code) {
     FileInput fileInput = parseWithoutSymbols(code);
     var symbolTable = new SymbolTableBuilderV2(fileInput)
       .build();
-    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(PROJECT_LEVEL_TYPE_TABLE, pythonFile, symbolTable, "").inferTypes(fileInput);
     FunctionDef functionDef = (FunctionDef) PythonTestUtils.getAllDescendant(fileInput, t -> t.is(Tree.Kind.FUNCDEF)).get(0);
     return (FunctionType) functionDef.name().typeV2();
   }

diff --git a/python-frontend/src/test/java/org/sonar/python/types/v2/TypesTestUtils.java b/python-frontend/src/test/java/org/sonar/python/types/v2/TypesTestUtils.java
@@ -56,7 +56,7 @@ public static FileInput parseAndInferTypes(PythonFile pythonFile, String... code
   public static FileInput parseAndInferTypes(ProjectLevelTypeTable typeTable, PythonFile pythonFile, String... code) {
     FileInput fileInput = PythonTestUtils.parseWithoutSymbols(code);
     var symbolTable = new SymbolTableBuilderV2(fileInput).build();
-    new TypeInferenceV2(typeTable, pythonFile, symbolTable).inferTypes(fileInput);
+    new TypeInferenceV2(typeTable, pythonFile, symbolTable, "my_package").inferTypes(fileInput);
     return fileInput;
   }
 }