9743ab0280c0ece58d5163a86dfb0c0c553d9999
[idea/community.git] / python / src / com / jetbrains / python / packaging / PyPackageUtil.java
1 // Copyright 2000-2020 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.
2 package com.jetbrains.python.packaging;
3
4 import com.google.common.collect.Sets;
5 import com.intellij.execution.ExecutionException;
6 import com.intellij.openapi.Disposable;
7 import com.intellij.openapi.application.Application;
8 import com.intellij.openapi.application.ApplicationManager;
9 import com.intellij.openapi.application.PathManager;
10 import com.intellij.openapi.application.ReadAction;
11 import com.intellij.openapi.diagnostic.Logger;
12 import com.intellij.openapi.editor.Document;
13 import com.intellij.openapi.fileEditor.FileDocumentManager;
14 import com.intellij.openapi.module.Module;
15 import com.intellij.openapi.project.Project;
16 import com.intellij.openapi.projectRoots.Sdk;
17 import com.intellij.openapi.roots.ModuleRootManager;
18 import com.intellij.openapi.roots.OrderRootType;
19 import com.intellij.openapi.roots.ProjectFileIndex;
20 import com.intellij.openapi.roots.ProjectRootManager;
21 import com.intellij.openapi.util.Pair;
22 import com.intellij.openapi.util.Ref;
23 import com.intellij.openapi.util.text.StringUtil;
24 import com.intellij.openapi.vfs.*;
25 import com.intellij.openapi.vfs.newvfs.events.VFileContentChangeEvent;
26 import com.intellij.openapi.vfs.newvfs.events.VFileCreateEvent;
27 import com.intellij.openapi.vfs.newvfs.events.VFileEvent;
28 import com.intellij.openapi.vfs.newvfs.events.VFilePropertyChangeEvent;
29 import com.intellij.psi.PsiElement;
30 import com.intellij.psi.PsiFile;
31 import com.intellij.psi.PsiManager;
32 import com.intellij.psi.ResolveResult;
33 import com.jetbrains.python.PyBundle;
34 import com.jetbrains.python.PyNames;
35 import com.jetbrains.python.PyPsiPackageUtil;
36 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
37 import com.jetbrains.python.codeInsight.typing.PyTypeShed;
38 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
39 import com.jetbrains.python.psi.*;
40 import com.jetbrains.python.psi.impl.PyPsiUtils;
41 import com.jetbrains.python.psi.resolve.PyResolveContext;
42 import com.jetbrains.python.psi.types.TypeEvalContext;
43 import com.jetbrains.python.remote.PyCredentialsContribution;
44 import com.jetbrains.python.sdk.CredentialsTypeExChecker;
45 import com.jetbrains.python.sdk.PythonSdkUtil;
46 import one.util.streamex.StreamEx;
47 import org.jetbrains.annotations.ApiStatus;
48 import org.jetbrains.annotations.NotNull;
49 import org.jetbrains.annotations.Nullable;
50
51 import java.util.*;
52 import java.util.concurrent.atomic.AtomicBoolean;
53 import java.util.stream.Collectors;
54 import java.util.stream.Stream;
55
56 /**
57  * @author vlan
58  */
59 public final class PyPackageUtil {
60   public static final String SETUPTOOLS = "setuptools";
61   public static final String PIP = "pip";
62   public static final String DISTRIBUTE = "distribute";
63   private static final Logger LOG = Logger.getInstance(PyPackageUtil.class);
64
65   private static class InterpreterChangeEvents {
66     private static final Logger LOG = Logger.getInstance(InterpreterChangeEvents.class);
67   }
68
69   @NotNull
70   private static final String REQUIRES = "requires";
71
72   @NotNull
73   private static final String INSTALL_REQUIRES = "install_requires";
74
75   private static final String @NotNull [] SETUP_PY_REQUIRES_KWARGS_NAMES = new String[]{
76     REQUIRES, INSTALL_REQUIRES, "setup_requires", "tests_require"
77   };
78
79   @NotNull
80   private static final String DEPENDENCY_LINKS = "dependency_links";
81
82   private PyPackageUtil() {
83   }
84
85   public static boolean hasSetupPy(@NotNull Module module) {
86     return findSetupPy(module) != null;
87   }
88
89   @Nullable
90   public static PyFile findSetupPy(@NotNull Module module) {
91     for (VirtualFile root : PyUtil.getSourceRoots(module)) {
92       final VirtualFile child = root.findChild("setup.py");
93       if (child != null) {
94         final PsiFile file = ReadAction.compute(() -> PsiManager.getInstance(module.getProject()).findFile(child));
95         if (file instanceof PyFile) {
96           return (PyFile)file;
97         }
98       }
99     }
100     return null;
101   }
102
103   public static boolean hasRequirementsTxt(@NotNull Module module) {
104     return findRequirementsTxt(module) != null;
105   }
106
107   @Nullable
108   public static VirtualFile findRequirementsTxt(@NotNull Module module) {
109     final String requirementsPath = PyPackageRequirementsSettings.getInstance(module).getRequirementsPath();
110     if (!requirementsPath.isEmpty()) {
111       final VirtualFile file = LocalFileSystem.getInstance().findFileByPath(requirementsPath);
112       if (file != null) {
113         return file;
114       }
115       final ModuleRootManager manager = ModuleRootManager.getInstance(module);
116       for (VirtualFile root : manager.getContentRoots()) {
117         final VirtualFile fileInRoot = root.findFileByRelativePath(requirementsPath);
118         if (fileInRoot != null) {
119           return fileInRoot;
120         }
121       }
122     }
123     return null;
124   }
125
126   @Nullable
127   private static PsiElement findSetupPyInstallRequires(@Nullable PyCallExpression setupCall) {
128     if (setupCall == null) return null;
129
130     return StreamEx
131       .of(REQUIRES, INSTALL_REQUIRES)
132       .map(setupCall::getKeywordArgument)
133       .map(PyPackageUtil::resolveValue)
134       .findFirst(Objects::nonNull)
135       .orElse(null);
136   }
137
138   @Nullable
139   public static List<PyRequirement> findSetupPyRequires(@NotNull Module module) {
140     final PyCallExpression setupCall = findSetupCall(module);
141     if (setupCall == null) return null;
142
143     final List<PyRequirement> requirementsFromRequires = getSetupPyRequiresFromArguments(setupCall, SETUP_PY_REQUIRES_KWARGS_NAMES);
144     final List<PyRequirement> requirementsFromLinks = getSetupPyRequiresFromArguments(setupCall, DEPENDENCY_LINKS);
145
146     return mergeSetupPyRequirements(requirementsFromRequires, requirementsFromLinks);
147   }
148
149   @Nullable
150   public static Map<String, List<PyRequirement>> findSetupPyExtrasRequire(@NotNull Module module) {
151     final PyCallExpression setupCall = findSetupCall(module);
152     if (setupCall == null) return null;
153
154     final PyDictLiteralExpression extrasRequire =
155       PyUtil.as(resolveValue(setupCall.getKeywordArgument("extras_require")), PyDictLiteralExpression.class);
156     if (extrasRequire == null) return null;
157
158     final Map<String, List<PyRequirement>> result = new HashMap<>();
159
160     for (PyKeyValueExpression extraRequires : extrasRequire.getElements()) {
161       final Pair<String, List<PyRequirement>> extraResult = getExtraRequires(extraRequires.getKey(), extraRequires.getValue());
162       if (extraResult != null) {
163         result.put(extraResult.first, extraResult.second);
164       }
165     }
166
167     return result;
168   }
169
170   @Nullable
171   private static Pair<String, List<PyRequirement>> getExtraRequires(@NotNull PyExpression extra, @Nullable PyExpression requires) {
172     if (extra instanceof PyStringLiteralExpression) {
173       final List<String> requiresValue = resolveRequiresValue(requires);
174
175       if (requiresValue != null) {
176         return Pair.createNonNull(((PyStringLiteralExpression)extra).getStringValue(),
177                                   PyRequirementParser.fromText(StringUtil.join(requiresValue, "\n")));
178       }
179     }
180
181     return null;
182   }
183
184   @NotNull
185   private static List<PyRequirement> getSetupPyRequiresFromArguments(@NotNull PyCallExpression setupCall,
186                                                                      String @NotNull ... argumentNames) {
187     return PyRequirementParser.fromText(
188       StreamEx
189         .of(argumentNames)
190         .map(setupCall::getKeywordArgument)
191         .flatCollection(PyPackageUtil::resolveRequiresValue)
192         .joining("\n")
193     );
194   }
195
196   @NotNull
197   private static List<PyRequirement> mergeSetupPyRequirements(@NotNull List<PyRequirement> requirementsFromRequires,
198                                                               @NotNull List<PyRequirement> requirementsFromLinks) {
199     if (!requirementsFromLinks.isEmpty()) {
200       final Map<String, List<PyRequirement>> nameToRequirements =
201         requirementsFromRequires.stream().collect(Collectors.groupingBy(PyRequirement::getName, LinkedHashMap::new, Collectors.toList()));
202
203       for (PyRequirement requirementFromLinks : requirementsFromLinks) {
204         nameToRequirements.replace(requirementFromLinks.getName(), Collections.singletonList(requirementFromLinks));
205       }
206
207       return nameToRequirements.values().stream().flatMap(Collection::stream).collect(Collectors.toCollection(ArrayList::new));
208     }
209
210     return requirementsFromRequires;
211   }
212
213   /**
214    * @param expression expression to resolve
215    * @return {@code expression} if it is not a reference or element that is found by following assignments chain.
216    * <em>Note: if result is {@link PyExpression} then parentheses around will be flattened.</em>
217    */
218   @Nullable
219   private static PsiElement resolveValue(@Nullable PyExpression expression) {
220     final PsiElement elementToAnalyze = PyPsiUtils.flattenParens(expression);
221
222     if (elementToAnalyze instanceof PyReferenceExpression) {
223       final TypeEvalContext context = TypeEvalContext.deepCodeInsight(elementToAnalyze.getProject());
224       final PyResolveContext resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context);
225
226       return StreamEx
227         .of(((PyReferenceExpression)elementToAnalyze).multiFollowAssignmentsChain(resolveContext))
228         .map(ResolveResult::getElement)
229         .findFirst(Objects::nonNull)
230         .map(e -> e instanceof PyExpression ? PyPsiUtils.flattenParens((PyExpression)e) : e)
231         .orElse(null);
232     }
233
234     return elementToAnalyze;
235   }
236
237   @Nullable
238   private static List<String> resolveRequiresValue(@Nullable PyExpression expression) {
239     final PsiElement elementToAnalyze = resolveValue(expression);
240
241     if (elementToAnalyze instanceof PyStringLiteralExpression) {
242       return Collections.singletonList(((PyStringLiteralExpression)elementToAnalyze).getStringValue());
243     }
244     else if (elementToAnalyze instanceof PyListLiteralExpression || elementToAnalyze instanceof PyTupleExpression) {
245       return StreamEx
246         .of(((PySequenceExpression)elementToAnalyze).getElements())
247         .map(PyPackageUtil::resolveValue)
248         .select(PyStringLiteralExpression.class)
249         .map(PyStringLiteralExpression::getStringValue)
250         .toList();
251     }
252
253     return null;
254   }
255
256   @NotNull
257   public static List<String> getPackageNames(@NotNull Module module) {
258     // TODO: Cache found module packages, clear cache on module updates
259     final List<String> packageNames = new ArrayList<>();
260     final Project project = module.getProject();
261     VirtualFile[] roots = ModuleRootManager.getInstance(module).getSourceRoots();
262     if (roots.length == 0) {
263       roots = ModuleRootManager.getInstance(module).getContentRoots();
264     }
265     for (VirtualFile root : roots) {
266       collectPackageNames(project, root, packageNames);
267     }
268     return packageNames;
269   }
270
271   @NotNull
272   public static String requirementsToString(@NotNull List<? extends PyRequirement> requirements) {
273     return StringUtil.join(requirements, requirement -> String.format("'%s'", requirement.getPresentableText()), ", ");
274   }
275
276   @Nullable
277   private static PyCallExpression findSetupCall(@NotNull PyFile file) {
278     final Ref<PyCallExpression> result = new Ref<>(null);
279     file.acceptChildren(new PyRecursiveElementVisitor() {
280       @Override
281       public void visitPyCallExpression(@NotNull PyCallExpression node) {
282         final PyExpression callee = node.getCallee();
283         final String name = PyUtil.getReadableRepr(callee, true);
284         if ("setup".equals(name)) {
285           result.set(node);
286         }
287       }
288
289       @Override
290       public void visitPyElement(@NotNull PyElement node) {
291         if (!(node instanceof ScopeOwner)) {
292           super.visitPyElement(node);
293         }
294       }
295     });
296     return result.get();
297   }
298
299   @Nullable
300   public static PyCallExpression findSetupCall(@NotNull Module module) {
301     return Optional
302       .ofNullable(findSetupPy(module))
303       .map(PyPackageUtil::findSetupCall)
304       .orElse(null);
305   }
306
307   private static void collectPackageNames(@NotNull final Project project,
308                                           @NotNull final VirtualFile root,
309                                           @NotNull final List<String> results) {
310     final ProjectFileIndex fileIndex = ProjectRootManager.getInstance(project).getFileIndex();
311     VfsUtilCore.visitChildrenRecursively(root, new VirtualFileVisitor<Void>() {
312       @Override
313       public boolean visitFile(@NotNull VirtualFile file) {
314         if (file.equals(root)) {
315           return true;
316         }
317         if (!fileIndex.isExcluded(file) && file.isDirectory() && file.findChild(PyNames.INIT_DOT_PY) != null) {
318           results.add(VfsUtilCore.getRelativePath(file, root, '.'));
319           return true;
320         }
321         return false;
322       }
323     });
324   }
325
326   public static boolean packageManagementEnabled(@Nullable Sdk sdk) {
327     if (!PythonSdkUtil.isRemote(sdk)) {
328       return true;
329     }
330     return new CredentialsTypeExChecker() {
331       @Override
332       protected boolean checkLanguageContribution(PyCredentialsContribution languageContribution) {
333         return languageContribution.isPackageManagementEnabled();
334       }
335     }.check(sdk);
336   }
337
338   /**
339    * Refresh the list of installed packages inside the specified SDK if it hasn't been updated yet
340    * displaying modal progress bar in the process, return cached packages otherwise.
341    * <p>
342    * Note that <strong>you shall never call this method from a write action</strong>, since such modal
343    * tasks are executed directly on EDT and network operations on the dispatch thread are prohibited
344    * (see the implementation of ApplicationImpl#runProcessWithProgressSynchronously() for details).
345    */
346   @NotNull
347   public static List<PyPackage> refreshAndGetPackagesModally(@NotNull Sdk sdk) {
348
349     final Application app = ApplicationManager.getApplication();
350     assert !(app.isWriteAccessAllowed()) :
351       "This method can't be called on WriteAction because " +
352       "refreshAndGetPackages would be called on AWT thread in this case (see runProcessWithProgressSynchronously) " +
353       "and may lead to freeze";
354
355
356     final Ref<List<PyPackage>> packagesRef = Ref.create();
357     final Throwable callStacktrace = new Throwable();
358     LOG.debug("Showing modal progress for collecting installed packages", new Throwable());
359     PyUtil.runWithProgress(null, PyBundle.message("sdk.scanning.installed.packages"), true, false, indicator -> {
360       if (PythonSdkUtil.isDisposed(sdk)) {
361         packagesRef.set(Collections.emptyList());
362         return;
363       }
364
365       indicator.setIndeterminate(true);
366       try {
367         final PyPackageManager manager = PyPackageManager.getInstance(sdk);
368         packagesRef.set(manager.refreshAndGetPackages(false));
369       }
370       catch (ExecutionException e) {
371         packagesRef.set(Collections.emptyList());
372         e.initCause(callStacktrace);
373         LOG.warn(e);
374       }
375     });
376     return packagesRef.get();
377   }
378
379   /**
380    * Run unconditional update of the list of packages installed in SDK. Normally only one such of updates should run at time.
381    * This behavior in enforced by the parameter isUpdating.
382    *
383    * @param manager    package manager for SDK
384    * @param isUpdating flag indicating whether another refresh is already running
385    * @return whether packages were refreshed successfully, e.g. this update wasn't cancelled because of another refresh in progress
386    */
387   public static boolean updatePackagesSynchronouslyWithGuard(@NotNull PyPackageManager manager, @NotNull AtomicBoolean isUpdating) {
388     assert !ApplicationManager.getApplication().isDispatchThread();
389     if (!isUpdating.compareAndSet(false, true)) {
390       return false;
391     }
392     try {
393       if (manager instanceof PyPackageManagerImpl) {
394         LOG.info("Refreshing installed packages for SDK " + ((PyPackageManagerImpl)manager).getSdk().getHomePath());
395       }
396       manager.refreshAndGetPackages(true);
397     }
398     catch (ExecutionException ignored) {
399     }
400     finally {
401       isUpdating.set(false);
402     }
403     return true;
404   }
405
406
407   public static boolean hasManagement(@NotNull List<PyPackage> packages) {
408     return (PyPsiPackageUtil.findPackage(packages, SETUPTOOLS) != null || PyPsiPackageUtil.findPackage(packages, DISTRIBUTE) != null) ||
409            PyPsiPackageUtil.findPackage(packages, PIP) != null;
410   }
411
412   @Nullable
413   public static List<PyRequirement> getRequirementsFromTxt(@NotNull Module module) {
414     final VirtualFile requirementsTxt = findRequirementsTxt(module);
415     if (requirementsTxt != null) {
416       return PyRequirementParser.fromFile(requirementsTxt);
417     }
418     return null;
419   }
420
421   public static void addRequirementToTxtOrSetupPy(@NotNull Module module,
422                                                   @NotNull String requirementName,
423                                                   @NotNull LanguageLevel languageLevel) {
424     final VirtualFile requirementsTxt = findRequirementsTxt(module);
425     if (requirementsTxt != null && requirementsTxt.isWritable()) {
426       final Document document = FileDocumentManager.getInstance().getDocument(requirementsTxt);
427       if (document != null) {
428         document.insertString(0, requirementName + "\n");
429       }
430       return;
431     }
432
433     final PyFile setupPy = findSetupPy(module);
434     if (setupPy == null) return;
435
436     final PyCallExpression setupCall = findSetupCall(setupPy);
437     if (setupCall == null) return;
438
439     final PsiElement installRequires = findSetupPyInstallRequires(setupCall);
440     if (installRequires != null) {
441       addRequirementToInstallRequires(installRequires, requirementName, languageLevel);
442     }
443     else {
444       final PyArgumentList argumentList = setupCall.getArgumentList();
445       final PyKeywordArgument requiresArg = generateRequiresKwarg(setupPy, requirementName, languageLevel);
446
447       if (argumentList != null && requiresArg != null) {
448         argumentList.addArgument(requiresArg);
449       }
450     }
451   }
452
453   private static void addRequirementToInstallRequires(@NotNull PsiElement installRequires,
454                                                       @NotNull String requirementName,
455                                                       @NotNull LanguageLevel languageLevel) {
456     final PyElementGenerator generator = PyElementGenerator.getInstance(installRequires.getProject());
457     final PyExpression newRequirement = generator.createExpressionFromText(languageLevel, "'" + requirementName + "'");
458
459     if (installRequires instanceof PyListLiteralExpression) {
460       installRequires.add(newRequirement);
461     }
462     else if (installRequires instanceof PyTupleExpression) {
463       final String newInstallRequiresText = StreamEx
464         .of(((PyTupleExpression)installRequires).getElements())
465         .append(newRequirement)
466         .map(PyExpression::getText)
467         .joining(",", "(", ")");
468
469       final PyExpression expression = generator.createExpressionFromText(languageLevel, newInstallRequiresText);
470
471       Optional
472         .ofNullable(PyUtil.as(expression, PyParenthesizedExpression.class))
473         .map(PyParenthesizedExpression::getContainedExpression)
474         .map(e -> PyUtil.as(e, PyTupleExpression.class))
475         .ifPresent(e -> installRequires.replace(e));
476     }
477     else if (installRequires instanceof PyStringLiteralExpression) {
478       final PyListLiteralExpression newInstallRequires = generator.createListLiteral();
479
480       newInstallRequires.add(installRequires);
481       newInstallRequires.add(newRequirement);
482
483       installRequires.replace(newInstallRequires);
484     }
485   }
486
487   @Nullable
488   private static PyKeywordArgument generateRequiresKwarg(@NotNull PyFile setupPy,
489                                                          @NotNull String requirementName,
490                                                          @NotNull LanguageLevel languageLevel) {
491     final String keyword = PyPsiUtils.containsImport(setupPy, "setuptools") ? INSTALL_REQUIRES : REQUIRES;
492     final String text = String.format("foo(%s=['%s'])", keyword, requirementName);
493     final PyExpression generated = PyElementGenerator.getInstance(setupPy.getProject()).createExpressionFromText(languageLevel, text);
494
495     if (generated instanceof PyCallExpression) {
496       final PyCallExpression callExpression = (PyCallExpression)generated;
497
498       return Stream
499         .of(callExpression.getArguments())
500         .filter(PyKeywordArgument.class::isInstance)
501         .map(PyKeywordArgument.class::cast)
502         .filter(kwarg -> keyword.equals(kwarg.getKeyword()))
503         .findFirst()
504         .orElse(null);
505     }
506
507     return null;
508   }
509
510   /**
511    * @deprecated Use {@link #runOnChangeUnderInterpreterPaths(Sdk, Disposable, Runnable)} passing a package manager as
512    * the parent disposable.
513    */
514   @ApiStatus.ScheduledForRemoval(inVersion = "2021.1")
515   @Deprecated
516   public static void runOnChangeUnderInterpreterPaths(@NotNull Sdk sdk, @NotNull Runnable runnable) {
517     Application app = ApplicationManager.getApplication();
518     Disposable parentDisposable = sdk instanceof Disposable ? (Disposable)sdk : app;
519     runOnChangeUnderInterpreterPaths(sdk, parentDisposable, runnable);
520   }
521
522   /**
523    * Execute the given executable on a pooled thread whenever there is a VFS event happening under some of the roots of the SDK.
524    *
525    * @param sdk              SDK those roots need to be watched
526    * @param parentDisposable disposable for the registered event listeners
527    * @param runnable         executable that's going to be executed
528    */
529   public static void runOnChangeUnderInterpreterPaths(@NotNull Sdk sdk,
530                                                       @NotNull Disposable parentDisposable,
531                                                       @NotNull Runnable runnable) {
532     final Application app = ApplicationManager.getApplication();
533     VirtualFileManager.getInstance().addAsyncFileListener(new AsyncFileListener() {
534       @Nullable
535       @Override
536       public ChangeApplier prepareChange(@NotNull List<? extends VFileEvent> events) {
537         final Set<VirtualFile> roots = getPackagingAwareSdkRoots(sdk);
538         if (roots.isEmpty()) return null;
539         allEvents:
540         for (VFileEvent event : events) {
541           if (event instanceof VFileContentChangeEvent || event instanceof VFilePropertyChangeEvent) continue;
542           // In case of create event getFile() returns null as the file hasn't been created yet
543           VirtualFile parent = null;
544           if (event instanceof VFileCreateEvent) {
545             parent = ((VFileCreateEvent)event).getParent();
546           }
547           else {
548             VirtualFile file = event.getFile();
549             if (file != null) parent = file.getParent();
550           }
551
552           if (parent != null && roots.contains(parent)) {
553             InterpreterChangeEvents.LOG.debug("Interpreter change in " + parent + " indicated by " + event +
554                                               " (all events: " + events + ")");
555             app.executeOnPooledThread(runnable);
556             break allEvents;
557           }
558         }
559         // No continuation in write action is needed
560         return null;
561       }
562     }, parentDisposable);
563   }
564
565   @NotNull
566   private static Set<VirtualFile> getPackagingAwareSdkRoots(@NotNull Sdk sdk) {
567     final Set<VirtualFile> result = Sets.newHashSet(sdk.getRootProvider().getFiles(OrderRootType.CLASSES));
568     final String skeletonsPath = PythonSdkUtil.getSkeletonsPath(PathManager.getSystemPath(), sdk.getHomePath());
569     final VirtualFile skeletonsRoot = LocalFileSystem.getInstance().findFileByPath(skeletonsPath);
570     result.removeIf(vf -> vf.equals(skeletonsRoot) ||
571                           vf.equals(PyUserSkeletonsUtil.getUserSkeletonsDirectory()) ||
572                           PyTypeShed.INSTANCE.isInside(vf));
573     return result;
574   }
575 }