Skip to content

Commit

Permalink
use py check_signals() to catch CTRL-C
Browse files Browse the repository at this point in the history
  • Loading branch information
ctb committed Jul 12, 2024
1 parent fb378be commit 49f57eb
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/fastgather.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
/// fastgather: Run gather with a query against a list of files.
use pyo3::Python;
use anyhow::Result;
use sourmash::prelude::Select;
use sourmash::selection::Selection;
Expand All @@ -10,6 +11,7 @@ use crate::utils::{

#[allow(clippy::too_many_arguments)]
pub fn fastgather(
py: Python,
query_filepath: String,
against_filepath: String,
threshold_bp: usize,
Expand Down Expand Up @@ -94,6 +96,7 @@ pub fn fastgather(

// run the gather!
consume_query_by_gather(
py,
query_sig,
scaled as u64,
matchlist,
Expand Down
3 changes: 3 additions & 0 deletions src/fastmultigather.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/// fastmultigather: Run gather for multiple queries against a list of files.
use anyhow::Result;
use rayon::prelude::*;
use pyo3::Python;

use sourmash::selection::Selection;

Expand All @@ -17,6 +18,7 @@ use crate::utils::{
};

pub fn fastmultigather(
py: Python,
query_filepath: String,
against_filepath: String,
threshold_bp: usize,
Expand Down Expand Up @@ -98,6 +100,7 @@ pub fn fastmultigather(

// Now, do the gather!
consume_query_by_gather(
py,
query_sig.clone(),
scaled as u64,
matchlist,
Expand Down
10 changes: 10 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use camino::Utf8PathBuf as PathBuf;

#[pyfunction]
fn do_manysearch(
py: Python,
querylist_path: String,
siglist_path: String,
threshold: f64,
Expand Down Expand Up @@ -73,6 +74,7 @@ fn do_manysearch(
#[pyfunction]
#[allow(clippy::too_many_arguments)]
fn do_fastgather(
py: Python,
query_filename: String,
siglist_path: String,
threshold_bp: usize,
Expand All @@ -86,6 +88,7 @@ fn do_fastgather(
let allow_failed_sigpaths = true;

match fastgather::fastgather(
py,
query_filename,
siglist_path,
threshold_bp,
Expand All @@ -105,6 +108,7 @@ fn do_fastgather(

#[pyfunction]
fn do_fastmultigather(
py: Python,
query_filenames: String,
siglist_path: String,
threshold_bp: usize,
Expand Down Expand Up @@ -138,6 +142,7 @@ fn do_fastmultigather(
bail!("output path specified, but not running fastmultigather against a rocksdb. See issue #239");
}
match fastmultigather::fastmultigather(
py,
query_filenames,
siglist_path,
threshold_bp,
Expand Down Expand Up @@ -173,6 +178,7 @@ fn set_global_thread_pool(num_threads: usize) -> PyResult<usize> {

#[pyfunction]
fn do_index(
py: Python,
siglist: String,
ksize: u8,
scaled: usize,
Expand Down Expand Up @@ -206,6 +212,7 @@ fn do_check(index: String, quick: bool) -> anyhow::Result<u8> {
#[pyfunction]
#[allow(clippy::too_many_arguments)]
fn do_multisearch(
py: Python,
querylist_path: String,
siglist_path: String,
threshold: f64,
Expand Down Expand Up @@ -238,6 +245,7 @@ fn do_multisearch(
#[pyfunction]
#[allow(clippy::too_many_arguments)]
fn do_pairwise(
py: Python,
siglist_path: String,
threshold: f64,
ksize: u8,
Expand Down Expand Up @@ -268,6 +276,7 @@ fn do_pairwise(

#[pyfunction]
fn do_manysketch(
py: Python,
filelist: String,
param_str: String,
output: String,
Expand All @@ -285,6 +294,7 @@ fn do_manysketch(

#[pyfunction]
fn do_cluster(
py: Python,
pairwise_csv: String,
output_clusters: String,
similarity_column: String,
Expand Down
3 changes: 3 additions & 0 deletions src/utils.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
/// Utility functions for sourmash_plugin_branchwater.
use rayon::prelude::*;
use pyo3::Python;
use sourmash::encodings::HashFunctions;
use sourmash::selection::Select;

Expand Down Expand Up @@ -926,6 +927,7 @@ pub fn branchwater_calculate_gather_stats(
/// removing matches in 'matchlist' from 'query'.
pub fn consume_query_by_gather(
py: Python,
query: SigStore,
scaled: u64,
matchlist: BinaryHeap<PrefetchResult>,
Expand Down Expand Up @@ -994,6 +996,7 @@ pub fn consume_query_by_gather(
);

while !matching_sketches.is_empty() {
py.check_signals()?;
let best_element = matching_sketches.peek().unwrap();

query_mh = query_mh.downsample_scaled(best_element.minhash.scaled())?;
Expand Down

0 comments on commit 49f57eb

Please sign in to comment.