From 7807c3450e4744e653b35fccbfefb8d7c6d51896 Mon Sep 17 00:00:00 2001 From: scosman Date: Wed, 11 Sep 2024 17:11:35 -0400 Subject: [PATCH] Add created_by and config infra --- libs/core/kiln_ai/datamodel/basemodel.py | 2 + libs/core/kiln_ai/datamodel/test_basemodel.py | 6 +- libs/core/kiln_ai/utils/config.py | 62 ++++++++++ libs/core/kiln_ai/utils/test_config.py | 113 ++++++++++++++++++ 4 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 libs/core/kiln_ai/utils/config.py create mode 100644 libs/core/kiln_ai/utils/test_config.py diff --git a/libs/core/kiln_ai/datamodel/basemodel.py b/libs/core/kiln_ai/datamodel/basemodel.py index c6e4f0c..6563df6 100644 --- a/libs/core/kiln_ai/datamodel/basemodel.py +++ b/libs/core/kiln_ai/datamodel/basemodel.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Optional, Self, Type, TypeVar +from kiln_ai.utils.config import Config from pydantic import ( BaseModel, ConfigDict, @@ -32,6 +33,7 @@ class KilnBaseModel(BaseModel): v: int = 1 # schema_version path: Optional[Path] = Field(default=None, exclude=True) created_at: datetime = Field(default_factory=datetime.now) + created_by: str = Field(default_factory=lambda: Config.shared().user_id) @computed_field() def model_type(self) -> str: diff --git a/libs/core/kiln_ai/datamodel/test_basemodel.py b/libs/core/kiln_ai/datamodel/test_basemodel.py index c63a80d..1cf32ab 100644 --- a/libs/core/kiln_ai/datamodel/test_basemodel.py +++ b/libs/core/kiln_ai/datamodel/test_basemodel.py @@ -73,13 +73,17 @@ def test_type_name(): assert model.model_type == "kiln_base_model" -def test_created_at(): +def test_created_atby(): model = KilnBaseModel() assert model.created_at is not None # Check it's within 2 seconds of now now = datetime.datetime.now() assert abs((model.created_at - now).total_seconds()) < 2 + # Created by + assert len(model.created_by) > 0 + # assert model.created_by == "scosman" + # Instance of the parented model for abstract methods class NamedParentedModel(KilnParentedModel): diff --git a/libs/core/kiln_ai/utils/config.py b/libs/core/kiln_ai/utils/config.py new file mode 100644 index 0000000..0f1724a --- /dev/null +++ b/libs/core/kiln_ai/utils/config.py @@ -0,0 +1,62 @@ +import os +import pwd +from typing import Any, Callable, Dict, Optional + + +class ConfigProperty: + def __init__( + self, + type_: type, + default: Any = None, + env_var: Optional[str] = None, + default_lambda: Optional[Callable[[], Any]] = None, + ): + self.type = type_ + self.default = default + self.env_var = env_var + self.default_lambda = default_lambda + + +class Config: + _shared_instance = None + + def __init__(self, properties: Dict[str, ConfigProperty] | None = None): + self._values: Dict[str, Any] = {} + self._properties: Dict[str, ConfigProperty] = properties or { + "user_id": ConfigProperty( + str, + env_var="KILN_USER_ID", + default_lambda=lambda: pwd.getpwuid(os.getuid()).pw_name, + ), + } + + @classmethod + def shared(cls): + if cls._shared_instance is None: + cls._shared_instance = cls() + return cls._shared_instance + + def __getattr__(self, name: str) -> Any: + if name not in self._properties: + raise AttributeError(f"Config has no attribute '{name}'") + + if name not in self._values: + prop = self._properties[name] + value = None + if prop.env_var and prop.env_var in os.environ: + value = os.environ[prop.env_var] + elif prop.default is not None: + value = prop.default + elif prop.default_lambda is not None: + value = prop.default_lambda() + self._values[name] = prop.type(value) if value is not None else None + + return self._values[name] + + def __setattr__(self, name: str, value: Any) -> None: + if name in ("_values", "_properties"): + super().__setattr__(name, value) + elif name in self._properties: + self._values[name] = self._properties[name].type(value) + else: + raise AttributeError(f"Config has no attribute '{name}'") diff --git a/libs/core/kiln_ai/utils/test_config.py b/libs/core/kiln_ai/utils/test_config.py new file mode 100644 index 0000000..cc5efa5 --- /dev/null +++ b/libs/core/kiln_ai/utils/test_config.py @@ -0,0 +1,113 @@ +import os + +import pytest + +from libs.core.kiln_ai.utils.config import Config, ConfigProperty + + +def TestConfig(): + return Config( + properties={ + "example_property": ConfigProperty( + str, default="default_value", env_var="EXAMPLE_PROPERTY" + ), + } + ) + + +@pytest.fixture +def reset_config(): + Config._shared_instance = None + yield + Config._shared_instance = None + + +def test_shared_instance(reset_config): + config1 = Config.shared() + config2 = Config.shared() + assert config1 is config2 + + +def test_separate_instances(reset_config): + config1 = TestConfig() + config2 = TestConfig() + assert config1 is not config2 + + +def test_property_default_value(reset_config): + config = TestConfig() + assert config.example_property == "default_value" + + +def test_property_env_var(reset_config): + os.environ["EXAMPLE_PROPERTY"] = "env_value" + config = TestConfig() + assert config.example_property == "env_value" + del os.environ["EXAMPLE_PROPERTY"] + + +def test_property_setter(reset_config): + config = TestConfig() + config.example_property = "new_value" + assert config.example_property == "new_value" + + +def test_nonexistent_property(reset_config): + config = TestConfig() + with pytest.raises(AttributeError): + config.nonexistent_property + + +def test_property_type_conversion(reset_config): + Config._shared_instance = None + + config = Config(properties={"int_property": ConfigProperty(int, default="42")}) + assert isinstance(config.int_property, int) + assert config.int_property == 42 + + +def test_property_priority(reset_config): + os.environ["EXAMPLE_PROPERTY"] = "env_value" + config = TestConfig() + + # Environment variable takes precedence over default + assert config.example_property == "env_value" + + # Setter takes precedence over environment variable + config.example_property = "new_value" + assert config.example_property == "new_value" + + del os.environ["EXAMPLE_PROPERTY"] + + +def test_lazy_loading(reset_config): + config = TestConfig() + assert "example_property" not in config._values + _ = config.example_property + assert "example_property" in config._values + + +def test_default_lambda(reset_config): + Config._shared_instance = None + + def default_lambda(): + return "lambda_value" + + config = Config( + properties={ + "lambda_property": ConfigProperty(str, default_lambda=default_lambda) + } + ) + + assert config.lambda_property == "lambda_value" + + # Test that the lambda is only called once + assert "lambda_property" in config._values + config._properties["lambda_property"].default_lambda = lambda: "new_lambda_value" + assert config.lambda_property == "lambda_value" + + +def test_user_id_default(reset_config): + config = Config() + # assert config.user_id == "scosman" + assert len(config.user_id) > 0