value_wrapper.py

# Copyright 2025 vesoft-inc
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
 
import logging
from datetime import date, datetime, time
from decimal import Decimal
from typing import Any, List, Optional, Type, TypeVar, Union, overload
 
from nebulagraph_python.decoder.data_types import ColumnType
from nebulagraph_python.py_data_types import (
    ColumnToPy,
    CompositeDataObject,
    Edge,
    NDuration,
    Node,
    NRecord,
    NVector,
    Path,
    TargetType,
)
 
logger = logging.getLogger(__name__)
 
 
T = TypeVar("T", bound=TargetType)
 
 
class ValueWrapper:
    __slots__ = ("value", "data_type")
 
    def __init__(self, value: Any, data_type: ColumnType):
        self.value = value
        self.data_type = data_type
 
    def is_null(self) -> bool:
        """If the value is null
 
        Returns
        -------
            bool: true if value is null
 
        """
        return self.value is None
 
    def is_bool(self) -> bool:
        """Check if the Value is Boolean type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_BOOL
 
        """
        return self.data_type == ColumnType.BOOL
 
    def is_long(self) -> bool:
        """Check if the Value is Long type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_UINT64 or COLUMN_TYPE_INT64
 
        """
        return self.data_type in [ColumnType.UINT64, ColumnType.INT64]
 
    def is_int(self) -> bool:
        """Check if the Value is Int type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_UINT8 or COLUMN_TYPE_INT8 or COLUMN_TYPE_UINT16
                or COLUMN_TYPE_INT16 or COLUMN_TYPE_UINT32 or COLUMN_TYPE_INT32
 
        """
        return self.data_type in [
            ColumnType.UINT8,
            ColumnType.INT8,
            ColumnType.UINT16,
            ColumnType.INT16,
            ColumnType.UINT32,
            ColumnType.INT32,
        ]
 
    def is_float(self) -> bool:
        """Check if the Value is Float type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_FLOAT32
 
        """
        return self.data_type == ColumnType.FLOAT32
 
    def is_double(self) -> bool:
        """Check if the Value is Double type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_FLOAT64
 
        """
        return self.data_type == ColumnType.FLOAT64
 
    def is_string(self) -> bool:
        """Check if the Value is String type.
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_STRING
 
        """
        return self.data_type == ColumnType.STRING
 
    def is_list(self) -> bool:
        """Check if the Value is List type.
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_LIST
 
        """
        return self.data_type == ColumnType.LIST
 
    def is_node(self) -> bool:
        """Check if the Value is Node type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_NODE
 
        """
        return self.data_type == ColumnType.NODE
 
    def is_edge(self) -> bool:
        """Check if the Value is Edge type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_EDGE
 
        """
        return self.data_type == ColumnType.EDGE
 
    def is_local_time(self) -> bool:
        """Check if the Value is Local Time type.
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_LOCALTIME
 
        """
        return self.data_type == ColumnType.LOCALTIME
 
    def is_zoned_time(self) -> bool:
        """Check if the Value is Zoned Time type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_ZONEDTIME
 
        """
        return self.data_type == ColumnType.ZONEDTIME
 
    def is_local_datetime(self) -> bool:
        """Check if the Value is Local Datetime type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_LOCALDATETIME
 
        """
        return self.data_type == ColumnType.LOCALDATETIME
 
    def is_zoned_datetime(self) -> bool:
        """Check if the Value is Zoned Datetime type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_ZONEDDATETIME
 
        """
        return self.data_type == ColumnType.ZONEDDATETIME
 
    def is_date(self) -> bool:
        """Check if the Value is Date type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_DATE
 
        """
        return self.data_type == ColumnType.DATE
 
    def is_record(self) -> bool:
        """Check if the Value is Record type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_RECORD
 
        """
        return self.data_type == ColumnType.RECORD
 
    def is_duration(self) -> bool:
        """Check if the Value is Duration type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_DURATION
 
        """
        return self.data_type == ColumnType.DURATION
 
    def is_path(self) -> bool:
        """Check if the Value is Path type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_PATH
 
        """
        return self.data_type == ColumnType.PATH
 
    def is_decimal(self) -> bool:
        """Check if the Value is Decimal type
 
        Returns
        -------
            bool: true if Value's type is COLUMN_TYPE_DECIMAL
 
        """
        return self.data_type == ColumnType.DECIMAL
 
    def as_bool(self):
        return self.castcastcast(bool)
 
    def as_int(self):
        return self.castcastcast(int)
 
    def as_long(self):
        return self.castcastcast(int)
 
    def as_string(self):
        return self.castcastcast(str)
 
    def as_float(self):
        return self.castcastcast(float)
 
    def as_double(self):
        return self.castcastcast(float)
 
    def as_list(self):
        return self.castcastcast(list)
 
    def as_node(self):
        return self.castcastcast(Node)
 
    def as_edge(self):
        return self.castcastcast(Edge)
 
    def as_local_time(self):
        return self.castcastcast(time)
 
    def as_zoned_time(self):
        return self.castcastcast(time)
 
    def as_date(self):
        return self.castcastcast(date)
 
    def as_local_datetime(self):
        return self.castcastcast(datetime)
 
    def as_zoned_datetime(self):
        return self.castcastcast(datetime)
 
    def as_duration(self):
        return self.castcastcast(NDuration)
 
    def as_record(self):
        return self.castcastcast(NRecord)
 
    def as_path(self):
        return self.castcastcast(Path)
 
    def as_decimal(self):
        return self.castcastcast(Decimal)
 
    def as_embedding_vector(self):
        return self.castcastcast(NVector)
 
    @overload
    def cast(self) -> TargetType:
        """Get self.value"""
        ...
 
    @overload
    def cast(self, target_type: Type[T]) -> T:
        """Get self.value, with target_type check"""
        ...
 
    def cast(
        self,
        target_type: Optional[Type[T]] = None,
    ) -> Union[T, TargetType]:
        if self.value is None and target_type is None:
            return None
        target_type = ColumnToPy[self.data_type] if target_type is None else target_type  # type: ignore
 
        if target_type is not None:
            if isinstance(self.value, target_type):
                return self.value
            else:
                raise TypeError(
                    f"Cannot cast {type(self.value)} to {target_type}. ColumnType and value:",
                    self.data_type,
                    self.value,
                )
 
    def cast_primitive(self) -> Any:
        """Convert the wrapped value to primitive Python types recursively.
 
        For basic types, uses cast()
        For composite types (Path, Node, Edge), calls cast_primitive()
        For containers (Map, List, Record), recursively calls cast_primitive() on elements
        """
        outer = self.castcastcast()
 
        # Handle composite types
        if isinstance(outer, CompositeDataObject):
            return outer.cast_primitive()
 
        # Handle list recursively
        if isinstance(outer, list):
            ans = []
            for i, v in enumerate(outer):
                if isinstance(v, ValueWrapper):
                    ans.append(v.cast_primitive())
                else:
                    raise TypeError(
                        f"Cannot cast list, where list[{i}] have type {type(v)}, not ValueWrapper. list[{i}]: ",
                        v,
                    )
            return ans
 
        return outer
 
    def __eq__(self, other):
        if not isinstance(other, ValueWrapper):
            return False
        return self.value == other.value and self.data_type == other.data_type
 
    def __hash__(self) -> int:
        return hash((self.value, self.data_type))
 
 
class Row:
    __slots__ = ("values",)
 
    def __init__(self, values: Optional[List[ValueWrapper]] = None):
        self.values = values if values is not None else []
 
    def add_value(self, value: ValueWrapper):
        """Append one value into row.
 
        Args:
        ----
            value (ValueWrapper): one value of the row
 
        """
        self.values.append(value)
 
    def get_values(self) -> list[ValueWrapper]:
        """Get the values of this row.
 
        Returns
        -------
            List[ValueWrapper]: list of values in this row
 
        """
        return self.values