Update sdk mock

koxudaxi · Jun 20, 2021 · 9fc8fe2 · 9fc8fe2
1 parent 3feb2a8
commit 9fc8fe2
Showing 1 changed file with 55 additions and 40 deletions.
diff --git a/testSrc/com/jetbrains/python/fixtures/PythonMockSdk.java b/testSrc/com/jetbrains/python/fixtures/PythonMockSdk.java
@@ -1,4 +1,4 @@
-// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copyright 2000-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.fixtures;
 
 import com.intellij.openapi.projectRoots.Sdk;
@@ -10,7 +10,6 @@
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
-import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.psi.LanguageLevel;
@@ -21,83 +20,99 @@
 import org.jetbrains.annotations.Nullable;
 
 import java.io.File;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
-/**
- * @author yole
- */
-public class PythonMockSdk {
-    @NonNls private static final String MOCK_SDK_NAME = "Mock Python SDK";
 
+public final class PythonMockSdk {
+    public static final String PYTHON_SDK_ID_NAME = "Python SDK";
     private PythonMockSdk() {
     }
 
-    public static Sdk create(final String version, final VirtualFile @NotNull ... additionalRoots) {
-        final String mock_path = PythonTestUtil.getTestDataPath() + "/MockSdk" + version + "/";
+    public static @NotNull Sdk create(@NotNull String name) {
+        return create(name, LanguageLevel.getLatest());
+    }
 
-        String sdkHome = new File(mock_path, "bin/python" + version).getPath();
+    public static @NotNull Sdk create(@NotNull LanguageLevel level, VirtualFile @NotNull ... additionalRoots) {
+        return create("MockSdk", level, additionalRoots);
+    }
+
+    private static @NotNull Sdk create(@NotNull String name, @NotNull LanguageLevel level, VirtualFile @NotNull ... additionalRoots) {
+        return create(name, new PyMockSdkType(level), level, additionalRoots);
+    }
+
+    public static @NotNull Sdk create(@NotNull String pathSuffix, @NotNull SdkTypeId sdkType, @NotNull LanguageLevel level, VirtualFile @NotNull ... additionalRoots) {
+        String sdkName = "Mock " + PYTHON_SDK_ID_NAME + " " + level.toPythonVersion();
+        return create(sdkName, pathSuffix, sdkType, level, additionalRoots);
+    }
+
+    public static @NotNull Sdk create(@NotNull String name, @NotNull String pathSuffix, @NotNull SdkTypeId sdkType, @NotNull LanguageLevel level, VirtualFile @NotNull ... additionalRoots) {
+        final String mockSdkPath = PythonTestUtil.getTestDataPath() + "/" + pathSuffix;
 
         MultiMap<OrderRootType, VirtualFile> roots = MultiMap.create();
-        OrderRootType classes = OrderRootType.CLASSES;
+        roots.putValues(OrderRootType.CLASSES, createRoots(mockSdkPath, level));
+        roots.putValues(OrderRootType.CLASSES, Arrays.asList(additionalRoots));
 
-        ContainerUtil.putIfNotNull(classes, LocalFileSystem.getInstance().refreshAndFindFileByIoFile(new File(mock_path, "Lib")), roots);
+        MockSdk sdk = new MockSdk(
+                name,
+                mockSdkPath + "/bin/python",
+                toVersionString(level),
+                roots,
+                sdkType
+        );
+
+        // com.jetbrains.python.psi.resolve.PythonSdkPathCache.getInstance() corrupts SDK, so have to clone
+        return sdk.clone();
+    }
 
-        ContainerUtil.putIfNotNull(classes, PyUserSkeletonsUtil.getUserSkeletonsDirectory(), roots);
+    private static @NotNull List<VirtualFile> createRoots(@NotNull @NonNls String mockSdkPath, @NotNull LanguageLevel level) {
+        final var result = new ArrayList<VirtualFile>();
 
-        final LanguageLevel level = LanguageLevel.fromPythonVersion(version);
-        final VirtualFile typeShedDir = PyTypeShed.INSTANCE.getDirectory();
-        PyTypeShed.INSTANCE
-                .findRootsForLanguageLevel(level)
-                .forEach(path -> ContainerUtil.putIfNotNull(classes, typeShedDir.findFileByRelativePath(path), roots));
+        final var localFS = LocalFileSystem.getInstance();
+        ContainerUtil.addIfNotNull(result, localFS.refreshAndFindFileByIoFile(new File(mockSdkPath, "Lib")));
+        ContainerUtil.addIfNotNull(result, localFS.refreshAndFindFileByIoFile(new File(mockSdkPath, PythonSdkUtil.SKELETON_DIR_NAME)));
 
-        String mock_stubs_path = mock_path + PythonSdkUtil.SKELETON_DIR_NAME;
-        ContainerUtil.putIfNotNull(classes, LocalFileSystem.getInstance().refreshAndFindFileByPath(mock_stubs_path), roots);
+        ContainerUtil.addIfNotNull(result, PyUserSkeletonsUtil.getUserSkeletonsDirectory());
 
-        roots.putValues(classes, Arrays.asList(additionalRoots));
+        result.addAll(PyTypeShed.INSTANCE.findRootsForLanguageLevel(level));
 
-        MockSdk sdk = new MockSdk(MOCK_SDK_NAME + " " + version, sdkHome, "Python " + version + " Mock SDK", roots, new PyMockSdkType(version));
+        return result;
+    }
 
-        // com.jetbrains.python.psi.resolve.PythonSdkPathCache.getInstance() corrupts SDK, so have to clone
-        return sdk.clone();
+    private static @NotNull String toVersionString(@NotNull LanguageLevel level) {
+        return "Python " + level.toPythonVersion();
     }
 
-    private static class PyMockSdkType implements SdkTypeId {
+    private static final class PyMockSdkType implements SdkTypeId {
 
         @NotNull
-        private final String myVersionString;
-        private final String mySdkIdName;
+        private final LanguageLevel myLevel;
 
-        private PyMockSdkType(@NotNull String string) {
-            mySdkIdName = PyNames.PYTHON_SDK_ID_NAME;
-            myVersionString = string;
+        private PyMockSdkType(@NotNull LanguageLevel level) {
+            myLevel = level;
         }
 
         @NotNull
         @Override
         public String getName() {
-            return mySdkIdName;
+            return PYTHON_SDK_ID_NAME;
         }
 
         @Nullable
         @Override
-        public String getVersionString(@NotNull Sdk sdk) {
-            return myVersionString;
+        public @NotNull String getVersionString(@NotNull Sdk sdk) {
+            return toVersionString(myLevel);
         }
 
         @Override
         public void saveAdditionalData(@NotNull SdkAdditionalData additionalData, @NotNull Element additional) {
-
         }
 
         @Nullable
         @Override
         public SdkAdditionalData loadAdditionalData(@NotNull Sdk currentSdk, @NotNull Element additional) {
             return null;
         }
-
-        @Override
-        public boolean isLocalSdk(@NotNull Sdk sdk) {
-            return false;
-        }
     }
-}
+}