1use std::cell::RefCell;
9use std::path::Path;
10
11use magnus::prelude::*;
12use magnus::{Error, RArray, RHash, Ruby, function, method};
13
14use lindera_binding_core::{CoreTokenizer, CoreTokenizerBuilder};
15
16use crate::dictionary::{RbDictionary, RbUserDictionary};
17use crate::error::to_magnus_error;
18use crate::token::RbToken;
19use crate::util::rb_hash_to_json;
20
21#[magnus::wrap(class = "Lindera::TokenizerBuilder", free_immediately, size)]
26pub struct RbTokenizerBuilder {
27 inner: RefCell<CoreTokenizerBuilder>,
29}
30
31impl RbTokenizerBuilder {
32 fn new() -> Result<Self, Error> {
38 let ruby = Ruby::get().expect("Ruby runtime not initialized");
39 let inner =
40 CoreTokenizerBuilder::new().map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
41 Ok(Self {
42 inner: RefCell::new(inner),
43 })
44 }
45
46 fn from_file(file_path: String) -> Result<Self, Error> {
56 let ruby = Ruby::get().expect("Ruby runtime not initialized");
57 let inner = CoreTokenizerBuilder::from_file(Path::new(&file_path))
58 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
59 Ok(Self {
60 inner: RefCell::new(inner),
61 })
62 }
63
64 fn set_mode(&self, mode: String) -> Result<(), Error> {
70 let ruby = Ruby::get().expect("Ruby runtime not initialized");
71 let mut builder = self.inner.borrow_mut();
72 builder
73 .set_mode(&mode)
74 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
75 Ok(())
76 }
77
78 fn set_dictionary(&self, path: String) {
84 self.inner.borrow_mut().set_dictionary(&path);
85 }
86
87 fn set_user_dictionary(&self, uri: String) {
93 self.inner.borrow_mut().set_user_dictionary(&uri);
94 }
95
96 fn set_keep_whitespace(&self, keep_whitespace: bool) {
102 self.inner.borrow_mut().set_keep_whitespace(keep_whitespace);
103 }
104
105 fn append_character_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
112 let ruby = Ruby::get().expect("Ruby runtime not initialized");
113 let filter_args = if let Some(hash) = args {
114 rb_hash_to_json(&ruby, hash)?
115 } else {
116 serde_json::Value::Object(serde_json::Map::new())
117 };
118 self.inner
119 .borrow_mut()
120 .append_character_filter(&kind, &filter_args);
121 Ok(())
122 }
123
124 fn append_token_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
131 let ruby = Ruby::get().expect("Ruby runtime not initialized");
132 let filter_args = if let Some(hash) = args {
133 rb_hash_to_json(&ruby, hash)?
134 } else {
135 serde_json::Value::Object(serde_json::Map::new())
136 };
137 self.inner
138 .borrow_mut()
139 .append_token_filter(&kind, &filter_args);
140 Ok(())
141 }
142
143 fn build(&self) -> Result<RbTokenizer, Error> {
149 let ruby = Ruby::get().expect("Ruby runtime not initialized");
150 let inner = self
151 .inner
152 .borrow()
153 .build()
154 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
155 Ok(RbTokenizer { inner })
156 }
157}
158
159#[magnus::wrap(class = "Lindera::Tokenizer", free_immediately, size)]
163pub struct RbTokenizer {
164 inner: CoreTokenizer,
166}
167
168fn tokenizer_new(
180 dictionary: &RbDictionary,
181 mode: Option<String>,
182 user_dictionary: Option<&RbUserDictionary>,
183) -> Result<RbTokenizer, Error> {
184 let ruby = Ruby::get().expect("Ruby runtime not initialized");
185 let mode_str = mode.as_deref().unwrap_or("normal");
186
187 let dict = dictionary.inner.clone();
188 let user_dict = user_dictionary.map(|d| d.inner.clone());
189
190 let inner = CoreTokenizer::from_segmenter(mode_str, dict, user_dict)
191 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
192
193 Ok(RbTokenizer { inner })
194}
195
196impl RbTokenizer {
197 fn tokenize(&self, text: String) -> Result<RArray, Error> {
207 let ruby = Ruby::get().expect("Ruby runtime not initialized");
208 let views = self
209 .inner
210 .tokenize(&text)
211 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
212
213 let arr = ruby.ary_new_capa(views.len());
214 for view in views {
215 arr.push(ruby.into_value(RbToken::from_view(view)))?;
216 }
217 Ok(arr)
218 }
219
220 fn tokenize_nbest(
233 &self,
234 text: String,
235 n: usize,
236 unique: Option<bool>,
237 cost_threshold: Option<i64>,
238 ) -> Result<RArray, Error> {
239 let ruby = Ruby::get().expect("Ruby runtime not initialized");
240 let results = self
241 .inner
242 .tokenize_nbest(&text, n, unique.unwrap_or(false), cost_threshold)
243 .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
244
245 let rb_results = ruby.ary_new_capa(results.len());
246 for (views, cost) in results {
247 let token_arr = ruby.ary_new_capa(views.len());
248 for view in views {
249 token_arr.push(ruby.into_value(RbToken::from_view(view)))?;
250 }
251 let pair = ruby.ary_new_capa(2);
252 pair.push(token_arr)?;
253 pair.push(cost)?;
254 rb_results.push(pair)?;
255 }
256
257 Ok(rb_results)
258 }
259}
260
261pub fn define(ruby: &Ruby, module: &magnus::RModule) -> Result<(), Error> {
272 let builder_class = module.define_class("TokenizerBuilder", ruby.class_object())?;
273 builder_class.define_singleton_method("new", function!(RbTokenizerBuilder::new, 0))?;
274 builder_class
275 .define_singleton_method("from_file", function!(RbTokenizerBuilder::from_file, 1))?;
276 builder_class.define_method("set_mode", method!(RbTokenizerBuilder::set_mode, 1))?;
277 builder_class.define_method(
278 "set_dictionary",
279 method!(RbTokenizerBuilder::set_dictionary, 1),
280 )?;
281 builder_class.define_method(
282 "set_user_dictionary",
283 method!(RbTokenizerBuilder::set_user_dictionary, 1),
284 )?;
285 builder_class.define_method(
286 "set_keep_whitespace",
287 method!(RbTokenizerBuilder::set_keep_whitespace, 1),
288 )?;
289 builder_class.define_method(
290 "append_character_filter",
291 method!(RbTokenizerBuilder::append_character_filter, 2),
292 )?;
293 builder_class.define_method(
294 "append_token_filter",
295 method!(RbTokenizerBuilder::append_token_filter, 2),
296 )?;
297 builder_class.define_method("build", method!(RbTokenizerBuilder::build, 0))?;
298
299 let tokenizer_class = module.define_class("Tokenizer", ruby.class_object())?;
300 tokenizer_class.define_singleton_method("new", function!(tokenizer_new, 3))?;
301 tokenizer_class.define_method("tokenize", method!(RbTokenizer::tokenize, 1))?;
302 tokenizer_class.define_method("tokenize_nbest", method!(RbTokenizer::tokenize_nbest, 4))?;
303
304 Ok(())
305}