diff --git a/eng/tools/azure-sdk-tools/azpysdk/apistub.py b/eng/tools/azure-sdk-tools/azpysdk/apistub.py index 31faae7caaf9..659a8833f25c 100644 --- a/eng/tools/azure-sdk-tools/azpysdk/apistub.py +++ b/eng/tools/azure-sdk-tools/azpysdk/apistub.py @@ -69,6 +69,12 @@ def register( default=None, help="Destination directory for generated API stub files.", ) + p.add_argument( + "--generate-from-pypi", + dest="generate_from_pypi", + default=None, + help="Generate the stub from this released PyPI version instead of local source code.", + ) p.add_argument( "--install-deps", dest="install_deps", @@ -95,6 +101,31 @@ def ensure_apistub_dependencies(self, executable: str, package_dir: str, staging package_dir, ) + def download_pypi_wheel(self, executable: str, package_name: str, version: str, staging_directory: str) -> str: + """Download a released wheel from PyPI into the staging directory and return its path.""" + logger.info(f"Downloading {package_name}=={version} from PyPI.") + self.run_venv_command( + executable, + [ + "-m", + "pip", + "download", + f"{package_name}=={version}", + "--no-deps", + "--only-binary=:all:", + "-d", + staging_directory, + ], + cwd=staging_directory, + check=True, + ) + found_whl = find_whl(staging_directory, package_name, version) + if not found_whl: + raise FileNotFoundError( + f"No wheel found for package {package_name} version {version} after downloading from PyPI." + ) + return os.path.join(staging_directory, found_whl) + def run(self, args: argparse.Namespace) -> int: """Run the apistub check command.""" logger.info("Running apistub check...") @@ -133,23 +164,28 @@ def run(self, args: argparse.Namespace) -> int: logger.error(f"Failed to install APIView dependencies: {e}") return getattr(e, "returncode", 1) - if not os.getenv("PREBUILT_WHEEL_DIR"): - create_package_and_install( - distribution_directory=staging_directory, - target_setup=package_dir, - skip_install=True, - cache_dir=None, - work_dir=staging_directory, - force_create=False, - package_type="wheel", - pre_download_disabled=False, - python_executable=executable, - ) + generate_from_pypi = getattr(args, "generate_from_pypi", None) + + if generate_from_pypi: + pkg_path = self.download_pypi_wheel(executable, package_name, generate_from_pypi, staging_directory) + else: + if not os.getenv("PREBUILT_WHEEL_DIR"): + create_package_and_install( + distribution_directory=staging_directory, + target_setup=package_dir, + skip_install=True, + cache_dir=None, + work_dir=staging_directory, + force_create=False, + package_type="wheel", + pre_download_disabled=False, + python_executable=executable, + ) + pkg_path = get_package_wheel_path(package_dir) if install_deps: self.pip_freeze(executable) - pkg_path = get_package_wheel_path(package_dir) pkg_path = os.path.abspath(pkg_path) out_token_path = os.path.abspath(getattr(args, "dest_dir", None) or package_dir) diff --git a/eng/tools/azure-sdk-tools/tests/test_apistub.py b/eng/tools/azure-sdk-tools/tests/test_apistub.py index 4dd4d5bb017a..4799383577c4 100644 --- a/eng/tools/azure-sdk-tools/tests/test_apistub.py +++ b/eng/tools/azure-sdk-tools/tests/test_apistub.py @@ -8,6 +8,24 @@ from azpysdk.apistub import apistub, get_package_wheel_path, get_cross_language_mapping_path + +def _build_parser(): + parser = argparse.ArgumentParser(prog="azpysdk") + subparsers = parser.add_subparsers(title="commands", dest="command") + apistub().register(subparsers) + return parser + + +class TestApistubRegistration: + def test_generate_from_pypi_flag_sets_version(self): + parser = _build_parser() + + args = parser.parse_args(["apistub", "--generate-from-pypi", "1.0.0"]) + + assert args.command == "apistub" + assert args.generate_from_pypi == "1.0.0" + + # ── get_package_wheel_path() ───────────────────────────────────────────── @@ -79,6 +97,7 @@ def _make_args( isolate=False, install_deps=False, dest_dir=None, + generate_from_pypi=None, ): return argparse.Namespace( target=".", @@ -88,6 +107,7 @@ def _make_args( token_file=token_file, install_deps=install_deps, dest_dir=dest_dir, + generate_from_pypi=generate_from_pypi, ) @patch( @@ -418,3 +438,70 @@ def fake_apistub_run(exe, cmds, **kwargs): assert "--skip-pylint" not in captured_cmds[0] assert os.path.exists(os.path.join(str(tmp_path), "azure-core_python.json")) pwsh_run.assert_not_called() + + @patch( + "azpysdk.apistub.REPO_ROOT", os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) + ) + @patch("azpysdk.apistub.get_cross_language_mapping_path", return_value=None) + @patch("azpysdk.apistub.create_package_and_install") + @patch("azpysdk.apistub.install_into_venv") + @patch("azpysdk.apistub.set_envvar_defaults") + def test_pypi_version_downloads_wheel_instead_of_building( + self, _env, _install, create_package_and_install, _get_mapping, tmp_path, monkeypatch + ): + """When a PyPI version is passed, the wheel is downloaded from PyPI and local build is skipped.""" + monkeypatch.chdir(os.getcwd()) + stub = apistub() + staging = str(tmp_path / "staging") + os.makedirs(staging, exist_ok=True) + fake_parsed = MagicMock() + fake_parsed.folder = str(tmp_path) + fake_parsed.name = "azure-core" + + captured_cmds = [] + + def fake_apistub_run(exe, cmds, **kwargs): + captured_cmds.append(cmds) + out_idx = cmds.index("--out-path") + out_dir = cmds[out_idx + 1] + open(os.path.join(out_dir, "azure-core_python.json"), "w").close() + + with patch.object(stub, "get_targeted_directories", return_value=[fake_parsed]), patch.object( + stub, "get_executable", return_value=(sys.executable, staging) + ), patch.object(stub, "install_dev_reqs"), patch.object(stub, "pip_freeze"), patch.object( + stub, "ensure_apistub_dependencies" + ), patch.object( + stub, "download_pypi_wheel", return_value="/fake/azure_core-1.0.0-py3-none-any.whl" + ) as download_pypi_wheel, patch.object( + stub, "run_venv_command", side_effect=fake_apistub_run + ): + stub.run(self._make_args(token_file=True, generate_from_pypi="1.0.0")) + + download_pypi_wheel.assert_called_once_with(sys.executable, "azure-core", "1.0.0", staging) + create_package_and_install.assert_not_called() + assert len(captured_cmds) == 1 + pkg_idx = captured_cmds[0].index("--pkg-path") + assert captured_cmds[0][pkg_idx + 1] == os.path.abspath("/fake/azure_core-1.0.0-py3-none-any.whl") + + @patch("azpysdk.apistub.find_whl", return_value="azure_core-1.0.0-py3-none-any.whl") + def test_download_pypi_wheel_runs_pip_download(self, _find_whl, tmp_path): + """download_pypi_wheel should pip download the wheel and return its path.""" + stub = apistub() + staging = str(tmp_path) + + with patch.object(stub, "run_venv_command") as run_venv_command: + result = stub.download_pypi_wheel(sys.executable, "azure-core", "1.0.0", staging) + + run_venv_command.assert_called_once() + cmds = run_venv_command.call_args.args[1] + assert cmds[0:4] == ["-m", "pip", "download", "azure-core==1.0.0"] + assert "--no-deps" in cmds + assert result == os.path.join(staging, "azure_core-1.0.0-py3-none-any.whl") + + @patch("azpysdk.apistub.find_whl", return_value=None) + def test_download_pypi_wheel_raises_when_no_wheel(self, _find_whl, tmp_path): + """download_pypi_wheel should raise FileNotFoundError when no wheel is downloaded.""" + stub = apistub() + with patch.object(stub, "run_venv_command"): + with pytest.raises(FileNotFoundError, match="No wheel found"): + stub.download_pypi_wheel(sys.executable, "azure-core", "1.0.0", str(tmp_path))