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
56fn 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#[pyfunction]
72#[pyo3(signature = (reference, query, **kwargs))]
73fn align(
74 reference: Bound<'_, PyAny>, 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#[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}