vortex/pyvortex/src/python_repr.rs at 33a9c9147f0e10b85fada35b3b0834d0b1eee4b0 · vortex-data/vortex · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
use std::convert::AsRef;
use std::fmt::{Display, Formatter};

use itertools::Itertools;
use vortex_dtype::{DType, ExtID, ExtMetadata, Nullability, PType};

pub trait PythonRepr {
    fn python_repr(&self) -> impl Display;
}

struct DTypePythonRepr<'a>(&'a DType);

impl PythonRepr for DType {
    fn python_repr(&self) -> impl Display {
        DTypePythonRepr(self)
    }
}

impl Display for DTypePythonRepr<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let DTypePythonRepr(dtype) = self;
        match dtype {
            DType::Null => write!(f, "null()"),
            DType::Bool(n) => write!(f, "bool({})", n.python_repr()),
            DType::Primitive(ptype, n) => match ptype {
                PType::U8 | PType::U16 | PType::U32 | PType::U64 => {
                    write!(f, "uint({}, {})", ptype.bit_width(), n.python_repr())
                }
                PType::I8 | PType::I16 | PType::I32 | PType::I64 => {
                    write!(f, "int({}, {})", ptype.bit_width(), n.python_repr())
                }
                PType::F16 | PType::F32 | PType::F64 => {
                    write!(f, "float({}, {})", ptype.bit_width(), n.python_repr())
                }
            },
            DType::Utf8(n) => write!(f, "utf8({})", n.python_repr()),
            DType::Binary(n) => write!(f, "binary({})", n.python_repr()),
            DType::Struct(st, n) => write!(
                f,
                "struct({{{}}}, {})",
                st.names()
                    .iter()
                    .zip(st.dtypes().iter())
                    .map(|(n, dt)| format!("\"{}\": {}", n, dt.python_repr()))
                    .join(", "),
                n.python_repr()
            ),
            DType::List(edt, n) => write!(f, "list({}, {})", edt.python_repr(), n.python_repr()),
            DType::Extension(ext) => {
                write!(
                    f,
                    "ext(\"{}\", {}, ",
                    ext.id().python_repr(),
                    ext.storage_dtype().python_repr()
                )?;
                match ext.metadata() {
                    None => write!(f, "None")?,
                    Some(metadata) => write!(f, "{}", metadata.python_repr())?,
                };
                write!(f, ")")
            }
        }
    }
}

struct NullabilityPythonRepr<'a>(&'a Nullability);

impl PythonRepr for Nullability {
    fn python_repr(&self) -> impl Display {
        NullabilityPythonRepr(self)
    }
}

impl Display for NullabilityPythonRepr<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let NullabilityPythonRepr(x) = self;
        match x {
            Nullability::NonNullable => write!(f, "False"),
            Nullability::Nullable => write!(f, "True"),
        }
    }
}

struct ExtMetadataPythonRepr<'a>(&'a ExtMetadata);

impl PythonRepr for ExtMetadata {
    fn python_repr(&self) -> impl Display {
        ExtMetadataPythonRepr(self)
    }
}

impl Display for ExtMetadataPythonRepr<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let ExtMetadataPythonRepr(metadata) = self;
        write!(f, "\"{}\"", metadata.as_ref().escape_ascii())
    }
}

struct ExtIDPythonRepr<'a>(&'a ExtID);

impl PythonRepr for ExtID {
    fn python_repr(&self) -> impl Display {
        ExtIDPythonRepr(self)
    }
}

impl Display for ExtIDPythonRepr<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let ExtIDPythonRepr(ext_id) = self;
        write!(f, "\"{}\"", ext_id.as_ref().escape_default())
    }
}