Support Exclude Package with custom unsafe packages

piptools/_compat/pip_compat.py

@@ -3,18 +3,18 @@
 
 import pip
 from pip._internal.index.package_finder import PackageFinder
-from pip._internal.network.session import PipSession
 from pip._internal.req import InstallRequirement
 from pip._internal.req import parse_requirements as _parse_requirements
 from pip._internal.req.constructors import install_req_from_parsed_requirement
 from pip._vendor.packaging.version import parse as parse_version
+from pip._vendor.requests.sessions import Session
 
 PIP_VERSION = tuple(map(int, parse_version(pip.__version__).base_version.split(".")))
 
 
 def parse_requirements(
     filename: str,
-    session: PipSession,
+    session: Session,
     finder: Optional[PackageFinder] = None,
     options: Optional[optparse.Values] = None,
     constraint: bool = False,

piptools/repositories/local.py

@@ -4,9 +4,9 @@
 
 from pip._internal.index.package_finder import PackageFinder
 from pip._internal.models.candidate import InstallationCandidate
+from pip._internal.network.session import PipSession
 from pip._internal.req import InstallRequirement
 from pip._internal.utils.hashes import FAVORITE_HASH
-from pip._vendor.requests import Session
 
 from piptools.utils import as_tuple, key_from_ireq, make_install_requirement
 
@@ -58,7 +58,7 @@ def finder(self) -> PackageFinder:
         return self.repository.finder
 
     @property
-    def session(self) -> Session:
+    def session(self) -> PipSession:
         return self.repository.session
 
     def clear_caches(self) -> None:

piptools/resolver.py

@@ -113,6 +113,8 @@ def __init__(
         prereleases: Optional[bool] = False,
         clear_caches: bool = False,
         allow_unsafe: bool = False,
+        unsafe_packages: Optional[Set[str]] = None,
+        allow_unsafe_recursive: bool = False,
     ) -> None:
         """
         This class resolves a given set of constraints (a collection of
@@ -127,6 +129,8 @@ def __init__(
         self.clear_caches = clear_caches
         self.allow_unsafe = allow_unsafe
         self.unsafe_constraints: Set[InstallRequirement] = set()
+        self.unsafe_packages = unsafe_packages or UNSAFE_PACKAGES
+        self.allow_unsafe_recursive = allow_unsafe_recursive
 
     @property
     def constraints(self) -> Set[InstallRequirement]:
@@ -189,17 +193,24 @@ def resolve(self, max_rounds: int = 10) -> Set[InstallRequirement]:
         # Filter out unsafe requirements.
         self.unsafe_constraints = set()
         if not self.allow_unsafe:
-            # reverse_dependencies is used to filter out packages that are only
-            # required by unsafe packages. This logic is incomplete, as it would
-            # fail to filter sub-sub-dependencies of unsafe packages. None of the
-            # UNSAFE_PACKAGES currently have any dependencies at all (which makes
-            # sense for installation tools) so this seems sufficient.
-            reverse_dependencies = self.reverse_dependencies(results)
+            if not self.allow_unsafe_recursive:
+                # reverse_dependencies is used to filter out packages that are only
+                # required by unsafe packages. This logic is incomplete, as it would
+                # fail to filter sub-sub-dependencies of unsafe packages. None of the
+                # UNSAFE_PACKAGES currently have any dependencies at all (which makes
+                # sense for installation tools) so this seems sufficient.
+                reverse_dependencies = self.reverse_dependencies(results)
+
+                for req in results.copy():
+                    required_by = reverse_dependencies.get(req.name.lower(), set())
+                    if required_by and all(
+                        name in self.unsafe_packages for name in required_by
+                    ):
+                        self.unsafe_constraints.add(req)
+                        results.remove(req)
+
             for req in results.copy():
-                required_by = reverse_dependencies.get(req.name.lower(), set())
-                if req.name in UNSAFE_PACKAGES or (
-                    required_by and all(name in UNSAFE_PACKAGES for name in required_by)
-                ):
+                if req.name in self.unsafe_packages:
                     self.unsafe_constraints.add(req)
                     results.remove(req)
 

piptools/scripts/compile.py

@@ -236,6 +236,20 @@ def _get_default_option(option_name: str) -> Any:
     default=True,
     help="Add options to generated file",
 )
+@click.option(
+    "--unsafe-packages",
+    multiple=True,
+    help="List of packages to consider unsafe. Replaces "
+    f"{', '.join(sorted(UNSAFE_PACKAGES))}; may be used more than once",
+)
+@click.option(
+    "--allow-unsafe-recursive/--no-allow-unsafe-recursive",
+    is_flag=True,
+    default=False,
+    help="Determines whether dependencies that solely belong to unsafe packages "
+    "should be treated as unsafe. Default is false. This only covers direct dependencies. "
+    "Dependencies of dependencies of unsafe packages will not be marked unsafe.",
+)
 def cli(
     ctx: click.Context,
     verbose: int,
@@ -269,6 +283,8 @@ def cli(
     pip_args_str: Optional[str],
     emit_index_url: bool,
     emit_options: bool,
+    unsafe_packages: Tuple[str, ...],
+    allow_unsafe_recursive: bool,
 ) -> None:
     """Compiles requirements.txt from requirements.in specs."""
     log.verbosity = verbose - quiet
@@ -462,6 +478,8 @@ def cli(
             cache=DependencyCache(cache_dir),
             clear_caches=rebuild,
             allow_unsafe=allow_unsafe,
+            unsafe_packages=set(unsafe_packages),
+            allow_unsafe_recursive=allow_unsafe_recursive,
         )
         results = resolver.resolve(max_rounds=max_rounds)
         hashes = resolver.resolve_hashes(results) if generate_hashes else None

tests/test_resolver.py

@@ -249,6 +249,78 @@ def test_resolver__allows_unsafe_deps(
     assert output == {str(line) for line in expected}
 
 
+@pytest.mark.parametrize(
+    (
+        "input",
+        "expected",
+        "unsafe_packages",
+        "allow_unsafe_recursive",
+        "unsafe_constraints",
+    ),
+    (
+        (
+            ["fake-piptools-test-with-pinned-deps==0.1"],
+            {
+                "fake-piptools-test-with-pinned-deps==0.1",
+                "pytz==2016.4 (from celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "billiard==3.3.0.23 (from "
+                "celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "celery==3.1.18 (from fake-piptools-test-with-pinned-deps==0.1)",
+            },
+            {"kombu"},
+            False,
+            {
+                "kombu==3.0.35 (from celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "anyjson==0.3.3 (from "
+                "kombu==3.0.35->celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "amqp==1.4.9 (from "
+                "kombu==3.0.35->celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+            },
+        ),
+        (
+            ["fake-piptools-test-with-pinned-deps==0.1"],
+            {
+                "fake-piptools-test-with-pinned-deps==0.1",
+                "pytz==2016.4 (from celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "billiard==3.3.0.23 (from "
+                "celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "celery==3.1.18 (from fake-piptools-test-with-pinned-deps==0.1)",
+                "anyjson==0.3.3 (from "
+                "kombu==3.0.35->celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+                "amqp==1.4.9 (from "
+                "kombu==3.0.35->celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+            },
+            {"kombu"},
+            True,
+            {
+                "kombu==3.0.35 (from celery==3.1.18->fake-piptools-test-with-pinned-deps==0.1)",
+            },
+        ),
+    ),
+)
+def test_resolver__custom_unsafe_deps(
+    resolver,
+    from_line,
+    input,
+    expected,
+    unsafe_packages,
+    allow_unsafe_recursive,
+    unsafe_constraints,
+):
+    input = [line if isinstance(line, tuple) else (line, False) for line in input]
+    input = [from_line(req[0], constraint=req[1]) for req in input]
+    resolver = resolver(
+        input,
+        unsafe_packages=unsafe_packages,
+        allow_unsafe_recursive=allow_unsafe_recursive,
+    )
+    output = resolver.resolve()
+    output = {str(line) for line in output}
+
+    assert output == expected
+    assert {str(line) for line in resolver.unsafe_constraints} == unsafe_constraints
+
+
 def test_resolver__max_number_rounds_reached(resolver, from_line):
     """
     Resolver should raise an exception if max round has been reached.