feat(python): add typing information

python-sdk/pyproject.toml

@@ -19,6 +19,13 @@ dynamic = ["version"]
 [project.urls]
 "Bug Tracker" = "https://github.com/Eppo-exp/rust-sdk/issues"
 
+[project.optional-dependencies]
+test = [
+  "pytest",
+  "cachetools",
+  "types-cachetools"
+]
+
 [tool.maturin]
 features = ["pyo3/extension-module"]
 python-source = "python"

python-sdk/python/eppo_client/__init__.py

@@ -1,3 +1,5 @@
+from typing import Any, Dict, Set, Union, Type
+
 # Rust currently does not define submodules as packages, so Rust
 # submodules are not importable from Python.[1] There is a hacky way
 # to make submodules re-exportable (by tweaking sys.modules) but it
@@ -10,12 +12,18 @@
 #
 # [1]: https://github.com/PyO3/pyo3/issues/759
 # [2]: https://www.maturin.rs/project_layout#pure-rust-project
-from ._eppo_client import *
+import eppo_client._eppo_client as _eppo_client
+from eppo_client._eppo_client import *
+from eppo_client._eppo_client import __version__
 
 # re-exports
 from eppo_client.assignment_logger import AssignmentCacheLogger
 from eppo_client.bandit import BanditResult
 
+Attribute = Union[str, int, float, bool, None]
+Attributes = Dict[str, Attribute]
+
 __doc__ = _eppo_client.__doc__
+__all__ = ["AssignmentCacheLogger", "BanditResult", "Attribute", "Attributes"]
 if hasattr(_eppo_client, "__all__"):
-    __all__ = _eppo_client.__all__
+    __all__.extend(_eppo_client.__all__)

python-sdk/python/eppo_client/_eppo_client.pyi

@@ -0,0 +1,172 @@
+from typing import Dict, Any, Set, Union
+
+# version: str
+__version__: str
+
+def init(config: ClientConfig) -> EppoClient: ...
+def get_instance() -> EppoClient: ...
+
+class Configuration:
+    def __init__(self, flags_configuration: bytes) -> None: ...
+    def get_flags_configuration(self) -> bytes: ...
+    def get_flag_keys(self) -> Set[str]: ...
+    def get_bandit_keys(self) -> Set[str]: ...
+
+class ClientConfig:
+    api_key: str
+    base_url: str
+    assignment_logger: AssignmentLogger
+    is_graceful_mode: bool
+    poll_interval_seconds: int | None
+    poll_jitter_seconds: int
+    initial_configuration: Configuration | None
+
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        base_url: str = ...,
+        assignment_logger: AssignmentLogger,
+        is_graceful_mode: bool = True,
+        poll_interval_seconds: int | None = ...,
+        poll_jitter_seconds: int = ...,
+        initial_configuration: Configuration | None = None
+    ): ...
+
+class AssignmentLogger:
+    def log_assignment(self, event: Dict) -> None: ...
+    def log_bandit_action(self, event: Dict) -> None: ...
+
+class EppoClient:
+    def get_string_assignment(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: str,
+    ) -> str: ...
+    def get_integer_assignment(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: int,
+    ) -> int: ...
+    def get_numeric_assignment(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: float,
+    ) -> float: ...
+    def get_boolean_assignment(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: bool,
+    ) -> bool: ...
+    def get_json_assignment(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: Any,
+    ) -> Any: ...
+    def get_string_assignment_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: str,
+    ) -> EvaluationResult: ...
+    def get_integer_assignment_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: int,
+    ) -> EvaluationResult: ...
+    def get_numeric_assignment_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: float,
+    ) -> EvaluationResult: ...
+    def get_boolean_assignment_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: bool,
+    ) -> EvaluationResult: ...
+    def get_json_assignment_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_attributes: Dict[str, Union[str, int, float, bool, None]],
+        default: Any,
+    ) -> EvaluationResult: ...
+    def get_bandit_action(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_context: (
+            ContextAttributes | Dict[str, Union[str, int, float, bool, None]]
+        ),
+        actions: (
+            Dict[str, ContextAttributes]
+            | Dict[str, Dict[str, Union[str, int, float, bool, None]]]
+        ),
+        default: str,
+    ) -> EvaluationResult: ...
+    def get_bandit_action_details(
+        self,
+        flag_key: str,
+        subject_key: str,
+        subject_context: (
+            ContextAttributes | Dict[str, Union[str, int, float, bool, None]]
+        ),
+        actions: (
+            Dict[str, ContextAttributes]
+            | Dict[str, Dict[str, Union[str, int, float, bool, None]]]
+        ),
+        default: str,
+    ) -> EvaluationResult: ...
+    def get_configuration(self) -> Configuration: ...
+    def set_configuration(self, configuration: Configuration): ...
+    def get_flag_keys(self) -> Set[str]: ...
+    def get_bandit_keys(self) -> Set[str]: ...
+    def set_is_graceful_mode(self, is_graceful_mode: bool): ...
+    def is_initialized(self) -> bool: ...
+    def wait_for_initialization(self) -> None: ...
+
+class ContextAttributes:
+    def __new__(
+        cls,
+        numeric_attributes: Dict[str, float],
+        categorical_attributes: Dict[str, str],
+    ): ...
+    @staticmethod
+    def empty() -> ContextAttributes: ...
+    @staticmethod
+    def from_dict(
+        attributes: Dict[str, Union[str, int, float, bool, None]]
+    ) -> ContextAttributes: ...
+    @property
+    def numeric_attributes(self) -> Dict[str, float]: ...
+    @property
+    def categorical_attributes(self) -> Dict[str, str]: ...
+
+class EvaluationResult:
+    variation: Any
+    action: str | None
+    evaluation_details: Any | None
+    def __new__(
+        cls,
+        variation: Any,
+        action: str | None = None,
+        evaluation_details: Any | None = None,
+    ): ...
+    def to_string(self) -> str: ...

