Add --save and --save-to options to install

pip/commands/install.py

@@ -6,6 +6,9 @@
 import tempfile
 import shutil
 import warnings
+
+from pip._vendor import six
+
 try:
     import wheel
 except ImportError:
@@ -28,6 +31,11 @@
 logger = logging.getLogger(__name__)
 
 
+if six.PY2:
+    class FileNotFoundError(OSError):
+        pass
+
+
 class InstallCommand(RequirementCommand):
     """
     Install packages from:
@@ -159,6 +167,21 @@ def __init__(self, *args, **kw):
             help="Do not compile py files to pyc",
         )
 
+        cmd_opts.add_option(
+            '--save',
+            action='store_true',
+            dest='save',
+            default=False,
+            help='Add package(s) to requirements file'
+        )
+
+        cmd_opts.add_option(
+            '--save-to',
+            dest='save_to',
+            default='requirements.txt',
+            help='Path to the requirements file'
+        )
+
         cmd_opts.add_option(cmdoptions.use_wheel())
         cmd_opts.add_option(cmdoptions.no_use_wheel())
         cmd_opts.add_option(cmdoptions.no_binary())
@@ -176,6 +199,18 @@ def __init__(self, *args, **kw):
         self.parser.insert_option_group(0, cmd_opts)
 
     def run(self, options, args):
+        if options.save:
+            requirements_fpath = options.save_to
+            if not os.path.isabs(requirements_fpath):
+                requirements_fpath = os.path.join(os.getcwd(),
+                                                  requirements_fpath)
+
+            basedir = os.path.dirname(requirements_fpath)
+            if not os.path.exists(basedir):
+                raise FileNotFoundError(
+                    'Directory for the requirements file doesn\'t exist: '
+                    '{basedir}.'.format(basedir=basedir))
+
         cmdoptions.resolve_wheel_no_use_binary(options)
         cmdoptions.check_install_build_global(options)
 
@@ -394,4 +429,39 @@ def run(self, options, args):
                         target_item_dir
                     )
             shutil.rmtree(temp_target_dir)
+
+        if options.save:
+            lines = []
+            if os.path.exists(requirements_fpath):
+                saved_packages = {}
+                with open(requirements_fpath, 'r') as requirements_file:
+                    req_lines = requirements_file.readlines()
+                    for line_number, line in enumerate(req_lines):
+                        lines.append(line)
+                        if '==' in line:
+                            name, _ = line.rstrip().split('==')
+                            saved_packages[name] = line_number
+
+            for requirement in requirement_set.requirements.values():
+                if not requirement.comes_from:
+                    pkg_name = requirement.name
+                    pkg_version = requirement.installed_version
+                    pkg_output_line = '{pkg_name}=={pkg_version}\n'.format(
+                        pkg_name=pkg_name,
+                        pkg_version=pkg_version)
+
+                    if len(lines) == 0:
+                        lines.append(pkg_output_line)
+                        continue
+
+                    if pkg_name in saved_packages:
+                        # if same version, nothing bad will happen
+                        line_number = saved_packages[pkg_name]
+                        lines[line_number] = pkg_output_line
+                    else:
+                        lines.append(pkg_output_line)
+
+            with open(requirements_fpath, 'w') as requirements_file:
+                requirements_file.writelines(lines)
+
         return requirement_set

tests/functional/test_install_save.py

@@ -0,0 +1,56 @@
+def test_save_installed_package_in_requirements(script):
+    freezed_requests = 'requests==1.0.0'
+    script.pip('install', '--save', freezed_requests)
+
+    requirements_fpath = script.cwd.join('requirements.txt')
+    with open(requirements_fpath, 'r') as requirements_file:
+        assert freezed_requests in requirements_file.read()
+
+
+def test_save_can_specify_save_to_file(script, tmpdir):
+    requirements_fpath = tmpdir.join('base.txt')
+    freezed_requests = 'requests==1.0.0'
+    script.pip('install', '--save',
+               '--save-to', requirements_fpath, freezed_requests)
+
+    with open(requirements_fpath, 'r') as requirements_file:
+        assert freezed_requests in requirements_file.read()
+
+
+def test_save_fails_if_base_directory_not_exists(script, tmpdir):
+    requirements_fpath = tmpdir.join('base/requirements.txt')
+    freezed_requests = 'requests==1.0.0'
+    result = script.pip('install', '--save',
+                        '--save-to', requirements_fpath, freezed_requests,
+                        expect_error=True)
+    assert result.returncode == 2
+    assert 'FileNotFoundError' in result.stderr
+
+
+def test_save_only_updates_package_line_if_already_exists(script, tmpdir):
+    requirements_fpath = tmpdir.join('requirements.txt')
+    existing_requests_line = 'requests==1.0.0'
+    with open(requirements_fpath, 'w') as requirements_file:
+        requirements_file.write(existing_requests_line)
+
+    updated_requests_line = 'requests==1.0.1'
+    script.pip('install', '--save',
+               '--save-to', requirements_fpath, updated_requests_line)
+
+    with open(requirements_fpath, 'r') as requirements_file:
+        assert updated_requests_line in requirements_file.read()
+
+
+def test_save_preserves_original_content_of_requirements_file(script, tmpdir):
+    requirements_fpath = tmpdir.join('dev-requirements.txt')
+    original_line = '-r requirements.txt'
+    with open(requirements_fpath, 'w') as requirements_file:
+        requirements_file.write(original_line + '\n')
+
+    requests_line = 'requests==1.0.1'
+    script.pip('install', '--save',
+               '--save-to', requirements_fpath, requests_line)
+    with open(requirements_fpath, 'r') as requirements_file:
+        file_contents = requirements_file.read()
+        assert original_line in file_contents
+        assert requests_line in file_contents