use std::cell::RefCell;
use std::path::Path;
use magnus::prelude::*;
use magnus::{Error, RArray, RHash, Ruby, function, method};
use lindera_binding_core::{CoreTokenizer, CoreTokenizerBuilder};
use crate::dictionary::{RbDictionary, RbUserDictionary};
use crate::error::to_magnus_error;
use crate::token::RbToken;
use crate::util::rb_hash_to_json;
#[magnus::wrap(class = "Lindera::TokenizerBuilder", free_immediately, size)]
pub struct RbTokenizerBuilder {
inner: RefCell<CoreTokenizerBuilder>,
}
impl RbTokenizerBuilder {
fn new() -> Result<Self, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let inner =
CoreTokenizerBuilder::new().map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
Ok(Self {
inner: RefCell::new(inner),
})
}
fn from_file(file_path: String) -> Result<Self, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let inner = CoreTokenizerBuilder::from_file(Path::new(&file_path))
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
Ok(Self {
inner: RefCell::new(inner),
})
}
fn set_mode(&self, mode: String) -> Result<(), Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let mut builder = self.inner.borrow_mut();
builder
.set_mode(&mode)
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
Ok(())
}
fn set_dictionary(&self, path: String) {
self.inner.borrow_mut().set_dictionary(&path);
}
fn set_user_dictionary(&self, uri: String) {
self.inner.borrow_mut().set_user_dictionary(&uri);
}
fn set_keep_whitespace(&self, keep_whitespace: bool) {
self.inner.borrow_mut().set_keep_whitespace(keep_whitespace);
}
fn append_character_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let filter_args = if let Some(hash) = args {
rb_hash_to_json(&ruby, hash)?
} else {
serde_json::Value::Object(serde_json::Map::new())
};
self.inner
.borrow_mut()
.append_character_filter(&kind, &filter_args);
Ok(())
}
fn append_token_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let filter_args = if let Some(hash) = args {
rb_hash_to_json(&ruby, hash)?
} else {
serde_json::Value::Object(serde_json::Map::new())
};
self.inner
.borrow_mut()
.append_token_filter(&kind, &filter_args);
Ok(())
}
fn build(&self) -> Result<RbTokenizer, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let inner = self
.inner
.borrow()
.build()
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
Ok(RbTokenizer { inner })
}
}
#[magnus::wrap(class = "Lindera::Tokenizer", free_immediately, size)]
pub struct RbTokenizer {
inner: CoreTokenizer,
}
fn tokenizer_new(
dictionary: &RbDictionary,
mode: Option<String>,
user_dictionary: Option<&RbUserDictionary>,
) -> Result<RbTokenizer, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let mode_str = mode.as_deref().unwrap_or("normal");
let dict = dictionary.inner.clone();
let user_dict = user_dictionary.map(|d| d.inner.clone());
let inner = CoreTokenizer::from_segmenter(mode_str, dict, user_dict)
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
Ok(RbTokenizer { inner })
}
impl RbTokenizer {
fn tokenize(&self, text: String) -> Result<RArray, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let views = self
.inner
.tokenize(&text)
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
let arr = ruby.ary_new_capa(views.len());
for view in views {
arr.push(ruby.into_value(RbToken::from_view(view)))?;
}
Ok(arr)
}
fn tokenize_nbest(
&self,
text: String,
n: usize,
unique: Option<bool>,
cost_threshold: Option<i64>,
) -> Result<RArray, Error> {
let ruby = Ruby::get().expect("Ruby runtime not initialized");
let results = self
.inner
.tokenize_nbest(&text, n, unique.unwrap_or(false), cost_threshold)
.map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
let rb_results = ruby.ary_new_capa(results.len());
for (views, cost) in results {
let token_arr = ruby.ary_new_capa(views.len());
for view in views {
token_arr.push(ruby.into_value(RbToken::from_view(view)))?;
}
let pair = ruby.ary_new_capa(2);
pair.push(token_arr)?;
pair.push(cost)?;
rb_results.push(pair)?;
}
Ok(rb_results)
}
}
pub fn define(ruby: &Ruby, module: &magnus::RModule) -> Result<(), Error> {
let builder_class = module.define_class("TokenizerBuilder", ruby.class_object())?;
builder_class.define_singleton_method("new", function!(RbTokenizerBuilder::new, 0))?;
builder_class
.define_singleton_method("from_file", function!(RbTokenizerBuilder::from_file, 1))?;
builder_class.define_method("set_mode", method!(RbTokenizerBuilder::set_mode, 1))?;
builder_class.define_method(
"set_dictionary",
method!(RbTokenizerBuilder::set_dictionary, 1),
)?;
builder_class.define_method(
"set_user_dictionary",
method!(RbTokenizerBuilder::set_user_dictionary, 1),
)?;
builder_class.define_method(
"set_keep_whitespace",
method!(RbTokenizerBuilder::set_keep_whitespace, 1),
)?;
builder_class.define_method(
"append_character_filter",
method!(RbTokenizerBuilder::append_character_filter, 2),
)?;
builder_class.define_method(
"append_token_filter",
method!(RbTokenizerBuilder::append_token_filter, 2),
)?;
builder_class.define_method("build", method!(RbTokenizerBuilder::build, 0))?;
let tokenizer_class = module.define_class("Tokenizer", ruby.class_object())?;
tokenizer_class.define_singleton_method("new", function!(tokenizer_new, 3))?;
tokenizer_class.define_method("tokenize", method!(RbTokenizer::tokenize, 1))?;
tokenizer_class.define_method("tokenize_nbest", method!(RbTokenizer::tokenize_nbest, 4))?;
Ok(())
}