1use std::fs;
2use std::path::Path;
3use std::str::FromStr;
4
5use pyo3::{exceptions::PyValueError, prelude::*};
6
7use lindera::dictionary::{
8 Dictionary, DictionaryKind, UserDictionary, load_dictionary_from_kind,
9 load_dictionary_from_path, load_user_dictionary_from_bin, load_user_dictionary_from_csv,
10 resolve_builder, resolve_metadata,
11};
12
13#[pyclass(name = "Dictionary")]
14#[derive(Clone)]
15pub struct PyDictionary {
16 pub inner: Dictionary,
17}
18
19#[pyclass(name = "UserDictionary")]
20#[derive(Clone)]
21pub struct PyUserDictionary {
22 pub inner: UserDictionary,
23}
24
25#[pyfunction]
26#[pyo3(signature = (kind, input_dir, output_dir))]
27pub fn build_dictionary(kind: &str, input_dir: &str, output_dir: &str) -> PyResult<()> {
28 let dict_kind =
29 DictionaryKind::from_str(kind).map_err(|_err| PyValueError::new_err("Invalid kind"))?;
30
31 let metadata = resolve_metadata(dict_kind.clone())
32 .map_err(|err| PyValueError::new_err(format!("Failed to resolve metadata: {err}")))?;
33
34 let builder = resolve_builder(dict_kind)
35 .map_err(|err| PyValueError::new_err(format!("Failed to resolve builder: {err}")))?;
36
37 fs::create_dir_all(output_dir).map_err(|err| {
39 PyValueError::new_err(format!("Failed to create output directory: {err}"))
40 })?;
41
42 builder
43 .build_dictionary(&metadata, Path::new(input_dir), Path::new(output_dir))
44 .map_err(|err| PyValueError::new_err(format!("Failed to build dictionary: {err}")))?;
45
46 Ok(())
47}
48
49#[pyfunction]
50#[pyo3(signature = (kind, input_file, output_dir))]
51pub fn build_user_dictionary(kind: &str, input_file: &str, output_dir: &str) -> PyResult<()> {
52 let dict_kind =
53 DictionaryKind::from_str(kind).map_err(|_err| PyValueError::new_err("Invalid kind"))?;
54
55 let metadata = resolve_metadata(dict_kind.clone())
56 .map_err(|err| PyValueError::new_err(format!("Failed to resolve metadata: {err}")))?;
57
58 let builder = resolve_builder(dict_kind)
59 .map_err(|err| PyValueError::new_err(format!("Failed to resolve builder: {err}")))?;
60
61 fs::create_dir_all(output_dir).map_err(|err| {
63 PyValueError::new_err(format!("Failed to create output directory: {err}"))
64 })?;
65
66 let output_file = if let Some(filename) = Path::new(input_file).file_name() {
71 let mut output_file = Path::new(output_dir).join(filename);
72 output_file.set_extension("bin");
73 output_file
74 } else {
75 return Err(PyValueError::new_err("Failed to determine output filename"));
76 };
77
78 builder
79 .build_user_dictionary(&metadata, Path::new(input_file), output_file.as_path())
80 .map_err(|err| PyValueError::new_err(format!("Failed to build user dictionary: {err}")))?;
81
82 Ok(())
83}
84
85#[pyfunction]
86#[pyo3(signature = (kind=None, path=None))]
87pub fn load_dictionary(kind: Option<&str>, path: Option<&str>) -> PyResult<PyDictionary> {
88 match (kind, path) {
89 (Some(kind_str), None) => {
90 let k = DictionaryKind::from_str(kind_str)
91 .map_err(|_err| PyValueError::new_err("Invalid kind"))?;
92 let dictionary = load_dictionary_from_kind(k).map_err(|err| {
93 PyValueError::new_err(format!("Failed to load dictionary: {err}"))
94 })?;
95
96 Ok(PyDictionary { inner: dictionary })
97 }
98 (None, Some(path_str)) => {
99 let p = Path::new(path_str);
100 let dictionary = load_dictionary_from_path(p).map_err(|err| {
101 PyValueError::new_err(format!("Failed to load dictionary: {err}"))
102 })?;
103
104 Ok(PyDictionary { inner: dictionary })
105 }
106 _ => Err(PyValueError::new_err("Invalid arguments")),
107 }
108}
109
110#[pyfunction]
111#[pyo3(signature = (path, kind=None))]
112pub fn load_user_dictionary(path: &str, kind: Option<&str>) -> PyResult<PyUserDictionary> {
113 let p = Path::new(path);
114 let ext = p
115 .extension()
116 .and_then(|ext| ext.to_str())
117 .ok_or_else(|| PyValueError::new_err("Invalid file path"))?;
118 match ext {
119 "csv" => match kind {
120 Some(kind) => {
121 let k = DictionaryKind::from_str(kind)
122 .map_err(|_err| PyValueError::new_err("Invalid kind"))?;
123 let user_dictionary = load_user_dictionary_from_csv(k, p).map_err(|err| {
124 PyValueError::new_err(format!("Failed to load user dictionary: {err}"))
125 })?;
126
127 Ok(PyUserDictionary {
128 inner: user_dictionary,
129 })
130 }
131 None => Err(PyValueError::new_err(
132 "Dictionary type must be specified if CSV file specified",
133 )),
134 },
135 "bin" => match kind {
136 Some(_kind) => Err(PyValueError::new_err(
137 "Dictionary type must be None if Binary file specified",
138 )),
139 None => {
140 let user_dictionary = load_user_dictionary_from_bin(p).map_err(|err| {
141 PyValueError::new_err(format!("Failed to load user dictionary: {err}"))
142 })?;
143
144 Ok(PyUserDictionary {
145 inner: user_dictionary,
146 })
147 }
148 },
149 _ => Err(PyValueError::new_err(format!(
150 "Unsupported file: path:{path}, kind:{kind:?}"
151 ))),
152 }
153}