python-sdk/python/eppo_client/config.py

@@ -1,9 +1 @@
 from eppo_client import ClientConfig as Config, AssignmentLogger
-
-import warnings
-
-warnings.warn(
-    "the eppo_client.config module is deprecated, use eppo_client instead",
-    DeprecationWarning,
-    stacklevel=2,
-)

python-sdk/src/client.rs

@@ -428,6 +428,12 @@ impl EppoClient {
         EvaluationResult::from_bandit_result(py, result, Some(details))
     }
 
+    fn get_configuration(&self) -> Option<Configuration> {
+        self.configuration_store
+            .get_configuration()
+            .map(Configuration::new)
+    }
+
     fn set_configuration(&self, configuration: &Configuration) {
         self.configuration_store
             .set_configuration(Arc::clone(&configuration.configuration));
@@ -457,14 +463,10 @@ impl EppoClient {
         }
     }
 
-    fn get_configuration(&self) -> Option<Configuration> {
-        self.configuration_store
-            .get_configuration()
-            .map(Configuration::new)
-    }
-
-    // Returns a set of all flag keys that have been initialized.
-    // This can be useful to debug the initialization process.
+    /// Returns a set of all flag keys that have been initialized.
+    /// This can be useful to debug the initialization process.
+    ///
+    /// Deprecated. Use EppoClient.get_configuration() instead.
     fn get_flag_keys<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<PySet>> {
         let config = self.configuration_store.get_configuration();
         match config {
@@ -473,8 +475,10 @@ impl EppoClient {
         }
     }
 
-    // Returns a set of all bandit keys that have been initialized.
-    // This can be useful to debug the initialization process.
+    /// Returns a set of all bandit keys that have been initialized.
+    /// This can be useful to debug the initialization process.
+    ///
+    /// Deprecated. Use EppoClient.get_configuration() instead.
     fn get_bandit_keys<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<PySet>> {
         let config = self.configuration_store.get_configuration();
         match config {