lindera_py/
dictionary.rs

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    // Ensure output directory exists
38    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    // Ensure output directory exists
62    fs::create_dir_all(output_dir).map_err(|err| {
63        PyValueError::new_err(format!("Failed to create output directory: {err}"))
64    })?;
65
66    // Determine output file name based on input file
67    // If the input file has no name, we cannot determine the output file name.
68    // In that case, we return an error.
69    // e.g., /path/to/input/file.txt -> /path/to/output/file.bin
70    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}