Add type stubs

cel_expr_python/BUILD

@@ -1,6 +1,6 @@
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
-load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:py_test.bzl", "py_test")
 
 package(default_visibility = ["//visibility:private"])
 
@@ -45,6 +45,7 @@ pybind_extension(
         "py_message_factory.cc",
         "py_message_factory.h",
     ],
+    data = ["cel.pyi"],
     visibility = ["//visibility:public"],
     deps = [
         ":cel_extension",

cel_expr_python/cel.pyi

@@ -0,0 +1,77 @@
+from collections.abc import Callable, Mapping, Sequence
+import datetime
+from typing import Any, ClassVar
+from google.protobuf import descriptor_pool as proto_descriptor_pool
+from google.protobuf import message as proto_message
+
+class Activation: pass
+
+class CelExtension(CelExtensionBase):
+  def __init__(self, name: str, functions: Sequence[FunctionDecl]) -> None: ...
+
+class CelExtensionBase:
+  def __init__(self, name: str) -> None: ...
+
+class Env:
+  def Activation(self, data: Mapping[str, Any] | None = ..., functions: Sequence[Function] | None = ..., arena: _InternalArena = ...) -> Activation: ...
+  def compile(self, expression: str, disable_check: bool = ...) -> Expression: ...
+  def deserialize(self, serialized: str) -> Expression: ...
+
+class Expression:
+  def eval(self, activation: Activation | None = ..., data: Mapping[str, Any] | None = ..., functions=..., arena: _InternalArena = ...) -> Value: ...
+  def return_type(self) -> Type: ...
+  def serialize(self) -> bytes: ...
+
+class Function:
+  def __init__(self, function_name: str, parameters: Sequence[Type], is_member: bool, impl: Callable[..., Any], return_type: Type = ...) -> None: ...
+
+class FunctionDecl:
+  def __init__(self, name: str, overloads: Sequence[Overload]) -> None: ...
+
+class Overload:
+  def __init__(self, overload_id: str, return_type: Type, parameters: Sequence[Type], is_member: bool = ..., impl: Callable[..., Any] = ...) -> None: ...
+
+class Type:
+  BOOL: ClassVar[Type] = ...
+  BYTES: ClassVar[Type] = ...
+  DOUBLE: ClassVar[Type] = ...
+  DURATION: ClassVar[Type] = ...
+  DYN: ClassVar[Type] = ...
+  ERROR: ClassVar[Type] = ...
+  INT: ClassVar[Type] = ...
+  LIST: ClassVar[Type] = ...
+  MAP: ClassVar[Type] = ...
+  NULL: ClassVar[Type] = ...
+  STRING: ClassVar[Type] = ...
+  TIMESTAMP: ClassVar[Type] = ...
+  TYPE: ClassVar[Type] = ...
+  UINT: ClassVar[Type] = ...
+  UNKNOWN: ClassVar[Type] = ...
+  def __init__(self, name: str) -> None: ...
+  @staticmethod
+  def AbstractType(name: str, params: Sequence[Type] = ...) -> Type: ...
+  @staticmethod
+  def List(element_type: Type) -> Type: ...
+  @staticmethod
+  def Map(key_type: Type, element_type: Type) -> Type: ...
+  @staticmethod
+  def Type(element_type: Type) -> Type: ...
+  def is_assignable_from(self, arg0: Type) -> bool: ...
+  def is_message(self) -> bool: ...
+  def name(self) -> str: ...
+  def __eq__(self, arg0: Type) -> bool: ...
+
+type BaseValue = None | bool | int | float | str | bytes | proto_message.Message | datetime.datetime | datetime.timedelta | Type
+type PlainValue = BaseValue | list[PlainValue] | dict[PlainValue, PlainValue]
+type CelValue = BaseValue | list[Value] | dict[PlainValue, Value]
+
+class Value:
+  def plain_value(self) -> PlainValue: ...
+  def type(self) -> Type: ...
+  def value(self) -> CelValue: ...
+
+class _InternalArena: pass
+
+def Arena() -> _InternalArena: ...
+
+def NewEnv(descriptor_pool: proto_descriptor_pool.DescriptorPool = ..., variables: Mapping[str, Type] | None = ..., extensions: Sequence[CelExtensionBase] | None = ..., container: str | None = ...) -> Env: ...

cel_expr_python/ext/BUILD

@@ -8,6 +8,7 @@ pybind_extension(
         "ext_bindings.cc",
     ],
     data = [
+        "ext_bindings.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],
@@ -24,6 +25,7 @@ pybind_extension(
         "ext_encoders.cc",
     ],
     data = [
+        "ext_encoders.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],
@@ -43,6 +45,7 @@ pybind_extension(
         "ext_math.cc",
     ],
     data = [
+        "ext_math.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],
@@ -64,6 +67,7 @@ pybind_extension(
         "ext_optional.cc",
     ],
     data = [
+        "ext_optional.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],
@@ -86,6 +90,7 @@ pybind_extension(
         "ext_proto.cc",
     ],
     data = [
+        "ext_proto.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],
@@ -102,6 +107,7 @@ pybind_extension(
         "ext_strings.cc",
     ],
     data = [
+        "ext_strings.pyi",
         "//cel_expr_python:cel",
     ],
     visibility = ["//visibility:public"],

cel_expr_python/ext/ext_bindings.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtBindings(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

cel_expr_python/ext/ext_encoders.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtEncoders(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

cel_expr_python/ext/ext_math.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtMath(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

cel_expr_python/ext/ext_optional.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtOptional(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

cel_expr_python/ext/ext_proto.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtProto(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

cel_expr_python/ext/ext_strings.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class ExtStrings(cel.CelExtensionBase):
+  def __init__(self) -> None: ...

conformance/BUILD

@@ -1,4 +1,4 @@
-load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:py_test.bzl", "py_test")
 
 package(default_visibility = ["//visibility:private"])
 

custom_ext/BUILD

@@ -1,13 +1,14 @@
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
-load("@rules_python//python:defs.bzl", "py_test")
 load("@rules_python//python:py_library.bzl", "py_library")
+load("@rules_python//python:py_test.bzl", "py_test")
 
 package(default_visibility = ["//visibility:private"])
 
 pybind_extension(
     name = "sample_cel_ext_cc",
     srcs = ["sample_cel_ext.cc"],
     data = [
+        "sample_cel_ext.pyi",
         "//cel_expr_python:cel",
     ],
     deps = [

custom_ext/sample_cel_ext.pyi

@@ -0,0 +1,4 @@
+from cel_expr_python import cel
+
+class SampleCelExtension(cel.CelExtensionBase):
+  def __init__(self) -> None: ...