feat(spider-py): Add support for parsing TDL types into native Python types.

python/spider-py/pyproject.toml

@@ -9,6 +9,7 @@ authors = [
   { name = "Zhihao Lin", email = "lin.zhihao@yscope.com" },
 ]
 dependencies = [
+  "lark>=1.2.2",
   "mariadb>=1.1.13",
   "msgpack-types>=0.5.0",
 ]

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

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

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

@@ -0,0 +1,102 @@
+"""Parse TDL type string."""
+
+from lark import Lark, LarkError, Token, Transformer, v_args
+
+from spider_py.type.tdl_type import (
+    BoolType,
+    ClassType,
+    DoubleType,
+    FloatType,
+    Int8Type,
+    Int16Type,
+    Int32Type,
+    Int64Type,
+    ListType,
+    MapType,
+    TdlType,
+)
+
+grammar = r"""
+type: map_type | list_type | base_type
+
+map_type: "Map" "<" type "," type ">"
+list_type: "List" "<" type ">"
+base_type: ID
+
+ID: /[a-zA-Z_][a-zA-Z0-9_]*(\.[a-zA-Z_][a-zA-Z0-9_]*)*/
+
+%import common.WS
+%ignore WS
+"""
+
+PrimitiveTypeDict = {
+    "bool": BoolType(),
+    "double": DoubleType(),
+    "float": FloatType(),
+    "int8": Int8Type(),
+    "int16": Int16Type(),
+    "int32": Int32Type(),
+    "int64": Int64Type(),
+}
+
+
+class TypeTransformer(Transformer[Token, TdlType]):
+    """Transforms Lark parse tree into TDL type."""
+
+    @v_args(inline=True)
+    def type(self, value_type: TdlType) -> TdlType:
+        """
+        Unwraps the type node to return the TdlType.
+        :param value_type: The type node to unwrap.
+        :return: The unwrapped TdlType.
+        """
+        return value_type
+
+    @v_args(inline=True)
+    def map_type(self, key_type: TdlType, value_type: TdlType) -> TdlType:
+        """
+        Transforms map node into Map type.
+        :param key_type: The key type of the map.
+        :param value_type: The value type of the map.
+        :return: The MapType with the given `key_type` and `value_type`.
+        """
+        return MapType(key_type, value_type)
+
+    @v_args(inline=True)
+    def list_type(self, element_type: TdlType) -> TdlType:
+        """
+        Transforms list node into List type.
+        :param element_type: The element type of the list.
+        :return: The ListType with the given `element_type`.
+        """
+        return ListType(element_type)
+
+    @v_args(inline=True)
+    def base_type(self, token_id: str) -> TdlType:
+        """
+        Transforms primitive node into primitive type.
+        :param token_id: The token id of the primitive type.
+        :return: The corresponding TdlType of the `token_id`.
+        """
+        if token_id in PrimitiveTypeDict:
+            return PrimitiveTypeDict[token_id]
+        return ClassType(token_id)
+
+
+_parser = Lark(grammar, start="type", parser="lalr")
+_transformer = TypeTransformer(visit_tokens=False)
+
+
+def parse_tdl_type(type_string: str) -> TdlType:
+    """
+    Parses TDL type string into TDL type.
+    :param type_string: TDL type string.
+    :return: Parsed TDL type.
+    :raises TypeError: If string is not a valid TDL type.
+    """
+    try:
+        tree = _parser.parse(type_string)
+        return _transformer.transform(tree)
+    except LarkError as e:
+        msg = f"Cannot parse TDL type `{type_string}`"
+        raise TypeError(msg) from e

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

@@ -1,9 +1,13 @@
 """Spider TDL types."""
 
 from abc import ABC, abstractmethod
+from types import GenericAlias
 
 from typing_extensions import override
 
+from spider_py.type.type import Double, Float, Int8, Int16, Int32, Int64
+from spider_py.type.utils import get_class_by_name
+
 
 class TdlType(ABC):
     """Abstract base class for all TDL types."""
@@ -12,6 +16,10 @@ class TdlType(ABC):
     def type_str(self) -> str:
         """:return: String representation of the TDL type."""
 
+    @abstractmethod
+    def native_type(self) -> type | GenericAlias:
+        """:return: Native Python type of the TDL type."""
+
 
 class DoubleType(TdlType):
     """TDL double type."""
@@ -20,6 +28,10 @@ class DoubleType(TdlType):
     def type_str(self) -> str:
         return "double"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Double
+
 
 class FloatType(TdlType):
     """TDL float type."""
@@ -28,6 +40,10 @@ class FloatType(TdlType):
     def type_str(self) -> str:
         return "float"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Float
+
 
 class Int8Type(TdlType):
     """TDL int8 type."""
@@ -36,6 +52,10 @@ class Int8Type(TdlType):
     def type_str(self) -> str:
         return "int8"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Int8
+
 
 class Int16Type(TdlType):
     """TDL int16 type."""
