pyridis-message 0.3.5

iridis, with python
Documentation 
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

from .pyridis_message import *

import pyarrow as pa
import numpy as np

from typing import Union
from enum import Enum

from typing import TypeVar, Type, Optional

T = TypeVar("T", bound="ArrowMessage")

class ArrowMessage:
    @classmethod
    def from_arrow(cls: Type[T], data: pa.Array) -> Optional[T]:
        if data.type == pa.null():
            return None

        if issubclass(cls, Enum):
            value = data.to_numpy(zero_copy_only=False)[0]

            for encoding in cls:
                if encoding.value == value:
                    return encoding

            return None
        else:
            names = [field.name for field in data.type]
            indices = [index for index in data.type.type_codes]

            map = dict(zip(names, indices))

            kwargs = {}
            for name in names:
                field_type = cls.__dataclass_fields__[name].type  # pyright: ignore

                if hasattr(field_type, "__origin__") and field_type.__origin__ is Union:
                    field_type = field_type.__args__[0]

                if issubclass(field_type, np.ndarray):
                    kwargs[name] = data.field(map[name]).to_numpy(zero_copy_only=False)

                elif issubclass(field_type, ArrowMessage):
                    kwargs[name] = field_type.from_arrow(data.field(map[name]))

                else:
                    kwargs[name] = data.field(map[name]).to_numpy(zero_copy_only=False)[
                        0
                    ]

            return cls(**kwargs)

    def to_arrow(self) -> pa.Array:
        if self is None:
            return pa.array([], type=pa.null())

        if isinstance(self, Enum):
            return pa.array([self.value], type=pa.utf8())

        fields = list(self.__dataclass_fields__.keys())  # pyright: ignore

        children = []
        for field in fields:
            value = getattr(self, field)
            if value is None:
                children.append(pa.array([], type=pa.null()))
            elif isinstance(value, np.ndarray):
                value.flags.writeable = False
                children.append(pa.array(value))
            elif isinstance(value, ArrowMessage):
                children.append(value.to_arrow())
            else:
                children.append(pa.array([value]))

        return pa.UnionArray.from_dense(
            types=pa.array([], type=pa.int8()),
            value_offsets=pa.array([], type=pa.int32()),
            children=children,
            field_names=fields,
        )