Skip to main content

novel_segment_napi/
lib.rs

1use napi::bindgen_prelude::*;
2use napi_derive::napi;
3use novel_segment::{DoSegmentOptions, Segment, SegmentOptions};
4use std::sync::Mutex;
5
6#[napi(object)]
7pub struct JsSegmentOptions {
8    pub auto_cjk: Option<bool>,
9    pub all_mod: Option<bool>,
10    pub node_novel_mode: Option<bool>,
11    pub convert_synonym: Option<bool>,
12}
13
14#[napi(object)]
15pub struct JsDoSegmentOptions {
16    pub simple: Option<bool>,
17    pub strip_punctuation: Option<bool>,
18    pub convert_synonym: Option<bool>,
19    pub strip_stopword: Option<bool>,
20    pub strip_space: Option<bool>,
21}
22
23#[napi(object)]
24pub struct JsWord {
25    pub w: String,
26    pub p: Option<u32>,
27    pub f: Option<f64>,
28}
29
30#[napi]
31pub struct NativeSegment {
32    inner: Mutex<Segment>,
33}
34
35fn map_err(e: novel_segment::Error) -> Error {
36    Error::from_reason(e.to_string())
37}
38
39#[napi]
40impl NativeSegment {
41    #[napi(factory)]
42    pub fn create(options: Option<JsSegmentOptions>) -> Result<Self> {
43        let o = options.unwrap_or(JsSegmentOptions {
44            auto_cjk: Some(true),
45            all_mod: Some(true),
46            node_novel_mode: Some(false),
47            convert_synonym: Some(true),
48        });
49        let mut seg = Segment::new(SegmentOptions {
50            auto_cjk: o.auto_cjk.unwrap_or(true),
51            all_mod: o.all_mod.unwrap_or(true),
52            node_novel_mode: o.node_novel_mode.unwrap_or(false),
53            options_do_segment: DoSegmentOptions {
54                convert_synonym: o.convert_synonym.or(Some(true)),
55                ..Default::default()
56            },
57            ..Default::default()
58        });
59        seg.use_default().map_err(map_err)?;
60        Ok(Self {
61            inner: Mutex::new(seg),
62        })
63    }
64
65    #[napi(factory)]
66    pub fn with_node_novel_default() -> Result<Self> {
67        let seg = Segment::with_node_novel_default().map_err(map_err)?;
68        Ok(Self {
69            inner: Mutex::new(seg),
70        })
71    }
72
73    #[napi]
74    pub fn do_segment(&self, text: String, options: Option<JsDoSegmentOptions>) -> Result<Vec<JsWord>> {
75        let seg = self.inner.lock().map_err(|e| Error::from_reason(e.to_string()))?;
76        let opts = options
77            .map(|o| DoSegmentOptions {
78                simple: o.simple,
79                strip_punctuation: o.strip_punctuation,
80                convert_synonym: o.convert_synonym,
81                strip_stopword: o.strip_stopword,
82                strip_space: o.strip_space,
83                disable_modules: Vec::new(),
84            })
85            .unwrap_or_default();
86        Ok(seg
87            .do_segment(&text, opts)
88            .into_iter()
89            .map(|w| JsWord {
90                w: w.w,
91                p: w.p,
92                f: w.f,
93            })
94            .collect())
95    }
96
97    #[napi]
98    pub fn do_segment_simple(&self, text: String, options: Option<JsDoSegmentOptions>) -> Result<Vec<String>> {
99        Ok(self
100            .do_segment(text, options)?
101            .into_iter()
102            .map(|w| w.w)
103            .collect())
104    }
105
106    #[napi]
107    pub fn stringify(&self, text: String, options: Option<JsDoSegmentOptions>) -> Result<String> {
108        Ok(self.do_segment_simple(text, options)?.concat())
109    }
110
111    #[napi]
112    pub fn add_word(&self, spec: String, p: Option<u32>, f: Option<f64>) -> Result<()> {
113        let mut seg = self.inner.lock().map_err(|e| Error::from_reason(e.to_string()))?;
114        seg.add_word(&spec, p, f).map_err(map_err)?;
115        Ok(())
116    }
117
118    #[napi]
119    pub fn add_synonym(&self, canonical: String, variants: Vec<String>) -> Result<()> {
120        let mut seg = self.inner.lock().map_err(|e| Error::from_reason(e.to_string()))?;
121        let refs: Vec<&str> = variants.iter().map(|s| s.as_str()).collect();
122        seg.add_synonym(&canonical, &refs);
123        Ok(())
124    }
125
126    #[napi]
127    pub fn add_blacklist(&self, word: String) -> Result<()> {
128        let mut seg = self.inner.lock().map_err(|e| Error::from_reason(e.to_string()))?;
129        seg.add_blacklist(&word);
130        Ok(())
131    }
132}