view rust/hg-cpython/src/update.rs @ 52190:e6a44bc91bc2 stable

rust-update: make `update_from_null` respect `worker.numcpu` config option

This was overlooked in the original series.

This is important for tests (because we run many at once), and for the
occasional end user that wants to keep their CPU usage in check.

A future series should clean up this `worker` parameter tunelling business by
rewriting the config in Rust, but doing so on stable would be a very bad
idea.

// debug.rs
//
// Copyright 2024 Mercurial developers
//
// This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version.

//! Module for updating a repository.
use cpython::{PyDict, PyModule, PyObject, PyResult, Python};
use hg::{
    progress::{HgProgressBar, Progress},
    update::update_from_null,
    BaseRevision,
};

use crate::{
    exceptions::FallbackError,
    utils::{hgerror_to_pyerr, repo_from_path},
};

pub fn update_from_null_fast_path(
    py: Python,
    repo_path: PyObject,
    to: BaseRevision,
    num_cpus: Option<usize>,
) -> PyResult<usize> {
    log::trace!("Using update from null fastpath");
    let repo = repo_from_path(py, repo_path)?;
    let progress: &dyn Progress = &HgProgressBar::new("updating");
    hgerror_to_pyerr(
        py,
        update_from_null(&repo, to.into(), progress, num_cpus),
    )
}

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let dotted_name = &format!("{}.update", package);
    let m = PyModule::new(py, dotted_name)?;

    m.add(py, "__package__", package)?;
    m.add(py, "__doc__", "Rust module for updating a repository")?;
    m.add(py, "FallbackError", py.get_type::<FallbackError>())?;
    m.add(
        py,
        "update_from_null",
        py_fn!(
            py,
            update_from_null_fast_path(
                repo_path: PyObject,
                to: BaseRevision,
                num_cpus: Option<usize>
            )
        ),
    )?;

    let sys = PyModule::import(py, "sys")?;
    let sys_modules: PyDict = sys.get(py, "modules")?.extract(py)?;
    sys_modules.set_item(py, dotted_name, &m)?;

    Ok(m)
}