Fixed several potential issues and cleaned up code.

max-krichenbauer · max-krichenbauer · commit ed27d00713d4 · 2019-11-21T14:50:05.000+09:00
diff --git a/ament_virtualenv/ament_virtualenv/build_venv.py b/ament_virtualenv/ament_virtualenv/build_venv.py
@@ -81,6 +81,7 @@ def main(argv=sys.argv[1:]):
     )
     parser.add_argument(
         '--python-version',
+        required=True,
         help="Build the virtualenv with which python major version."
     )
     parser.add_argument(
@@ -100,22 +101,28 @@ def main(argv=sys.argv[1:]):
         type=str,
         help="Extra pip args for install."
     )
-
     args, unknown = parser.parse_known_args()
+    return build_venv(
+        root_dir=args.root_dir,
+        python_version=args.python_version,
+        requirements_filename=args.requirements,
+        use_system_packages=args.use_system_packages,
+        extra_pip_args=args.extra_pip_args[1:-1],
+        retries=args.retries
+    )
 
-    root_dir = os.path.realpath(args.root_dir)
-
+def build_venv(root_dir, python_version, requirements_filename, use_system_packages=False, extra_pip_args="", retries=0):
+    root_dir = os.path.realpath(root_dir)
+    python_executable = find_python(python_version)
     os.environ['DH_VIRTUALENV_INSTALL_ROOT'] = os.path.dirname(root_dir)
 