@@ -44,6 +64,10 @@ class Int16Type(TdlType):
     def type_str(self) -> str:
         return "int16"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Int16
+
 
 class Int32Type(TdlType):
     """TDL int32 type."""
@@ -52,6 +76,10 @@ class Int32Type(TdlType):
     def type_str(self) -> str:
         return "int32"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Int32
+
 
 class Int64Type(TdlType):
     """TDL int64 type."""
@@ -60,6 +88,10 @@ class Int64Type(TdlType):
     def type_str(self) -> str:
         return "int64"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return Int64
+
 
 class BoolType(TdlType):
     """TDL bool type."""
@@ -68,6 +100,10 @@ class BoolType(TdlType):
     def type_str(self) -> str:
         return "bool"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return bool
+
 
 class ClassType(TdlType):
     """TDL Custom class type."""
@@ -83,6 +119,14 @@ def __init__(self, name: str) -> None:
     def type_str(self) -> str:
         return self._name
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        """
+        :return: Native Python type of the class.
+        :raise: TypeError if `class_name` is not a valid class.
+        """
+        return get_class_by_name(self._name)
+
 
 class ListType(TdlType):
     """TDL List type."""
@@ -98,6 +142,10 @@ def __init__(self, element_type: TdlType) -> None:
     def type_str(self) -> str:
         return f"List<{self.element_type.type_str()}>"
 
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return list[self.element_type.native_type()]  # type: ignore[misc]
+
 
 def is_integral(tdl_type: TdlType) -> bool:
     """
@@ -142,3 +190,7 @@ def __init__(self, key_type: TdlType, value_type: TdlType) -> None:
     @override
     def type_str(self) -> str:
         return f"Map<{self.key_type.type_str()},{self.value_type.type_str()}>"
+
+    @override
+    def native_type(self) -> type | GenericAlias:
+        return dict[self.key_type.native_type(), self.value_type.native_type()]  # type: ignore[misc]

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

@@ -1,9 +1,30 @@
 """Utility for TDL types."""
 
+from importlib import import_module
+from typing import cast
+
 
 def get_class_name(cls: type) -> str:
     """
     :param cls:
     :return: Full name of `cls`.
     """
     return f"{cls.__module__}.{cls.__qualname__}"
+
+
+def get_class_by_name(name: str) -> type:
+    """
+    Gets class by name.
+    :param name:
+    :return:
+    :raise: TypeError if `class_name` is not a valid class.
+    """
+    parts = name.split(".")
+    module_name = ".".join(parts[:-1])
+    class_name = parts[-1]
+    try:
+        module = import_module(module_name)
+        return cast("type", getattr(module, class_name))
+    except (ValueError, ModuleNotFoundError, AttributeError) as exc:
+        msg = f"{name} is not a valid TDL type."
+        raise TypeError(msg) from exc

python/spider-py/tests/type/test_to_native.py

@@ -0,0 +1,54 @@
+"""Test converting TDL type to native type."""
+
+from types import GenericAlias
+
+import pytest
+
+import spider_py
+from spider_py.type.tdl_parse import parse_tdl_type
+from spider_py.type.utils import get_class_name
+
+
+def string_to_native(s: str) -> type | GenericAlias:
+    """
+    Converts TDL string to native type.
+    :param s: TDL string to convert.
+    :return: The converted native type.
+    :raises: TypeError forwards from `parse_tdl_type` and `native_type`.
+    """
+    return parse_tdl_type(s).native_type()
+
+
+class TestToNativeType:
+    """Test converting TDL type to native type."""
+
+    def test_to_primitive_native_type(self) -> None:
+        """Test converting primitive TDL type to native type."""
+        assert string_to_native("bool") is bool
+        assert string_to_native("double") is spider_py.Double
+        assert string_to_native("float") is spider_py.Float
+        assert string_to_native("int8") is spider_py.Int8
+        assert string_to_native("int16") is spider_py.Int16
+        assert string_to_native("int32") is spider_py.Int32
+        assert string_to_native("int64") is spider_py.Int64
+
+    def test_to_class_type(self) -> None:
+        """Test converting class TDL type to native type."""
+        assert string_to_native(get_class_name(TestToNativeType)) == TestToNativeType
+        with pytest.raises(TypeError):
+            string_to_native("NonExistType")
+        with pytest.raises(TypeError):
+            string_to_native("NonExistModule.NonExistType")
+
+    def test_to_list_type(self) -> None:
+        """Test converting list TDL type to native type."""
+        assert string_to_native("List<int8>") == list[spider_py.Int8]
+        assert string_to_native("List<List<int8>>") == list[list[spider_py.Int8]]
+
+    def test_to_map_type(self) -> None:
+        """Test converting map TDL type to native type."""
+        assert string_to_native("Map<int8,int8>") == dict[spider_py.Int8, spider_py.Int8]
+        assert (
+            string_to_native("Map<List<int8>,Map<int8,double>>")
+            == dict[list[spider_py.Int8], dict[spider_py.Int8, spider_py.Double]]
+        )