tsalign/
lib.rs

1use std::io;
2
3use lib_tsalign::{
4    a_star_aligner::{
5        alignment_result::AlignmentResult,
6        configurable_a_star_align::{Config, a_star_align},
7        template_switch_distance::AlignmentType,
8    },
9    costs::U64Cost,
10};
11use lib_tsshow::plain_text::show_template_switches;
12use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
13use pythonize::{depythonize, pythonize};
14
15#[pyclass]
16struct TSPairwiseAlignment {
17    result: AlignmentResult<AlignmentType, U64Cost>,
18}
19
20#[pymethods]
21impl TSPairwiseAlignment {
22    fn viz_template_switches(&self) -> PyResult<()> {
23        show_template_switches(io::stdout(), &self.result, &None)
24            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
25        Ok(())
26    }
27
28    fn stats<'a>(&'a self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
29        Ok(pythonize(py, self.result.statistics())?)
30    }
31
32    fn cigar(&self) -> Option<String> {
33        match &self.result {
34            AlignmentResult::WithTarget { alignment, .. } => Some(alignment.cigar()),
35            AlignmentResult::WithoutTarget { .. } => None,
36        }
37    }
38
39    fn alignments<'a>(&'a self, py: Python<'a>) -> PyResult<Option<Bound<'a, PyAny>>> {
40        match &self.result {
41            AlignmentResult::WithTarget { alignment, .. } => {
42                let mut container = Vec::new();
43                alignment.iter_compact().for_each(|e| container.push(e));
44                Ok(Some(pythonize(py, &container)?))
45            }
46            AlignmentResult::WithoutTarget { .. } => Ok(None),
47        }
48    }
49}
50
51fn py_to_str(o: Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
52    let str = o.str()?.to_str()?.as_bytes().to_vec();
53    Ok(str)
54}
55
56/// Creates a config object by amending the default by values present in the python dictionary
57fn create_config(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<Config> {
58    let Some(kwargs) = kwargs else {
59        return Ok(Config::default());
60    };
61
62    let config = depythonize::<Config>(kwargs)?;
63
64    Ok(config)
65}
66
67/// Align two sequences, accounting for template switches
68///
69/// The function takes a reference and a query string, and performs a global alignment on both. The output alignment may contain (short-range) template switches.
70/// Optionally, settings can be specified. See the (configuration struct)[lib_tsalign::a_star_aligner::configurable_a_star_align::Config] for the available keys and values.
71#[pyfunction]
72#[pyo3(signature = (reference, query, **kwargs))]
73fn align(
74    reference: Bound<'_, PyAny>, // Accepting PyAny instead of PyString to allow using e.g. `Bio.Seq` types and alike. String representation will be used.
75    query: Bound<'_, PyAny>,
76    kwargs: Option<&Bound<'_, PyDict>>,
77) -> PyResult<Option<TSPairwiseAlignment>> {
78    let reference = py_to_str(reference)?;
79    let query = py_to_str(query)?;
80    let config = create_config(kwargs)?;
81
82    let r = a_star_align(&reference, &query, &config);
83
84    match r {
85        result @ AlignmentResult::WithTarget { .. } => {
86            let ts_alignment = TSPairwiseAlignment { result };
87            Ok(Some(ts_alignment))
88        }
89        AlignmentResult::WithoutTarget { .. } => Ok(None),
90    }
91}
92
93/// Bindings for the `lib_tsalign` library.
94#[pymodule]
95fn tsalign(m: &Bound<'_, PyModule>) -> PyResult<()> {
96    pyo3_log::init();
97    m.add_class::<TSPairwiseAlignment>()?;
98    m.add_function(wrap_pyfunction!(align, m)?)?;
99    Ok(())
100}