Python - Add ability to retrieve a range of string

This commit is contained in:
Anthony MOI
2019-12-29 01:37:03 -05:00
parent 049029dc42
commit 06d515d41b

View File

@ -1,6 +1,7 @@
extern crate tokenizers as tk;
use pyo3::prelude::*;
use pyo3::types::*;
#[pyclass(dict)]
#[repr(transparent)]
@ -26,6 +27,34 @@ impl Encoding {
self.encoding.get_normalized().get().to_owned()
}
#[args(kwargs = "**")]
fn get_range(
&self,
range: (usize, usize),
kwargs: Option<&PyDict>,
) -> PyResult<Option<String>> {
let mut original = false;
if let Some(kwargs) = kwargs {
if let Some(koriginal) = kwargs.get_item("original") {
original = koriginal.extract()?;
}
}
if original {
Ok(self
.encoding
.get_normalized()
.get_range_original(range.0..range.1)
.map(|s| s.to_owned()))
} else {
Ok(self
.encoding
.get_normalized()
.get_range(range.0..range.1)
.map(|s| s.to_owned()))
}
}
#[getter]
fn get_ids(&self) -> Vec<u32> {
self.encoding.get_ids().to_vec()