-    python_executable = find_python(args.python_version)
-
     deploy = Deployment(
         package=os.path.basename(root_dir),
-        requirements_filename=args.requirements,
+        requirements_filename=requirements_filename,
         upgrade_pip=True,
-        use_system_packages=args.use_system_packages,
+        use_system_packages=use_system_packages,
         python=python_executable,
-        extra_pip_arg=args.extra_pip_args[1:-1].split(' '),
+        extra_pip_arg=extra_pip_args.split(' '),
         log_file=None,
         builtin_venv=check_module(python_executable, 'venv'),
         builtin_pip=check_module(python_executable, 'pip'),
@@ -127,7 +134,7 @@ def main(argv=sys.argv[1:]):
             print('Generating virtualenv in {}'.format(deploy.package_dir))
             deploy.create_virtualenv()
 
-            print('Installing requirements')
+            print('Installing requirements from {}'.format(deploy.requirements_filename))
             deploy.install_dependencies()
 
             print('Fixing virtualenv root to {}'.format(deploy.virtualenv_install_dir))
@@ -143,8 +150,8 @@ def main(argv=sys.argv[1:]):
                 # Remove local folder
                 shutil.rmtree(local_dir)
         except Exception as e:
-            args.retries -= 1
-            if args.retries >= 0:
+            retries -= 1
+            if retries >= 0:
                 print("Error, clearing virtualenv and retrying: {}".format(e), file=sys.stderr)
                 try:
                     shutil.rmtree(root_dir)
@@ -153,12 +160,11 @@ def main(argv=sys.argv[1:]):
                 continue
             else:
                 raise
-
         break
 
     return 0
 #
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/ament_virtualenv/ament_virtualenv/combine_requirements.py b/ament_virtualenv/ament_virtualenv/combine_requirements.py
@@ -97,4 +97,4 @@ def main(argv=sys.argv[1:]):
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/ament_virtualenv/ament_virtualenv/glob_requirements.py b/ament_virtualenv/ament_virtualenv/glob_requirements.py
@@ -34,14 +34,13 @@
 
 
 def find_in_workspaces(project, file):
-    paths = os.environ.get('AMENT_PREFIX_PATH')
+    paths = os.environ.get('COLCON_PREFIX_PATH')
     if not paths:
         return None
     paths = paths.split(os.pathsep)
     workspaces = []
     for path in paths:
-        if '/install/' in path:
-            workspaces.append(path)
+        workspaces.append(os.path.join(path, project))
     # should be at share/ament_virtualenv/
     search_dirs = ['etc', 'include', 'libexec', 'share']
     if 'libexec' in search_dirs:
@@ -126,18 +125,18 @@ def glob_requirements(package_name, no_deps):
                 for dependency in dependencies:
                     package_queue.put(dependency.name)
 
-    print(';'.join(requirements_list))
-    return 0
+    return ';'.join(requirements_list)
 
 
 def main(argv=sys.argv[1:]):
     parser = argparse.ArgumentParser()
     parser.add_argument('--package-name', type=str, required=True)
     parser.add_argument('--no-deps', action="store_true")
     args, unknown = parser.parse_known_args()
-    return glob_requirements(**vars(args))
+    print(glob_requirements(**vars(args)))
+    return 0
 #
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/ament_virtualenv/ament_virtualenv/install.py b/ament_virtualenv/ament_virtualenv/install.py
@@ -26,13 +26,23 @@
 import sys
 import subprocess
 import shutil
-import importlib
-import unittest
 import psutil
+import argparse
 
+try:
+    from ament_virtualenv.glob_requirements import glob_requirements
+    from ament_virtualenv.combine_requirements import combine_requirements
+    from ament_virtualenv.build_venv import build_venv
+    ament_virtualenv_import_failed = False
+except:
+    ament_virtualenv_import_failed = True
+#
 
 def find_program(name='build_venv.py', package='ament_virtualenv'):
     '''
+    Helper function to find the modules that are part of this
+    package (glob_requirements, combine_requirements, build_venv),
+    in cases where a direct import failes due to python path issues.
     '''
     ament_prefix_path = os.environ.get("AMENT_PREFIX_PATH")
     if not ament_prefix_path:
@@ -52,58 +62,90 @@ def find_program(name='build_venv.py', package='ament_virtualenv'):
 def install_venv(install_base, package_name, python_version='2'):
     venv_install_dir = os.path.join(install_base, 'venv')
     bin_dir = os.path.join(install_base, 'bin')
-	# 
+    # 
     # Build the virtual environment
     python = shutil.which("python")
     if not python:
         print("ERROR: Failed to locate python", file=sys.stderr)
-        return
+        return 1
 
     # glob_requirements --package-name ament_cmake_haros
-    glob_requirements = find_program(name='glob_requirements.py', package='ament_virtualenv')
-    if not glob_requirements:
-        print("ERROR: Failed to locate glob_requirements", file=sys.stderr)
-        return
-    cmd = [
-        python, 
-        glob_requirements,
-        '--package-name',
-        package_name
-    ]
-    requirements_list = subprocess.check_output(cmd)
-    requirements_list = requirements_list.decode("utf-8").strip()
-
-    # combine_requirements --requirements-list a/requirements.txt;b/requirements.txt --output-file x/generated_requirements.txt
-    combine_requirements = find_program(name='combine_requirements.py', package='ament_virtualenv')
-    if not combine_requirements:
-        print("ERROR: Failed to locate combine_requirements", file=sys.stderr)
-        return
+    if ament_virtualenv_import_failed == True:
+        # Fallback: try to find the command line script and run it
+        glob_requirements_py = find_program(name='glob_requirements.py', package='ament_virtualenv')
+        if not glob_requirements_py:
+            print("ERROR: Failed to locate glob_requirements", file=sys.stderr)
+            return 1
+        cmd = [
+            python, 
+            glob_requirements_py,
+            '--package-name',
+            package_name
+        ]
+        requirements_list = subprocess.check_output(cmd)
+        requirements_list = requirements_list.decode("utf-8").strip()
+    else:
+        # Use the module directly
+        requirements_list = glob_requirements(package_name=package_name, no_deps=False)
+    # ^ glob_requirements
+    #
+    # combine_requirements --requirements-list a/requirements.txt;b/requirements.txt
+    #                      --output-file x/generated_requirements.txt
     generated_requirements = '/tmp/test_ament_virtualenv-generated_requirements.txt'
-    cmd = [
-        python,
-        combine_requirements,
-        '--requirements-list',
-        requirements_list,
-        '--output-file',
-        generated_requirements
-    ]
-    subprocess.check_output(cmd)
-
-    build_venv = find_program(name='build_venv.py', package='ament_virtualenv')
-    if not build_venv:
-        print("ERROR: Failed to locate build_venv", file=sys.stderr)
-        return
-    cmd = [
-        python,
-        build_venv,
-        '--root-dir', venv_install_dir,
-        '--requirements', generated_requirements,
-        '--retries', '3',
-        '--python-version', python_version,
-        # '--use-system-packages',
-        '--extra-pip-args', '\"-qq\"',
-    ]
-    ret = subprocess.check_output(cmd)    
+    if ament_virtualenv_import_failed:
+        combine_requirements_py = find_program(name='combine_requirements.py', package='ament_virtualenv')
+        if not combine_requirements_py:
+            print("ERROR: Failed to locate combine_requirements", file=sys.stderr)
+            return 1
+        cmd = [
+            python,
+            combine_requirements_py,
+            '--requirements-list',
+            requirements_list,
+            '--output-file',
+            generated_requirements
+        ]
+        subprocess.check_output(cmd)
+    else:
+        requirements_files = []
+        for requirements_file in requirements_list.split(';'):
+            requirements_files.append(open(requirements_file, 'r'))
+        generated_requirements_file = open(generated_requirements, 'w')
+        combine_requirements(
+            requirements_list=requirements_files,
+            output_file=generated_requirements_file
+        )
+        for requirements_file in requirements_files:
+            requirements_file.close()
+        generated_requirements_file.close()
+    # ^ combine_requirements
+    #
+    if ament_virtualenv_import_failed:
+        build_venv_py = find_program(name='build_venv.py', package='ament_virtualenv')
+        if not build_venv_py:
+            print("ERROR: Failed to locate build_venv", file=sys.stderr)
+            return 1
+        cmd = [
+            python,
+            build_venv_py,
+            '--root-dir', venv_install_dir,
+            '--requirements', generated_requirements,
+            '--retries', '3',
+            '--python-version', python_version,
+            # '--use-system-packages',
+            '--extra-pip-args', '\"-qq\"',
+        ]
+        ret = subprocess.check_output(cmd)
+    else:
+        ret = build_venv(
+            root_dir=venv_install_dir,
+            python_version=python_version,
+            requirements_filename=generated_requirements,
+            use_system_packages=False,
+            extra_pip_args="-qq",
+            retries=3
+        )
+        print('XXXXXXXX build_venv returned ' + str(ret))
     # 
     # Wrapper shell executables we installed
     for bin_file in os.listdir(bin_dir):
@@ -114,6 +156,7 @@ def install_venv(install_base, package_name, python_version='2'):
         if not os.path.isfile(bin_path):
             continue
         os.rename(bin_path, bin_path+'-venv')
+        venv_rel_path = os.path.relpath(venv_install_dir, bin_dir)
         # create new file with the name of the previous file
         with open(bin_path, "w") as f:
             f.write("#!/usr/bin/python3\n")
@@ -123,16 +166,28 @@ def install_venv(install_base, package_name, python_version='2'):
             f.write("if __name__ == '__main__':\n")
             f.write("    dir_path = os.path.dirname(os.path.realpath(__file__))\n")
             f.write("    bin_path = os.path.join(dir_path, '" + bin_file + "-venv')\n")
-            f.write("    cmd = '" + venv_install_dir + "/bin/python ' + bin_path\n")
+            f.write("    vpy_path = os.path.abspath(os.path.join(dir_path, '" + venv_rel_path +"'))\n")
+            f.write("    vpy_path = os.path.join(vpy_path, 'bin', 'python')\n")
+            f.write("    cmd = vpy_path + ' ' + bin_path\n")
             f.write("    sys.exit(subprocess.call(cmd, shell=True))\n")
         # change file permissions to executable
         st = os.stat(bin_path)
         os.chmod(bin_path, st.st_mode | stat.S_IEXEC | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
-
+    #
+    return 0
 
 def main(argv=sys.argv[1:]):
-	return 0 # todo
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--install-base', required=True)
+    parser.add_argument('--package-name', required=True)
+    parser.add_argument('--python-version', required=True)
+    args, unknown = parser.parse_known_args()
+    return install_venv(
+        install_base=args.install_base,
+        package_name=args.package_name,
+        python_version=args.python_version
+    )
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/test_ament_virtualenv/test_ament_virtualenv/test_ament_virtualenv.py b/test_ament_virtualenv/test_ament_virtualenv/test_ament_virtualenv.py
@@ -25,20 +25,34 @@ def main(args=None):
     # 1: Test if we're in a virtual environment at all.
     is_in_venv = hasattr(sys, 'real_prefix')
     if not is_in_venv:
-        print("[test_ament_virtualenv] FAILURE: Python virtual environment not activated.")
+        print(
+            "[test_ament_virtualenv] "
+            "FAILURE: Python virtual environment not activated."
+        )
         return 1
     # 2: Test the Python version.
     if sys.version_info.major != 2:
-        print("[test_ament_virtualenv] FAILURE: Wrong Python version.")
+        print(
+            "[test_ament_virtualenv] "
+            "FAILURE: Wrong Python version."
+        )
         return 1
     # 3: Test if proper requirements have been installed
     try:
         requests = importlib.import_module("requests")
         if requests.__version__ != '2.20.1':
-            print("[test_ament_virtualenv] FAILURE: Requirements not provided correctly.")
+            print(
+                "[test_ament_virtualenv] "
+                "FAILURE: Requirements not provided correctly "
+                "(expected 'requests==2.20.1', found "+requests.__version__+")"
+            )
             return 1
     except:
-        print("[test_ament_virtualenv] FAILURE: Requirements not provided.")
+        print(
+            "[test_ament_virtualenv] "
+            "FAILURE: Requirements not provided "
+            "(expected 'requests' to be present but could not find it)"
+        )
         return 1
     print("[test_ament_virtualenv] SUCCESS: All checks passed.")
     return 0

Original file line number	Diff line number	Diff line change
`@@ -97,4 +97,4 @@ def main(argv=sys.argv[1:]):`
`97`	`97`
`98`	`98`
`99`	`99`	`if __name__ == "__main__":`
`100`		`- main()`
	`100`	`+ sys.exit(main())`