polars_python/conversion/
categorical.rs

1use std::sync::Arc;
2
3use polars_dtype::categorical::{CatSize, Categories};
4use pyo3::{pyclass, pymethods};
5
6#[pyclass(frozen)]
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    fn __getnewargs__(&self) -> (String, String, String) {
29        (
30            self.categories.name().to_string(),
31            self.categories.namespace().to_string(),
32            self.categories.physical().as_str().to_owned(),
33        )
34    }
35
36    #[staticmethod]
37    pub fn global_categories() -> Self {
38        Self {
39            categories: Categories::global(),
40        }
41    }
42
43    #[staticmethod]
44    pub fn random(namespace: String, physical: String) -> Self {
45        Self {
46            categories: Categories::random(namespace.into(), physical.parse().unwrap()),
47        }
48    }
49
50    pub fn __eq__(&self, other: &Self) -> bool {
51        Arc::ptr_eq(&self.categories, &other.categories)
52    }
53
54    pub fn __hash__(&self) -> u64 {
55        self.categories.hash()
56    }
57
58    pub fn name(&self) -> &str {
59        self.categories.name()
60    }
61
62    pub fn namespace(&self) -> &str {
63        self.categories.namespace()
64    }
65
66    pub fn physical(&self) -> &str {
67        self.categories.physical().as_str()
68    }
69
70    pub fn get_cat(&self, s: &str) -> Option<CatSize> {
71        self.categories.mapping().get_cat(s)
72    }
73
74    pub fn cat_to_str(&self, cat: CatSize) -> Option<String> {
75        Some(self.categories.mapping().cat_to_str(cat)?.to_owned())
76    }
77
78    pub fn is_global(&self) -> bool {
79        self.categories.is_global()
80    }
81}
82
83impl From<Arc<Categories>> for PyCategories {
84    fn from(categories: Arc<Categories>) -> Self {
85        Self { categories }
86    }
87}