0fc1b2017a35e197e6f8f902a71f834dc89863cb
[idea/community.git] / python / src / com / jetbrains / python / packaging / PyPackageUtil.java
1 /*
2  * Copyright 2000-2016 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.packaging;
17
18 import com.intellij.execution.ExecutionException;
19 import com.intellij.openapi.application.ApplicationManager;
20 import com.intellij.openapi.diagnostic.Logger;
21 import com.intellij.openapi.editor.Document;
22 import com.intellij.openapi.fileEditor.FileDocumentManager;
23 import com.intellij.openapi.module.Module;
24 import com.intellij.openapi.project.Project;
25 import com.intellij.openapi.projectRoots.Sdk;
26 import com.intellij.openapi.roots.ModuleRootManager;
27 import com.intellij.openapi.roots.ProjectFileIndex;
28 import com.intellij.openapi.roots.ProjectRootManager;
29 import com.intellij.openapi.util.Ref;
30 import com.intellij.openapi.util.text.StringUtil;
31 import com.intellij.openapi.vfs.LocalFileSystem;
32 import com.intellij.openapi.vfs.VfsUtilCore;
33 import com.intellij.openapi.vfs.VirtualFile;
34 import com.intellij.openapi.vfs.VirtualFileVisitor;
35 import com.intellij.psi.PsiElement;
36 import com.intellij.psi.PsiFile;
37 import com.intellij.psi.PsiManager;
38 import com.jetbrains.python.PyBundle;
39 import com.jetbrains.python.PyNames;
40 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
41 import com.jetbrains.python.packaging.setupPy.SetupTaskIntrospector;
42 import com.jetbrains.python.psi.*;
43 import com.jetbrains.python.psi.resolve.PyResolveContext;
44 import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
45 import com.jetbrains.python.psi.types.TypeEvalContext;
46 import com.jetbrains.python.remote.PyCredentialsContribution;
47 import com.jetbrains.python.sdk.CredentialsTypeExChecker;
48 import com.jetbrains.python.sdk.PythonSdkType;
49 import org.jetbrains.annotations.NotNull;
50 import org.jetbrains.annotations.Nullable;
51
52 import java.util.*;
53 import java.util.concurrent.atomic.AtomicBoolean;
54 import java.util.stream.Collectors;
55 import java.util.stream.Stream;
56
57 /**
58  * @author vlan
59  */
60 public class PyPackageUtil {
61   public static final String SETUPTOOLS = "setuptools";
62   public static final String PIP = "pip";
63   public static final String DISTRIBUTE = "distribute";
64   private static final Logger LOG = Logger.getInstance(PyPackageUtil.class);
65
66   @NotNull
67   private static final String REQUIRES = "requires";
68
69   @NotNull
70   private static final String INSTALL_REQUIRES = "install_requires";
71
72   @NotNull
73   private static final String[] SETUP_PY_REQUIRES_KWARGS_NAMES = new String[] {
74     REQUIRES, INSTALL_REQUIRES, "setup_requires", "tests_require"
75   };
76
77   @NotNull
78   private static final String DEPENDENCY_LINKS = "dependency_links";
79
80   private PyPackageUtil() {
81   }
82
83   public static boolean hasSetupPy(@NotNull Module module) {
84     return findSetupPy(module) != null;
85   }
86
87   @Nullable
88   public static PyFile findSetupPy(@NotNull Module module) {
89     for (VirtualFile root : PyUtil.getSourceRoots(module)) {
90       final VirtualFile child = root.findChild("setup.py");
91       if (child != null) {
92         final PsiFile file = PsiManager.getInstance(module.getProject()).findFile(child);
93         if (file instanceof PyFile) {
94           return (PyFile)file;
95         }
96       }
97     }
98     return null;
99   }
100
101   public static boolean hasRequirementsTxt(@NotNull Module module) {
102     return findRequirementsTxt(module) != null;
103   }
104
105   @Nullable
106   public static VirtualFile findRequirementsTxt(@NotNull Module module) {
107     final String requirementsPath = PyPackageRequirementsSettings.getInstance(module).getRequirementsPath();
108     if (!requirementsPath.isEmpty()) {
109       final VirtualFile file = LocalFileSystem.getInstance().findFileByPath(requirementsPath);
110       if (file != null) {
111         return file;
112       }
113       final ModuleRootManager manager = ModuleRootManager.getInstance(module);
114       for (VirtualFile root : manager.getContentRoots()) {
115         final VirtualFile fileInRoot = root.findFileByRelativePath(requirementsPath);
116         if (fileInRoot != null) {
117           return fileInRoot;
118         }
119       }
120     }
121     return null;
122   }
123
124   @Nullable
125   private static PyListLiteralExpression findSetupPyInstallRequires(@NotNull Module module, @Nullable PyCallExpression setupCall) {
126     if (setupCall == null) {
127       return null;
128     }
129
130     return Stream
131       .of(REQUIRES, INSTALL_REQUIRES)
132       .map(setupCall::getKeywordArgument)
133       .map(requires -> resolveRequiresValue(module, requires))
134       .filter(requires -> requires != null)
135       .findFirst()
136       .orElse(null);
137   }
138
139   @Nullable
140   public static List<PyRequirement> findSetupPyRequires(@NotNull Module module) {
141     final PyCallExpression setupCall = findSetupCall(module);
142
143     if (setupCall == null) {
144       return null;
145     }
146
147     final List<PyRequirement> requirementsFromRequires = getSetupPyRequiresFromArguments(module, setupCall, SETUP_PY_REQUIRES_KWARGS_NAMES);
148     final List<PyRequirement> requirementsFromLinks = getSetupPyRequiresFromArguments(module, setupCall, DEPENDENCY_LINKS);
149
150     return mergeSetupPyRequirements(requirementsFromRequires, requirementsFromLinks);
151   }
152
153   @NotNull
154   private static List<PyRequirement> getSetupPyRequiresFromArguments(@NotNull Module module,
155                                                                      @NotNull PyCallExpression setupCall,
156                                                                      @NotNull String... argumentNames) {
157     return PyRequirement.fromText(
158       Stream
159         .of(argumentNames)
160         .map(setupCall::getKeywordArgument)
161         .map(requires -> resolveRequiresValue(module, requires))
162         .filter(requires -> requires != null)
163         .flatMap(requires -> Stream.of(requires.getElements()))
164         .filter(PyStringLiteralExpression.class::isInstance)
165         .map(requirement -> ((PyStringLiteralExpression)requirement).getStringValue())
166         .collect(Collectors.joining("\n"))
167     );
168   }
169
170   @NotNull
171   private static List<PyRequirement> mergeSetupPyRequirements(@NotNull List<PyRequirement> requirementsFromRequires,
172                                                               @NotNull List<PyRequirement> requirementsFromLinks) {
173     if (!requirementsFromLinks.isEmpty()) {
174       final Map<String, List<PyRequirement>> nameToRequirements =
175         requirementsFromRequires.stream().collect(Collectors.groupingBy(PyRequirement::getName, LinkedHashMap::new, Collectors.toList()));
176
177       for (PyRequirement requirementFromLinks : requirementsFromLinks) {
178         nameToRequirements.replace(requirementFromLinks.getName(), Collections.singletonList(requirementFromLinks));
179       }
180
181       return nameToRequirements.values().stream().flatMap(Collection::stream).collect(Collectors.toCollection(ArrayList::new));
182     }
183
184     return requirementsFromRequires;
185   }
186
187   @Nullable
188   private static PyListLiteralExpression resolveRequiresValue(@NotNull Module module, @Nullable PyExpression requires) {
189     if (requires instanceof PyListLiteralExpression) {
190       return (PyListLiteralExpression)requires;
191     }
192     if (requires instanceof PyReferenceExpression) {
193       final TypeEvalContext context = TypeEvalContext.deepCodeInsight(module.getProject());
194       final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
195       final QualifiedResolveResult result = ((PyReferenceExpression)requires).followAssignmentsChain(resolveContext);
196       final PsiElement element = result.getElement();
197       if (element instanceof PyListLiteralExpression) {
198         return (PyListLiteralExpression)element;
199       }
200     }
201     return null;
202   }
203
204   @NotNull
205   public static List<String> getPackageNames(@NotNull Module module) {
206     // TODO: Cache found module packages, clear cache on module updates
207     final List<String> packageNames = new ArrayList<String>();
208     final Project project = module.getProject();
209     VirtualFile[] roots = ModuleRootManager.getInstance(module).getSourceRoots();
210     if (roots.length == 0) {
211       roots = ModuleRootManager.getInstance(module).getContentRoots();
212     }
213     for (VirtualFile root : roots) {
214       collectPackageNames(project, root, packageNames);
215     }
216     return packageNames;
217   }
218
219   @NotNull
220   public static String requirementsToString(@NotNull List<PyRequirement> requirements) {
221     return StringUtil.join(requirements, requirement -> String.format("'%s'", requirement.toString()), ", ");
222   }
223
224   @Nullable
225   public static PyCallExpression findSetupCall(@NotNull PyFile file) {
226     final Ref<PyCallExpression> result = new Ref<PyCallExpression>(null);
227     file.acceptChildren(new PyRecursiveElementVisitor() {
228       @Override
229       public void visitPyCallExpression(PyCallExpression node) {
230         final PyExpression callee = node.getCallee();
231         final String name = PyUtil.getReadableRepr(callee, true);
232         if ("setup".equals(name)) {
233           result.set(node);
234         }
235       }
236
237       @Override
238       public void visitPyElement(PyElement node) {
239         if (!(node instanceof ScopeOwner)) {
240           super.visitPyElement(node);
241         }
242       }
243     });
244     return result.get();
245   }
246
247   @Nullable
248   public static PyCallExpression findSetupCall(@NotNull Module module) {
249     return Optional
250       .ofNullable(findSetupPy(module))
251       .map(PyPackageUtil::findSetupCall)
252       .orElse(null);
253   }
254
255   private static void collectPackageNames(@NotNull final Project project,
256                                           @NotNull final VirtualFile root,
257                                           @NotNull final List<String> results) {
258     final ProjectFileIndex fileIndex = ProjectRootManager.getInstance(project).getFileIndex();
259     VfsUtilCore.visitChildrenRecursively(root, new VirtualFileVisitor() {
260       @Override
261       public boolean visitFile(@NotNull VirtualFile file) {
262         if (!fileIndex.isExcluded(file) && file.isDirectory() && file.findChild(PyNames.INIT_DOT_PY) != null) {
263           results.add(VfsUtilCore.getRelativePath(file, root, '.'));
264         }
265         return true;
266       }
267     });
268   }
269
270   public static boolean packageManagementEnabled(@Nullable Sdk sdk) {
271     if (!PythonSdkType.isRemote(sdk)) {
272       return true;
273     }
274     return new CredentialsTypeExChecker() {
275       @Override
276       protected boolean checkLanguageContribution(PyCredentialsContribution languageContribution) {
277         return languageContribution.isPackageManagementEnabled();
278       }
279     }.withSshContribution(true).withVagrantContribution(true).withWebDeploymentContribution(true).check(sdk);
280   }
281
282   @Nullable
283   public static List<PyPackage> refreshAndGetPackagesModally(@NotNull Sdk sdk) {
284     final Ref<List<PyPackage>> packagesRef = Ref.create();
285     LOG.debug("Showing modal progress for collecting installed packages", new Throwable());
286     PyUtil.runWithProgress(null, PyBundle.message("sdk.scanning.installed.packages"), true, false, indicator -> {
287       indicator.setIndeterminate(true);
288       try {
289         packagesRef.set(PyPackageManager.getInstance(sdk).refreshAndGetPackages(false));
290       }
291       catch (ExecutionException e) {
292         LOG.warn(e);
293       }
294     });
295     return packagesRef.get();
296   }
297
298   /**
299    * Run unconditional update of the list of packages installed in SDK. Normally only one such of updates should run at time.
300    * This behavior in enforced by the parameter isUpdating.
301    * 
302    * @param manager    package manager for SDK
303    * @param isUpdating flag indicating whether another refresh is already running
304    * @return whether packages were refreshed successfully, e.g. this update wasn't cancelled because of another refresh in progress
305    */
306   public static boolean updatePackagesSynchronouslyWithGuard(@NotNull PyPackageManager manager, @NotNull AtomicBoolean isUpdating) {
307     assert !ApplicationManager.getApplication().isDispatchThread();
308     if (!isUpdating.compareAndSet(false, true)) {
309       return false;
310     }
311     try {
312       if (manager instanceof PyPackageManagerImpl) {
313         LOG.info("Refreshing installed packages for SDK " + ((PyPackageManagerImpl)manager).getSdk().getHomePath());
314       }
315       manager.refreshAndGetPackages(true);
316     }
317     catch (ExecutionException ignored) {
318     }
319     finally {
320       isUpdating.set(false);
321     }
322     return true;
323   }
324   
325
326   @Nullable
327   public static PyPackage findPackage(@NotNull List<PyPackage> packages, @NotNull String name) {
328     for (PyPackage pkg : packages) {
329       if (name.equalsIgnoreCase(pkg.getName())) {
330         return pkg;
331       }
332     }
333     return null;
334   }
335
336   public static boolean hasManagement(@NotNull List<PyPackage> packages) {
337     return (findPackage(packages, SETUPTOOLS) != null || findPackage(packages, DISTRIBUTE) != null) ||
338            findPackage(packages, PIP) != null;
339   }
340
341   @Nullable
342   public static List<PyRequirement> getRequirementsFromTxt(@NotNull Module module) {
343     final VirtualFile requirementsTxt = findRequirementsTxt(module);
344     if (requirementsTxt != null) {
345       return PyRequirement.fromFile(requirementsTxt);
346     }
347     return null;
348   }
349
350   public static void addRequirementToTxtOrSetupPy(@NotNull Module module,
351                                                   @NotNull String requirementName,
352                                                   @NotNull LanguageLevel languageLevel) {
353     final VirtualFile requirementsTxt = findRequirementsTxt(module);
354     if (requirementsTxt != null && requirementsTxt.isWritable()) {
355       final Document document = FileDocumentManager.getInstance().getDocument(requirementsTxt);
356       if (document != null) {
357         document.insertString(0, requirementName + "\n");
358       }
359       return;
360     }
361
362     final PyFile setupPy = findSetupPy(module);
363     if (setupPy == null) {
364       return;
365     }
366
367     final PyCallExpression setupCall = findSetupCall(setupPy);
368     final PyListLiteralExpression installRequires = findSetupPyInstallRequires(module, setupCall);
369     final PyElementGenerator generator = PyElementGenerator.getInstance(module.getProject());
370
371     if (installRequires != null && installRequires.isWritable()) {
372       final String text = String.format("'%s'", requirementName);
373       final PyExpression generated = generator.createExpressionFromText(languageLevel, text);
374       installRequires.add(generated);
375
376       return;
377     }
378
379     if (setupCall != null) {
380       final PyArgumentList argumentList = setupCall.getArgumentList();
381       final PyKeywordArgument requiresArg = generateRequiresKwarg(setupPy, requirementName, languageLevel, generator);
382
383       if (argumentList != null && requiresArg != null) {
384         argumentList.addArgument(requiresArg);
385       }
386     }
387   }
388
389   @Nullable
390   private static PyKeywordArgument generateRequiresKwarg(@NotNull PyFile setupPy,
391                                                          @NotNull String requirementName,
392                                                          @NotNull LanguageLevel languageLevel,
393                                                          @NotNull PyElementGenerator generator) {
394     final String keyword = SetupTaskIntrospector.usesSetuptools(setupPy) ? INSTALL_REQUIRES : REQUIRES;
395     final String text = String.format("foo(%s=['%s'])", keyword, requirementName);
396     final PyExpression generated = generator.createExpressionFromText(languageLevel, text);
397
398     if (generated instanceof PyCallExpression) {
399       final PyCallExpression callExpression = (PyCallExpression)generated;
400
401       return Stream
402         .of(callExpression.getArguments())
403         .filter(PyKeywordArgument.class::isInstance)
404         .map(PyKeywordArgument.class::cast)
405         .filter(kwarg -> keyword.equals(kwarg.getKeyword()))
406         .findFirst()
407         .orElse(null);
408     }
409
410     return null;
411   }
412 }