polars_python/conversion/
categorical.rs

1use std::sync::Arc;
2
3use polars_dtype::categorical::{CatSize, Categories};
4use pyo3::{pyclass, pymethods};
5
6#[pyclass]
7#[repr(transparent)]
8#[derive(Clone)]
9pub struct PyCategories {
10    categories: Arc<Categories>,
11}
12
13impl PyCategories {
14    pub fn categories(&self) -> &Arc<Categories> {
15        &self.categories
16    }
17}
18
19#[pymethods]
20impl PyCategories {
21    #[new]
22    pub fn __init__(name: String, namespace: String, physical: String) -> Self {
23        Self {
24            categories: Categories::new(name.into(), namespace.into(), physical.parse().unwrap()),
25        }
26    }
27
28    #[staticmethod]
29    pub fn global_categories() -> Self {
30        Self {
31            categories: Categories::global(),
32        }
33    }
34
35    #[staticmethod]
36    pub fn random(namespace: String, physical: String) -> Self {
37        Self {
38            categories: Categories::random(namespace.into(), physical.parse().unwrap()),
39        }
40    }
41
42    pub fn __eq__(&self, other: &Self) -> bool {
43        Arc::ptr_eq(&self.categories, &other.categories)
44    }
45
46    pub fn __hash__(&self) -> u64 {
47        self.categories.hash()
48    }
49
50    pub fn name(&self) -> &str {
51        self.categories.name()
52    }
53
54    pub fn namespace(&self) -> &str {
55        self.categories.namespace()
56    }
57
58    pub fn physical(&self) -> &str {
59        self.categories.physical().as_str()
60    }
61
62    pub fn get_cat(&self, s: &str) -> Option<CatSize> {
63        self.categories.mapping().get_cat(s)
64    }
65
66    pub fn cat_to_str(&self, cat: CatSize) -> Option<String> {
67        Some(self.categories.mapping().cat_to_str(cat)?.to_owned())
68    }
69
70    pub fn is_global(&self) -> bool {
71        self.categories.is_global()
72    }
73}
74
75impl From<Arc<Categories>> for PyCategories {
76    fn from(categories: Arc<Categories>) -> Self {
77        Self { categories }
78    }
79}