1 /*
2  * Copyright (C) 2016 The Android Open Source Project
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 
17 package com.android.tradefed.targetprep;
18 
19 import com.android.annotations.VisibleForTesting;
20 import com.android.tradefed.build.IBuildInfo;
21 import com.android.tradefed.command.remote.DeviceDescriptor;
22 import com.android.tradefed.config.Option;
23 import com.android.tradefed.config.OptionClass;
24 import com.android.tradefed.device.DeviceNotAvailableException;
25 import com.android.tradefed.device.ITestDevice;
26 import com.android.tradefed.invoker.IInvocationContext;
27 import com.android.tradefed.log.LogUtil.CLog;
28 import com.android.tradefed.targetprep.multi.IMultiTargetPreparer;
29 import com.android.tradefed.util.CommandResult;
30 import com.android.tradefed.util.CommandStatus;
31 import com.android.tradefed.util.EnvUtil;
32 import com.android.tradefed.util.FileUtil;
33 import com.android.tradefed.util.IRunUtil;
34 import com.android.tradefed.util.RunUtil;
35 import com.android.tradefed.util.VtsFileUtil;
36 import com.android.tradefed.util.VtsPythonRunnerHelper;
37 import com.android.tradefed.util.VtsVendorConfigFileUtil;
38 
39 import java.io.File;
40 import java.io.IOException;
41 import java.nio.file.FileVisitResult;
42 import java.nio.file.Files;
43 import java.nio.file.Path;
44 import java.nio.file.SimpleFileVisitor;
45 import java.nio.file.attribute.BasicFileAttributes;
46 import java.util.Collection;
47 import java.util.LinkedHashSet;
48 import java.util.Map;
49 import java.util.NoSuchElementException;
50 import java.util.TreeMap;
51 import java.util.TreeSet;
52 
53 /**
54  * Sets up a Python virtualenv on the host and installs packages. To activate it, the working
55  * directory is changed to the root of the virtualenv.
56  *
57  * This's a fork of PythonVirtualenvPreparer and is forked in order to simplify the change
58  * deployment process and reduce the deployment time, which are critical for VTS services.
59  * That means changes here will be upstreamed gradually.
60  */
61 @OptionClass(alias = "python-venv")
62 public class VtsPythonVirtualenvPreparer implements IMultiTargetPreparer {
63     private static final String LOCAL_PYPI_PATH_ENV_VAR_NAME = "VTS_PYPI_PATH";
64     private static final String LOCAL_PYPI_PATH_KEY = "pypi_packages_path";
65     private static final int SECOND_IN_MSECS = 1000;
66     private static final int MINUTE_IN_MSECS = 60 * SECOND_IN_MSECS;
67     protected static int PIP_RETRY = 3;
68     private static final int PIP_RETRY_WAIT = 3 * SECOND_IN_MSECS;
69     protected static final int PIP_INSTALL_DELAY = SECOND_IN_MSECS;
70     public static final String VIRTUAL_ENV_V3 = "VIRTUAL_ENV_V3";
71     public static final String VIRTUAL_ENV = "VIRTUAL_ENV";
72 
73     @Option(name = "venv-dir", description = "path of an existing virtualenv to use")
74     protected File mVenvDir = null;
75 
76     @Option(name = "requirements-file", description = "pip-formatted requirements file")
77     private File mRequirementsFile = null;
78 
79     @Option(name = "script-file", description = "scripts which need to be executed in advance")
80     private Collection<String> mScriptFiles = new TreeSet<>();
81 
82     @Option(name = "dep-module", description = "modules which need to be installed by pip")
83     protected Collection<String> mDepModules = new LinkedHashSet<>();
84 
85     @Option(name = "no-dep-module", description = "modules which should not be installed by pip")
86     private Collection<String> mNoDepModules = new TreeSet<>();
87 
88     @Option(name = "reuse",
89             description = "Reuse an exising virtualenv path if exists in "
90                     + "temp directory. When this option is enabled, virtualenv directory used or "
91                     + "created by this preparer will not be deleted after tests complete.")
92     protected boolean mReuse = true;
93 
94     @Option(name = "python-version",
95             description = "The version of a Python interpreter to use."
96                     + "Currently, only major version number is fully supported."
97                     + "Example: \"2\", or \"3\".")
98     private String mPythonVersion = "2";
99 
100     private IBuildInfo mBuildInfo = null;
101     private DeviceDescriptor mDescriptor = null;
102     private IRunUtil mRunUtil = new RunUtil();
103 
104     String mLocalPypiPath = null;
105     String mPipPath = null;
106 
107     // Since we allow virtual env path to be reused during a test plan/module, only the preparer
108     // which created the directory should be the one to delete it.
109     private boolean mIsDirCreator = false;
110 
111     // If the same object is used in multiple threads (in sharding mode), the class
112     // needs to know when it is safe to call the teardown method.
113     private int mNumOfInstances = 0;
114 
115     // A map of initially installed pip modules and versions. Newly installed modules are not
116     // currently added automatically.
117     private Map<String, String> mPipInstallList = null;
118 
119     /**
120      * {@inheritDoc}
121      */
122     @Override
setUp(IInvocationContext context)123     public synchronized void setUp(IInvocationContext context)
124             throws TargetSetupError, BuildError, DeviceNotAvailableException {
125         ++mNumOfInstances;
126         mBuildInfo = context.getBuildInfos().get(0);
127         if (mNumOfInstances == 1) {
128             CLog.i("Preparing python dependencies...");
129             ITestDevice device = context.getDevices().get(0);
130             mDescriptor = device.getDeviceDescriptor();
131             initVirtualenv(mBuildInfo);
132             CLog.d("Python virtualenv path is: " + mVenvDir);
133             VtsPythonRunnerHelper.activateVirtualenv(getRunUtil(), mVenvDir.getAbsolutePath());
134             setLocalPypiPath();
135             installDeps();
136         }
137         addPathToBuild(mBuildInfo);
138     }
139 
140     /**
141      * {@inheritDoc}
142      */
143     @Override
tearDown(IInvocationContext context, Throwable e)144     public synchronized void tearDown(IInvocationContext context, Throwable e)
145             throws DeviceNotAvailableException {
146         --mNumOfInstances;
147         if (mNumOfInstances > 0) {
148             // Since this is a host side preparer, no need to repeat
149             return;
150         }
151         if (!mReuse && mVenvDir != null && mIsDirCreator) {
152             try {
153                 recursiveDelete(mVenvDir.toPath());
154                 CLog.d("Deleted the virtual env's temp working dir, %s.", mVenvDir);
155             } catch (IOException exception) {
156                 CLog.e("Failed to delete %s: %s", mVenvDir, exception);
157             }
158             mVenvDir = null;
159         }
160     }
161 
162     /**
163      * This method sets mLocalPypiPath, the local PyPI package directory to
164      * install python packages from in the installDeps method.
165      */
setLocalPypiPath()166     protected void setLocalPypiPath() {
167         VtsVendorConfigFileUtil configReader = new VtsVendorConfigFileUtil();
168         if (configReader.LoadVendorConfig(mBuildInfo)) {
169             // First try to load local PyPI directory path from vendor config file
170             try {
171                 String pypiPath = configReader.GetVendorConfigVariable(LOCAL_PYPI_PATH_KEY);
172                 if (pypiPath.length() > 0 && dirExistsAndHaveReadAccess(pypiPath)) {
173                     mLocalPypiPath = pypiPath;
174                     CLog.d(String.format("Loaded %s: %s", LOCAL_PYPI_PATH_KEY, mLocalPypiPath));
175                 }
176             } catch (NoSuchElementException e) {
177                 /* continue */
178             }
179         }
180 
181         // If loading path from vendor config file is unsuccessful,
182         // check local pypi path defined by LOCAL_PYPI_PATH_ENV_VAR_NAME
183         if (mLocalPypiPath == null) {
184             CLog.d("Checking whether local pypi packages directory exists");
185             String pypiPath = System.getenv(LOCAL_PYPI_PATH_ENV_VAR_NAME);
186             if (pypiPath == null) {
187                 CLog.d("Local pypi packages directory not specified by env var %s",
188                         LOCAL_PYPI_PATH_ENV_VAR_NAME);
189             } else if (dirExistsAndHaveReadAccess(pypiPath)) {
190                 mLocalPypiPath = pypiPath;
191                 CLog.d("Set local pypi packages directory to %s", pypiPath);
192             }
193         }
194 
195         if (mLocalPypiPath == null) {
196             CLog.d("Failed to set local pypi packages path. Therefore internet connection to "
197                     + "https://pypi.python.org/simple/ must be available to run VTS tests.");
198         }
199     }
200 
201     /**
202      * This method returns whether the given path is a dir that exists and the user has read access.
203      */
dirExistsAndHaveReadAccess(String path)204     private boolean dirExistsAndHaveReadAccess(String path) {
205         File pathDir = new File(path);
206         if (!pathDir.exists() || !pathDir.isDirectory()) {
207             CLog.d("Directory %s does not exist.", pathDir);
208             return false;
209         }
210 
211         if (!EnvUtil.isOnWindows()) {
212             CommandResult c = getRunUtil().runTimedCmd(MINUTE_IN_MSECS, "ls", path);
213             if (c.getStatus() != CommandStatus.SUCCESS) {
214                 CLog.d(String.format("Failed to read dir: %s. Result %s. stdout: %s, stderr: %s",
215                         path, c.getStatus(), c.getStdout(), c.getStderr()));
216                 return false;
217             }
218             return true;
219         } else {
220             try {
221                 String[] pathDirList = pathDir.list();
222                 if (pathDirList == null) {
223                     CLog.d("Failed to read dir: %s. Please check access permission.", pathDir);
224                     return false;
225                 }
226             } catch (SecurityException e) {
227                 CLog.d(String.format(
228                         "Failed to read dir %s with SecurityException %s", pathDir, e));
229                 return false;
230             }
231             return true;
232         }
233     }
234 
235     /**
236      * Installs all python pip module dependencies specified in options.
237      * @throws TargetSetupError if failed
238      */
installDeps()239     protected void installDeps() throws TargetSetupError {
240         boolean hasDependencies = false;
241         if (!mScriptFiles.isEmpty()) {
242             for (String scriptFile : mScriptFiles) {
243                 CLog.d("Attempting to execute a script, %s", scriptFile);
244                 CommandResult c = getRunUtil().runTimedCmd(5 * MINUTE_IN_MSECS, scriptFile);
245                 if (c.getStatus() != CommandStatus.SUCCESS) {
246                     CLog.e("Executing script %s failed", scriptFile);
247                     throw new TargetSetupError("Failed to source a script", mDescriptor);
248                 }
249             }
250         }
251 
252         if (mRequirementsFile != null) {
253             hasDependencies = true;
254             boolean success = false;
255 
256             long retry_interval = PIP_RETRY_WAIT;
257             for (int try_count = 0; try_count < PIP_RETRY + 1; try_count++) {
258                 if (try_count > 0) {
259                     getRunUtil().sleep(retry_interval);
260                     retry_interval *= 3;
261                 }
262 
263                 if (installPipRequirementFile(mRequirementsFile)) {
264                     success = true;
265                     break;
266                 }
267             }
268 
269             if (!success) {
270                 throw new TargetSetupError(
271                         "Failed to install pip requirement file " + mRequirementsFile, mDescriptor);
272             }
273         }
274 
275         if (!mDepModules.isEmpty()) {
276             for (String dep : mDepModules) {
277                 hasDependencies = true;
278 
279                 if (mNoDepModules.contains(dep) || isPipModuleInstalled(dep)) {
280                     continue;
281                 }
282 
283                 boolean success = installPipModuleLocally(dep);
284 
285                 long retry_interval = PIP_RETRY_WAIT;
286                 for (int retry_count = 0; retry_count < PIP_RETRY + 1; retry_count++) {
287                     if (retry_count > 0) {
288                         getRunUtil().sleep(retry_interval);
289                         retry_interval *= 3;
290                     }
291 
292                     if (success || (!success && installPipModule(dep))) {
293                         success = true;
294                         getRunUtil().sleep(PIP_INSTALL_DELAY);
295                         break;
296                     }
297                 }
298 
299                 if (!success) {
300                     throw new TargetSetupError("Failed to install pip module " + dep, mDescriptor);
301                 }
302             }
303         }
304         if (!hasDependencies) {
305             CLog.d("No dependencies to install");
306         }
307     }
308 
309     /**
310      * Installs a pip requirement file from Internet.
311      * @param req pip module requirement file object
312      * @return true if success. False otherwise
313      */
installPipRequirementFile(File req)314     private boolean installPipRequirementFile(File req) {
315         CommandResult result = getRunUtil().runTimedCmd(10 * MINUTE_IN_MSECS, getPipPath(),
316                 "install", "-r", mRequirementsFile.getAbsolutePath());
317         CLog.d(String.format("Result %s. stdout: %s, stderr: %s", result.getStatus(),
318                 result.getStdout(), result.getStderr()));
319 
320         return result.getStatus() == CommandStatus.SUCCESS;
321     }
322 
323     /**
324      * Installs a pip module from local directory.
325      * @param name of the module
326      * @return true if the module is successfully installed; false otherwise.
327      */
installPipModuleLocally(String name)328     private boolean installPipModuleLocally(String name) {
329         if (mLocalPypiPath == null) {
330             return false;
331         }
332         CLog.d("Attempting installation of %s from local directory", name);
333         CommandResult result = getRunUtil().runTimedCmd(5 * MINUTE_IN_MSECS, getPipPath(),
334                 "install", name, "--no-index", "--find-links=" + mLocalPypiPath);
335         CLog.d(String.format("Result %s. stdout: %s, stderr: %s", result.getStatus(),
336                 result.getStdout(), result.getStderr()));
337 
338         return result.getStatus() == CommandStatus.SUCCESS;
339     }
340 
341     /**
342      * Install a pip module from Internet
343      * @param name of the module
344      * @return true if success. False otherwise
345      */
installPipModule(String name)346     private boolean installPipModule(String name) {
347         CLog.d("Attempting installation of %s from PyPI", name);
348         CommandResult result =
349                 getRunUtil().runTimedCmd(5 * MINUTE_IN_MSECS, getPipPath(), "install", name);
350         CLog.d("Result %s. stdout: %s, stderr: %s", result.getStatus(), result.getStdout(),
351                 result.getStderr());
352         if (result.getStatus() != CommandStatus.SUCCESS) {
353             CLog.e("Installing %s from PyPI failed.", name);
354             CLog.d("Attempting to upgrade %s", name);
355             result = getRunUtil().runTimedCmd(
356                     5 * MINUTE_IN_MSECS, getPipPath(), "install", "--upgrade", name);
357 
358             CLog.d(String.format("Result %s. stdout: %s, stderr: %s", result.getStatus(),
359                     result.getStdout(), result.getStderr()));
360         }
361 
362         return result.getStatus() == CommandStatus.SUCCESS;
363     }
364 
365     /**
366      * This method returns absolute pip path in virtualenv.
367      *
368      * This method is needed because although PATH is set in IRunUtil, IRunUtil will still
369      * use pip from system path.
370      *
371      * @return absolute pip path in virtualenv. null if virtualenv not available.
372      */
getPipPath()373     public String getPipPath() {
374         if (mPipPath != null) {
375             return mPipPath;
376         }
377 
378         String virtualenvPath = mVenvDir.getAbsolutePath();
379         if (virtualenvPath == null) {
380             return null;
381         }
382         mPipPath = new File(VtsPythonRunnerHelper.getPythonBinDir(virtualenvPath), "pip")
383                            .getAbsolutePath();
384         return mPipPath;
385     }
386 
387     /**
388      * Get the major python version from option.
389      *
390      * Currently, only 2 and 3 are supported.
391      *
392      * @return major version number
393      * @throws TargetSetupError
394      */
getConfiguredPythonVersionMajor()395     protected int getConfiguredPythonVersionMajor() throws TargetSetupError {
396         if (mPythonVersion.startsWith("3.") || mPythonVersion.equals("3")) {
397             return 3;
398         } else if (mPythonVersion.startsWith("2.") || mPythonVersion.equals("2")) {
399             return 2;
400         } else {
401             throw new TargetSetupError("Unsupported python version " + mPythonVersion);
402         }
403     }
404 
405     /**
406      * Add PYTHONPATH and VIRTUAL_ENV_PATH to BuildInfo.
407      * @param buildInfo
408      * @throws TargetSetupError
409      */
addPathToBuild(IBuildInfo buildInfo)410     protected void addPathToBuild(IBuildInfo buildInfo) throws TargetSetupError {
411         String target = null;
412         switch (getConfiguredPythonVersionMajor()) {
413             case 2:
414                 target = VtsPythonVirtualenvPreparer.VIRTUAL_ENV;
415                 break;
416             case 3:
417                 target = VtsPythonVirtualenvPreparer.VIRTUAL_ENV_V3;
418                 break;
419         }
420 
421         if (!buildInfo.getBuildAttributes().containsKey(target)) {
422             buildInfo.addBuildAttribute(target, mVenvDir.getAbsolutePath());
423         }
424     }
425 
426     /**
427      * Create virtualenv directory by executing virtualenv command.
428      * @param buildInfo
429      * @throws TargetSetupError
430      */
initVirtualenv(IBuildInfo buildInfo)431     protected void initVirtualenv(IBuildInfo buildInfo) throws TargetSetupError {
432         if (checkTestPlanLevelVirtualenv(buildInfo)) {
433             return;
434         }
435 
436         try {
437             if (checkHostReuseVirtualenv(buildInfo)) {
438                 return;
439             }
440 
441             if (createVirtualenv()) {
442                 return;
443             }
444 
445         } catch (IOException | RuntimeException e) {
446             CLog.e(e);
447         }
448 
449         CLog.e(String.format("Failed to create virtualenv at %s.", mVenvDir));
450         throw new TargetSetupError("Error creating virtualenv", mDescriptor);
451     }
452 
getVirtualenvCreationMarkFile()453     protected File getVirtualenvCreationMarkFile() {
454         return new File(mVenvDir, "complete");
455     }
456 
457     /**
458      * Completes the creation of virtualenv.
459      * @return true if the directory is successfully prepared as virutalenv; false otherwise
460      * @throws IOException if completion mark file creation failed.
461      */
createVirtualenv()462     protected boolean createVirtualenv() throws IOException {
463         CLog.d("Creating virtualenv at " + mVenvDir);
464 
465         String[] cmd = new String[] {
466                 "virtualenv", "-p", "python" + mPythonVersion, mVenvDir.getAbsolutePath()};
467 
468         long waitRetryCreate = 5 * SECOND_IN_MSECS;
469 
470         for (int try_count = 0; try_count < PIP_RETRY + 1; try_count++) {
471             if (try_count > 0) {
472                 getRunUtil().sleep(waitRetryCreate);
473             }
474             CommandResult c = getRunUtil().runTimedCmd(3 * MINUTE_IN_MSECS, cmd);
475 
476             if (c.getStatus() != CommandStatus.SUCCESS) {
477                 String message_lower = (c.getStdout() + c.getStderr()).toLowerCase();
478                 if (message_lower.contains("errno 26")
479                         || message_lower.contains("text file busy")) {
480                     // Race condition, retry.
481                     CLog.d("detected the virtualenv path is being created by other process.");
482 
483                     if (createVirtualenv_waitForOtherProcessToCreateVirtualEnv()) {
484                         CLog.d("detected the other process has created virtualenv.");
485                         return true;
486                     }
487                 } else {
488                     // Other error, abort.
489                     CLog.e(String.format("Exit code: %s, stdout: %s, stderr: %s", c.getStatus(),
490                             c.getStdout(), c.getStderr()));
491                     break;
492                 }
493 
494             } else {
495                 mIsDirCreator = true;
496                 getVirtualenvCreationMarkFile().createNewFile();
497                 CLog.d("Succesfully created virtualenv at " + mVenvDir);
498                 return true;
499             }
500         }
501 
502         return false;
503     }
504 
505     /**
506      * Checks whether a host-wise virutanenv directory can be used. If not, creates a empty one.
507      * @param buildInfo
508      * @return true if a host-wise virutanenv directory can be used; false otherwise.
509      * @throws IOException if failed to create empty directory for the virtualenv path.
510      */
checkHostReuseVirtualenv(IBuildInfo buildInfo)511     protected boolean checkHostReuseVirtualenv(IBuildInfo buildInfo) throws IOException {
512         if (mReuse) {
513             String tempDir = System.getProperty("java.io.tmpdir");
514             mVenvDir = new File(tempDir, "vts-virtualenv-" + mPythonVersion);
515             if (mVenvDir.exists()) {
516                 if (createVirtualenv_waitForOtherProcessToCreateVirtualEnv()) {
517                     CLog.d("Using existing virtualenv for version " + mPythonVersion);
518                     return true;
519                 }
520             }
521         } else {
522             mVenvDir = FileUtil.createTempDir("vts-virtualenv-" + mPythonVersion + "-"
523                     + VtsFileUtil.normalizeFileName(buildInfo.getTestTag()) + "_");
524         }
525 
526         return false;
527     }
528 
529     /**
530      * Checks whether a test plan-wise common virtualenv directory can be used.
531      * @param buildInfo
532      * @return true if a test plan-wise virtuanenv directory exists; false otherwise
533      * @throws TargetSetupError
534      */
checkTestPlanLevelVirtualenv(IBuildInfo buildInfo)535     protected boolean checkTestPlanLevelVirtualenv(IBuildInfo buildInfo) throws TargetSetupError {
536         if (mVenvDir == null) {
537             String venvDir = null;
538             switch (getConfiguredPythonVersionMajor()) {
539                 case 2:
540                     venvDir = buildInfo.getBuildAttributes().get(
541                             VtsPythonVirtualenvPreparer.VIRTUAL_ENV);
542                     break;
543                 case 3:
544                     venvDir = buildInfo.getBuildAttributes().get(
545                             VtsPythonVirtualenvPreparer.VIRTUAL_ENV_V3);
546                     break;
547             }
548 
549             if (venvDir != null) {
550                 mVenvDir = new File(venvDir);
551             }
552         }
553 
554         if (mVenvDir != null) {
555             return true;
556         }
557 
558         return false;
559     }
560 
561     /**
562      * Wait for another process to finish creating virtualenv path.
563      * @return true if creation is detected a success; false otherwise.
564      */
createVirtualenv_waitForOtherProcessToCreateVirtualEnv()565     protected boolean createVirtualenv_waitForOtherProcessToCreateVirtualEnv() {
566         long start = System.currentTimeMillis();
567         long totalWaitCheckComplete = 3 * MINUTE_IN_MSECS;
568         long waitRetryCheckComplete = SECOND_IN_MSECS / 2;
569 
570         while (true) {
571             if (getVirtualenvCreationMarkFile().exists()) {
572                 return true;
573             }
574 
575             if (System.currentTimeMillis() - start < totalWaitCheckComplete) {
576                 getRunUtil().sleep(waitRetryCheckComplete);
577             } else {
578                 break;
579             }
580         }
581 
582         return false;
583     }
584 
addDepModule(String module)585     protected void addDepModule(String module) {
586         mDepModules.add(module);
587     }
588 
setRequirementsFile(File f)589     protected void setRequirementsFile(File f) {
590         mRequirementsFile = f;
591     }
592 
593     /**
594      * Get an instance of {@link IRunUtil}.
595      */
596     @VisibleForTesting
getRunUtil()597     protected IRunUtil getRunUtil() {
598         if (mRunUtil == null) {
599             mRunUtil = new RunUtil();
600         }
601         return mRunUtil;
602     }
603 
604     /**
605      * This method recursively deletes a file tree without following symbolic links.
606      *
607      * @param rootPath the path to delete.
608      * @throws IOException if fails to traverse or delete the files.
609      */
recursiveDelete(Path rootPath)610     private static void recursiveDelete(Path rootPath) throws IOException {
611         Files.walkFileTree(rootPath, new SimpleFileVisitor<Path>() {
612             @Override
613             public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
614                     throws IOException {
615                 Files.delete(file);
616                 return FileVisitResult.CONTINUE;
617             }
618             @Override
619             public FileVisitResult postVisitDirectory(Path dir, IOException e) throws IOException {
620                 if (e != null) {
621                     throw e;
622                 }
623                 Files.delete(dir);
624                 return FileVisitResult.CONTINUE;
625             }
626         });
627     }
628 
629     /**
630      * Locally checks whether a pip module is installed.
631      *
632      * This read the installed module list from command "pip list" and check whether the
633      * module in requirement string is installed and its version satisfied.
634      *
635      * Note: This method is only a help method for speed optimization purpose.
636      *       It does not check dependencies of the module.
637      *       It replace dots "." in module name with dash "-".
638      *       If the "pip list" command failed, it will return false and will not throw exception
639      *       It can also only accept one ">=" version requirement string.
640      *       If this method returns false, the requirement should still be checked using pip itself.
641      *
642      * @param requirement such as "numpy", "pip>=9"
643      * @return True if module is installed locally with correct version. False otherwise
644      */
isPipModuleInstalled(String requirement)645     private boolean isPipModuleInstalled(String requirement) {
646         if (mPipInstallList == null) {
647             mPipInstallList = getInstalledPipModules();
648             if (mPipInstallList == null) {
649                 CLog.e("Failed to read local pip install list.");
650                 return false;
651             }
652         }
653 
654         String name;
655         String version = null;
656         if (requirement.contains(">=")) {
657             String[] tokens = requirement.split(">=");
658             if (tokens.length != 2) {
659                 return false;
660             }
661             name = tokens[0];
662             version = tokens[1];
663         } else if (requirement.contains("=") || requirement.contains("<")
664                 || requirement.contains(">")) {
665             return false;
666         } else {
667             name = requirement;
668         }
669 
670         name = name.replaceAll("\\.", "-");
671 
672         if (!mPipInstallList.containsKey(name)) {
673             return false;
674         }
675 
676         // TODO: support other comparison and multiple condition if there's a use case.
677         if (version != null && !isVersionGreaterEqual(mPipInstallList.get(name), version)) {
678             return false;
679         }
680 
681         return true;
682     }
683 
684     /**
685      * Compares whether version string 1 is greater or equal to version string 2
686      * @param version1
687      * @param version2
688      * @return True if the value of version1 >= version2
689      */
isVersionGreaterEqual(String version1, String version2)690     private static boolean isVersionGreaterEqual(String version1, String version2) {
691         version1 = version1.replaceAll("[^0-9.]+", "");
692         version2 = version2.replaceAll("[^0-9.]+", "");
693 
694         String[] tokens1 = version1.split("\\.");
695         String[] tokens2 = version2.split("\\.");
696 
697         int length = Math.max(tokens1.length, tokens2.length);
698         for (int i = 0; i < length; i++) {
699             try {
700                 int token1 = i < tokens1.length ? Integer.parseInt(tokens1[i]) : 0;
701                 int token2 = i < tokens2.length ? Integer.parseInt(tokens2[i]) : 0;
702                 if (token1 < token2) {
703                     return false;
704                 }
705             } catch (NumberFormatException e) {
706                 CLog.e("failed to compare pip module version: %s and %s", version1, version2);
707                 return false;
708             }
709         }
710 
711         return true;
712     }
713 
714     /**
715      * Gets map of installed pip packages and their versions.
716      * @return installed pip packages
717      */
718     private Map<String, String> getInstalledPipModules() {
719         CommandResult res = getRunUtil().runTimedCmd(30 * SECOND_IN_MSECS, getPipPath(), "list");
720         if (res.getStatus() != CommandStatus.SUCCESS) {
721             CLog.e(String.format("Failed to read pip installed list: "
722                             + "Result %s. stdout: %s, stderr: %s",
723                     res.getStatus(), res.getStdout(), res.getStderr()));
724             return null;
725         }
726         String raw = res.getStdout();
727         String[] lines = raw.split("\\r?\\n");
728 
729         TreeMap<String, String> pipInstallList = new TreeMap<>();
730 
731         for (String line : lines) {
732             line = line.trim();
733             if (line.length() == 0 || line.startsWith("Package ") || line.startsWith("-")) {
734                 continue;
735             }
736             String[] tokens = line.split("\\s+");
737             if (tokens.length != 2) {
738                 CLog.e("Error parsing pip installed package list. Line text: " + line);
739                 continue;
740             }
741             pipInstallList.put(tokens[0], tokens[1]);
742         }
743 
744         return pipInstallList;
745     }
746 }
747