Move re-editing logic from cli to utils.

pir/cli.py

@@ -1,14 +1,12 @@
 # -*- coding: utf-8 -*-
 from __future__ import division, print_function, absolute_import, unicode_literals
 
-import re
-
 import pip
 import click
 
 from . import echoes
 from .groups import AliasedGroup
-from .utils import classify_installed_or_not
+from .utils import classify_installed_or_not, re_edit_requirements
 
 
 @click.group(cls=AliasedGroup)
@@ -92,19 +90,7 @@ def uninstall(packages, requirement):
         else:
             uninstalled_packages.append(pkg)
 
-    output = []
-    pattern = re.compile('^[\w0-9\-.]+')
-    for line in requirement.readlines():
-        match_obj = pattern.match(line)
-        if match_obj:
-            matched_pkg = match_obj.group().lower()
-            if matched_pkg not in packages:
-                output.append(line)
-        else:
-            output.append(line)
-    content = '\n'.join(output)
-    if output[-1] != '\n':
-        output += '\n'
+    content = re_edit_requirements(requirement.readlines(), packages)
     with open(requirement.name, 'w') as f:
         f.write(content)
     msg = "Remove the following packages from {requirement}: {packages}"

pir/utils.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 from __future__ import division, print_function, absolute_import, unicode_literals
 
+import re
+
 import pip
 
 
@@ -13,3 +15,22 @@ def classify_installed_or_not(packages):
     will_install = set(packages) - installed
     need_upgrade = set(packages) & installed
     return will_install, need_upgrade
+
+
+def re_edit_requirements(lines, will_remove):
+    """
+    Re-Edit requirements file
+
+    :param lines: requirement.readlines()'s return value.
+    :param will_remove: Will be removed packages.
+    :return: The contents of the file after re-editing.
+    """
+    pattern = re.compile('^[\w0-9\-.]+')
+    re_editing = []
+    for line in lines:
+        matched = pattern.match(line)
+        if not matched or matched.group().lower() not in will_remove:
+            re_editing.append(line)
+    if re_editing[-1] != '\n':
+        re_editing.append('\n')
+    return ''.join(re_editing)

tests/test_utils.py

@@ -14,3 +14,19 @@ def mockreturn():
 
     rv = classify_installed_or_not(packages)
     assert rv == expected
+
+
+def test_re_edit_requirements():
+    from pir.utils import re_edit_requirements
+    before = [
+        "# requirements.in\n",
+        "flask\n",
+        "pytest\n",
+        "\n",
+    ]
+    after = """# requirements.in
+pytest
+
+"""
+    content = re_edit_requirements(before, ['flask'])
+    assert content == after