feat(spider-py): Define TDL types; Add support for converting Python types into TDL types. (#187)

Co-authored-by: Lin Zhihao <[email protected]>

Author: sitaowang1998

python/spider-py/pyproject.toml

@@ -60,13 +60,15 @@ ignore = [
   "FIX002",  # Allow todo statements
   "PERF401",  # Allow for loops when creating lists
   "PERF403",  # Allow for loops when creating dicts
+  "PLR1716",  # Allow chained boolean comparisons
   "S311",  # Allow usage of `random` package
   "SIM102",  # Allow collapsible if statements for readability
   "TD002",  # Author unnecessary for todo statement
   "TD003",  # Issue link unnecessary for todo statement
   "UP015",  # Explicit open modes are helpful
 ]
 isort.order-by-type = false
+typing-extensions = false  # Disable type extensions for 3.10 compatibility
 
 [tool.ruff.lint.per-file-ignores]
 "tests/**" = [
@@ -75,3 +77,6 @@ isort.order-by-type = false
   "S603",  # Allow use of `subprocess.Popen` (security warning)
   "T201",  # Allow use of `print` (testing)
 ]
+
+[tool.ruff.lint.pydocstyle]
+ignore-decorators = ["typing.override"]

python/spider-py/src/spider_py/__init__.py

@@ -1 +1,12 @@
 """Spider package root."""
+
+from spider_py.type import Double, Float, Int8, Int16, Int32, Int64
+
+__all__ = [
+    "Double",
+    "Float",
+    "Int8",
+    "Int16",
+    "Int32",
+    "Int64",
+]

python/spider-py/src/spider_py/type/__init__.py

@@ -0,0 +1,6 @@
+"""Spider type package."""
+
+from spider_py.type.tdl_convert import to_tdl_type_str
+from spider_py.type.type import Double, Float, Int8, Int16, Int32, Int64
+
+__all__ = ["Double", "Float", "Int8", "Int16", "Int32", "Int64", "to_tdl_type_str"]

python/spider-py/src/spider_py/type/tdl_convert.py

@@ -0,0 +1,99 @@
+"""Converts native types to TDL types."""
+
+from collections.abc import Collection
+from types import GenericAlias
+from typing import get_args, get_origin
+
+from spider_py.type.tdl_type import (
+    BoolType,
+    ClassType,
+    DoubleType,
+    FloatType,
+    Int8Type,
+    Int16Type,
+    Int32Type,
+    Int64Type,
+    ListType,
+    MapType,
+    TdlType,
+)
+from spider_py.type.type import Double, Float, Int8, Int16, Int32, Int64
+from spider_py.type.utils import get_class_name
+
+TypeDict = {
+    Int8: Int8Type(),
+    Int16: Int16Type(),
+    Int32: Int32Type(),
+    Int64: Int64Type(),
+    Float: FloatType(),
+    Double: DoubleType(),
+    bool: BoolType(),
+}
+
+
+def _to_primitive_tdl_type(native_type: type | GenericAlias) -> TdlType | None:
+    """
+    Converts a native type to primitive TDL type.
+    :param native_type:
+    :return:
+        - The converted TDL primitive if `native_type` is supported.
+        - None if `native_type` is not a primitive Python type.
+    :raises TypeError: If `native_type` is a primitive Python type not supported by TDL.
+    """
+    if isinstance(native_type, type) and native_type in TypeDict:
+        return TypeDict[native_type]
+
+    if native_type in (int, float, str, complex, bytes):
+        msg = f"{native_type} is not a TDL type. Please use the corresponding TDL primitive type."
+        raise TypeError(msg)
+
+    return None
+
+
+def to_tdl_type(native_type: type | GenericAlias) -> TdlType:
+    """
+    Converts a Python type to TDL type.
+    :param native_type:
+    :return: The converted TDL type.
+    :raises TypeError: If `native_type` is not a valid TDL type.
+    """
+    primitive_tdl_type = _to_primitive_tdl_type(native_type)
+    if primitive_tdl_type is not None:
+        return primitive_tdl_type
+
+    if isinstance(native_type, GenericAlias):
+        origin = get_origin(native_type)
+        if origin is list:
+            args = get_args(native_type)
+            if len(args) == 0:
+                msg = "List does not have an element type."
+                raise TypeError(msg)
+            arg = args[0]
+            return ListType(to_tdl_type(arg))
+
+        if origin is dict:
+            args = get_args(native_type)
+            msg = "Dict does not have a key/value type."
+            if len(args) != 2:  # noqa: PLR2004
+                raise TypeError(msg)
+            key = args[0]
+            value = args[1]
+            return MapType(to_tdl_type(key), to_tdl_type(value))
+
+        msg = f"{native_type} is not a valid TDL type."
+        raise TypeError(msg)
+
+    if issubclass(native_type, Collection):
+        msg = f"{native_type} is not a valid TDL type."
+        raise TypeError(msg)
+
+    return ClassType(get_class_name(native_type))
+
+
+def to_tdl_type_str(native_type: type | GenericAlias) -> str:
+    """
+    :param native_type: A Python native type.
+    :return: A string representation of the TDL type.
+    :raises TypeError: If `native_type` is not a valid TDL type.
+    """
+    return to_tdl_type(native_type).type_str()

python/spider-py/src/spider_py/type/tdl_type.py

@@ -0,0 +1,144 @@
+"""Spider TDL types."""
+
+from abc import ABC, abstractmethod
+
+from typing_extensions import override
+
+
+class TdlType(ABC):
+    """Abstract base class for all TDL types."""
+
+    @abstractmethod
+    def type_str(self) -> str:
+        """:return: String representation of the TDL type."""
+
+
+class DoubleType(TdlType):
+    """TDL double type."""
+
+    @override
+    def type_str(self) -> str:
+        return "double"
+
+
+class FloatType(TdlType):
+    """TDL float type."""
+
+    @override
+    def type_str(self) -> str:
+        return "float"
+
+
+class Int8Type(TdlType):
+    """TDL int8 type."""
+
+    @override
+    def type_str(self) -> str:
+        return "int8"
+
+
+class Int16Type(TdlType):
+    """TDL int16 type."""
+
+    @override
+    def type_str(self) -> str:
+        return "int16"
+
+
+class Int32Type(TdlType):
+    """TDL int32 type."""
+
+    @override
+    def type_str(self) -> str:
+        return "int32"
+
+
+class Int64Type(TdlType):
+    """TDL int64 type."""
+
+    @override
+    def type_str(self) -> str:
+        return "int64"
+
+
+class BoolType(TdlType):
+    """TDL bool type."""
+
+    @override
+    def type_str(self) -> str:
+        return "bool"
+
+
+class ClassType(TdlType):
+    """TDL Custom class type."""
+
+    def __init__(self, name: str) -> None:
+        """
+        Creates a TDL custom class type.
+        :param name: The name of the class.
+        """
+        self._name = name
+
+    @override
+    def type_str(self) -> str:
+        return self._name
+
+
+class ListType(TdlType):
+    """TDL List type."""
+
+    def __init__(self, element_type: TdlType) -> None:
+        """
+        Creates a TDL list type.
+        :param element_type:
+        """
+        self.element_type = element_type
+
+    @override
+    def type_str(self) -> str:
+        return f"List<{self.element_type.type_str()}>"
+
+
+def is_integral(tdl_type: TdlType) -> bool:
+    """
+    :param tdl_type:
+    :return: Whether `tdl_type` is a TDL integral type.
+    """
+    return isinstance(tdl_type, (Int8Type, Int16Type, Int32Type, Int64Type))
+
+
+def is_string(tdl_type: TdlType) -> bool:
+    """
+    :param tdl_type:
+    :return: Whether `tdl_type` is a TDL string type, i.e. `List<int8>`.
+    """
+    return isinstance(tdl_type, ListType) and isinstance(tdl_type.element_type, Int8Type)
+
+
+def is_map_key(tdl_type: TdlType) -> bool:
+    """
+    :param tdl_type:
+    :return: Whether `tdl_type` is a supported key type of a map.
+    """
+    return is_integral(tdl_type) or is_string(tdl_type)
+
+
+class MapType(TdlType):
+    """TDL Map type."""
+
+    def __init__(self, key_type: TdlType, value_type: TdlType) -> None:
+        """
+        Creates a TDL map type.
+        :param key_type:
+        :param value_type:
+        :raises TypeError: If key is not a supported type.
+        """
+        if not is_map_key(key_type):
+            msg = f"{key_type} is not a supported type for map key."
+            raise TypeError(msg)
+        self.key_type = key_type
+        self.value_type = value_type
+
+    @override
+    def type_str(self) -> str:
+        return f"Map<{self.key_type.type_str()},{self.value_type.type_str()}>"

python/spider-py/src/spider_py/type/type.py

@@ -0,0 +1,74 @@
+"""Custom type module for Spider."""
+
+from __future__ import annotations
+
+from typing import cast
+
+
+class BoundedInt(int):
+    """Bounded integer type."""
+
+    def __new__(cls, value: int, bits: int = 32) -> BoundedInt:
+        """Creates a bounded integer."""
+        if bits not in (8, 16, 32, 64):
+            msg = f"Unsupported bits size: {bits}. Supported sizes are 8, 16, 32, or 64."
+            raise ValueError(msg)
+
+        lower_bound = -(1 << (bits - 1))
+        upper_bound = (1 << (bits - 1)) - 1
+
+        if not (lower_bound <= value and value <= upper_bound):
+            msg = (
+                f"Bounded integer value ({value}) must be between {lower_bound} and {upper_bound}."
+            )
+            raise ValueError(msg)
+
+        return super().__new__(cls, value)
+
+
+class Int8(BoundedInt):
+    """8 bits integer type."""
+
+    def __new__(cls, value: int) -> Int8:
+        """Creates an int8 integer."""
+        return cast("Int8", super().__new__(cls, value, bits=8))
+
+
+class Int16(BoundedInt):
+    """16 bits integer type."""
+
+    def __new__(cls, value: int) -> Int16:
+        """Creates an int16 integer."""
+        return cast("Int16", super().__new__(cls, value, bits=16))
+
+
+class Int32(BoundedInt):
+    """32 bits integer type."""
+
+    def __new__(cls, value: int) -> Int32:
+        """Creates an int32 integer."""
+        return cast("Int32", super().__new__(cls, value, bits=32))
+
+
+class Int64(BoundedInt):
+    """64 bits integer type."""
+
+    def __new__(cls, value: int) -> Int64:
+        """Creates an int64 integer."""
+        return cast("Int64", super().__new__(cls, value, bits=64))
+
+
+class Float(float):
+    """Float type."""
+
+    def __new__(cls, value: float) -> Float:
+        """Creates a float number."""
+        return super().__new__(cls, value)
+
+
+class Double(float):
+    """Double type."""
+
+    def __new__(cls, value: float) -> Double:
+        """Creates a double number."""
+        return super().__new__(cls, value)

python/spider-py/src/spider_py/type/utils.py

@@ -0,0 +1,9 @@
+"""Utility for TDL types."""
+
+
+def get_class_name(cls: type) -> str:
+    """
+    :param cls:
+    :return: Full name of `cls`.
+    """
+    return f"{cls.__module__}.{cls.__qualname__}"