From 436b16bf3f90f3af4e9eb808777c53da3597ee87 Mon Sep 17 00:00:00 2001
From: ianpye <ianpangye@gmail.com>
Date: Fri, 26 Jan 2024 18:39:06 +0800
Subject: [PATCH] Add flags to disable pypi.org and add own private pip
 repositories

---
 conda_lock/conda_lock.py          | 24 ++++++++++++++++++++++++
 conda_lock/src_parser/__init__.py | 10 ++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/conda_lock/conda_lock.py b/conda_lock/conda_lock.py
index 48b1d7171..8d3784be6 100644
--- a/conda_lock/conda_lock.py
+++ b/conda_lock/conda_lock.py
@@ -270,6 +270,8 @@ def make_lock_files(  # noqa: C901
     metadata_yamls: Sequence[pathlib.Path] = (),
     with_cuda: Optional[str] = None,
     strip_auth: bool = False,
+    disable_pypi_requests: bool = False,
+    private_pip_repositories: Optional[Sequence[str]] = None,
 ) -> None:
     """
     Generate a lock file from the src files provided
@@ -324,6 +326,8 @@ def make_lock_files(  # noqa: C901
         channel_overrides=channel_overrides,
         platform_overrides=platform_overrides,
         required_categories=required_categories if filter_categories else None,
+        pip_repository_overrides=private_pip_repositories,
+        allow_pypi_requests_overrides=not disable_pypi_requests,
     )
 
     # Load existing lockfile if it exists
@@ -1132,6 +1136,8 @@ def run_lock(
     metadata_choices: AbstractSet[MetadataOption] = frozenset(),
     metadata_yamls: Sequence[pathlib.Path] = (),
     strip_auth: bool = False,
+    disable_pypi_requests: bool = False,
+    private_pip_repositories: Optional[Sequence[str]] = None,
 ) -> None:
     if len(environment_files) == 0:
         environment_files = handle_no_specified_source_files(lockfile_path)
@@ -1158,6 +1164,8 @@ def run_lock(
         metadata_choices=metadata_choices,
         metadata_yamls=metadata_yamls,
         strip_auth=strip_auth,
+        disable_pypi_requests=disable_pypi_requests,
+        private_pip_repositories=private_pip_repositories,
     )
 
 
@@ -1321,6 +1329,18 @@ def main() -> None:
     type=click.Path(),
     help="YAML or JSON file(s) containing structured metadata to add to metadata section of the lockfile.",
 )
+@click.option(
+    "--disable-pypi-requests",
+    is_flag=True,
+    default=False,
+    help="Disable requests to pypi.org",
+)
+@click.option(
+    "-r",
+    "--private_pip_repositories",
+    multiple=True,
+    help="Add private pip repositories",
+)
 @click.pass_context
 def lock(
     ctx: click.Context,
@@ -1335,6 +1355,8 @@ def lock(
     filename_template: str,
     lockfile: Optional[PathLike],
     strip_auth: bool,
+    disable_pypi_requests: bool,
+    private_pip_repositories: Sequence[str],
     extras: Sequence[str],
     filter_categories: bool,
     check_input_hash: bool,
@@ -1408,6 +1430,8 @@ def lock(
         metadata_choices=metadata_enum_choices,
         metadata_yamls=[pathlib.Path(path) for path in metadata_yamls],
         strip_auth=strip_auth,
+        disable_pypi_requests=disable_pypi_requests,
+        private_pip_repositories=private_pip_repositories,
     )
     if strip_auth:
         with tempfile.TemporaryDirectory() as tempdir:
diff --git a/conda_lock/src_parser/__init__.py b/conda_lock/src_parser/__init__.py
index 7fa4f3b98..f532daac0 100644
--- a/conda_lock/src_parser/__init__.py
+++ b/conda_lock/src_parser/__init__.py
@@ -79,6 +79,7 @@ def make_lock_spec(
     pip_repository_overrides: Optional[Sequence[str]] = None,
     platform_overrides: Optional[Sequence[str]] = None,
     required_categories: Optional[AbstractSet[str]] = None,
+    allow_pypi_requests_overrides: Optional[bool] = None,
 ) -> LockSpecification:
     """Generate the lockfile specs from a set of input src_files.  If required_categories is set filter out specs that do not match those"""
     platforms = (
@@ -124,11 +125,16 @@ def dep_has_category(d: Dependency, categories: AbstractSet[str]) -> bool:
             ]
             for platform, dependencies in aggregated_lock_spec.dependencies.items()
         }
-
+    if allow_pypi_requests_overrides is None:
+        allow_pypi_requests = aggregated_lock_spec.allow_pypi_requests
+    else:
+        allow_pypi_requests = (
+            allow_pypi_requests_overrides and aggregated_lock_spec.allow_pypi_requests
+        )
     return LockSpecification(
         dependencies=dependencies,
         channels=channels,
         pip_repositories=pip_repositories,
         sources=aggregated_lock_spec.sources,
-        allow_pypi_requests=aggregated_lock_spec.allow_pypi_requests,
+        allow_pypi_requests=allow_pypi_requests,
     )