Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make Settings Pydantic and use the power of BaseSettings to simplify CLI #700

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 33 additions & 16 deletions src/blueapi/cli/cli.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
import logging
import sys
from collections.abc import Sequence
from functools import wraps
from pathlib import Path
from pprint import pprint
Expand All @@ -11,45 +12,60 @@
from bluesky_stomp.models import Broker
from observability_utils.tracing import setup_tracing
from pydantic import ValidationError
from pydantic_settings.sources import PathType
from requests.exceptions import ConnectionError

from blueapi import __version__
from blueapi.cli.format import OutputFormat
from blueapi.client.client import BlueapiClient
from blueapi.client.event_bus import AnyEvent, BlueskyStreamingError, EventBusClient
from blueapi.client.rest import BlueskyRemoteControlError
from blueapi.config import ApplicationConfig, ConfigLoader
from blueapi.config import ApplicationConfig
from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent
from blueapi.worker import ProgressEvent, Task, WorkerEvent

from .scratch import setup_scratch
from .updates import CliEventRenderer


def parse_path_type(path_type: PathType) -> list[Path]:
"""
Parse a PathType parameter and return a list of Path objects.

:param path_type: The input which can be a Path, str, or a sequence of Path/str.
:return: A list of Path objects.
"""
if isinstance(path_type, str | Path):
# Single Path or string: Convert to Path and return as a single-element list
return [Path(path_type)]

if isinstance(path_type, Sequence):
# Sequence of Paths/strings: Convert each element to a Path
return [Path(item) for item in path_type if isinstance(item, str | Path)]

# If it doesn't match the expected types, raise an error
raise TypeError(f"Unsupported PathType: {type(path_type)}")


@click.group(invoke_without_command=True)
@click.version_option(version=__version__, prog_name="blueapi")
@click.option(
"-c", "--config", type=Path, help="Path to configuration YAML file", multiple=True
)
@click.pass_context
def main(ctx: click.Context, config: Path | None | tuple[Path, ...]) -> None:
# if no command is supplied, run with the options passed

config_loader = ConfigLoader(ApplicationConfig)
if config is not None:
configs = (config,) if isinstance(config, Path) else config
for path in configs:
if path.exists():
config_loader.use_values_from_yaml(path)
else:
raise FileNotFoundError(f"Cannot find file: {path}")

def main(ctx: click.Context, config: PathType) -> None:
# Override default yaml_file path in the model_config if `config` is provided
ApplicationConfig.model_config["yaml_file"] = config
DiamondJoseph marked this conversation as resolved.
Show resolved Hide resolved
app_config = ApplicationConfig() # Instantiates with customized sources
Comment on lines +58 to +59
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should: I'm concerned this. mutation of global state is an antipattern. Could you instead subclass the config?

Suggested change
ApplicationConfig.model_config["yaml_file"] = config
app_config = ApplicationConfig() # Instantiates with customized sources
class YamlConfig(ApplicationConfig):
model_config = {**ApplicationConfig.model_config, **{"yaml_file": config}}
app_config = YamlConfig()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that's how Pydantic uses it and this is not regular state but settings

also this is not Java

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mutation of global state is an antipattern in any programming language (good thread).

Now I have to remember to address this if I use ApplicationConfig somewhere else, just like you have:

ApplicationConfig.model_config["yaml_file"] = None

There is nothing that makes this obvious to me as a developer.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looking into the base settings docs now

re: mutation - it was there anyway and this config is just the init, and anywhere else you'd sooner use configManager getter methods, should you need access to the config.

image

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the mutable settings option is here:
and we're not using it
https://docs.pydantic.dev/latest/concepts/pydantic_settings/#removing-sources

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like what we're trying to do was just not a well-envisioned use case: pydantic/pydantic-settings#346

Should maybe reconsider the use of BaseSettings and/or raise a PR upstream?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that is true, turns out that pydantic-settings isn't as well-finished product like pydantic itself

print(f"Loaded configuration {app_config}")
ctx.ensure_object(dict)
loaded_config: ApplicationConfig = config_loader.load()
ctx.obj["config"] = app_config

ctx.obj["config"] = loaded_config
# note: this is the key result of the 'main' function, it loaded the config
# and due to 'pass context' flag above
# it's left for the handler of words that are later in the stdin
logging.basicConfig(
format="%(asctime)s - %(message)s", level=loaded_config.logging.level
format="%(asctime)s - %(message)s", level=app_config.logging.level
)

if ctx.invoked_subcommand is None:
Expand Down Expand Up @@ -173,6 +189,7 @@ def listen_to_events(obj: dict) -> None:
)
)
)

fmt = obj["fmt"]

def on_event(
Expand Down
109 changes: 23 additions & 86 deletions src/blueapi/config.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
from collections.abc import Mapping
from enum import Enum
from pathlib import Path
from typing import Any, Generic, Literal, TypeVar
from typing import Literal

import yaml
from bluesky_stomp.models import BasicAuthentication
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings, SettingsConfigDict, YamlConfigSettingsSource

from blueapi.utils import BlueapiBaseModel, InvalidConfigError
from blueapi.utils import BlueapiBaseModel

LogLevel = Literal["NOTSET", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]

DEFAULT_PATH = Path("config.yaml") # Default YAML file path


class SourceKind(str, Enum):
PLAN_FUNCTIONS = "planFunctions"
Expand Down Expand Up @@ -77,7 +78,8 @@ class ScratchConfig(BlueapiBaseModel):
repositories: list[ScratchRepository] = Field(default_factory=list)


class ApplicationConfig(BlueapiBaseModel):
# class ApplicationConfig(BaseSettings, cli_parse_args=True, cli_prog_name="blueapi"):
class ApplicationConfig(BaseSettings):
"""
Config for the worker application as a whole. Root of
config tree.
Expand All @@ -89,83 +91,18 @@ class ApplicationConfig(BlueapiBaseModel):
api: RestConfig = Field(default_factory=RestConfig)
scratch: ScratchConfig | None = None

def __eq__(self, other: object) -> bool:
if isinstance(other, ApplicationConfig):
return (
(self.stomp == other.stomp)
& (self.env == other.env)
& (self.logging == other.logging)
& (self.api == other.api)
)
return False


C = TypeVar("C", bound=BaseModel)


class ConfigLoader(Generic[C]):
"""
Small utility class for loading config from various sources.
You must define a config schema as a dataclass (or series of
nested dataclasses) that can then be loaded from some combination
of default values, dictionaries, YAML/JSON files etc.
"""

def __init__(self, schema: type[C]) -> None:
self._adapter = TypeAdapter(schema)
self._values: dict[str, Any] = {}

def use_values(self, values: Mapping[str, Any]) -> None:
"""
Use all values provided in the config, override any defaults
and values set by previous calls into this class.

Args:
values (Mapping[str, Any]): Dictionary of override values,
does not need to be exhaustive
if defaults provided.
"""

def recursively_update_map(old: dict[str, Any], new: Mapping[str, Any]) -> None:
for key in new:
if (
key in old
and isinstance(old[key], dict)
and isinstance(new[key], dict)
):
recursively_update_map(old[key], new[key])
else:
old[key] = new[key]

recursively_update_map(self._values, values)

def use_values_from_yaml(self, path: Path) -> None:
"""
Use all values provided in a YAML/JSON file in the
config, override any defaults and values set by
previous calls into this class.

Args:
path (Path): Path to YAML/JSON file
"""

with path.open("r") as stream:
values = yaml.load(stream, yaml.Loader)
self.use_values(values)

def load(self) -> C:
"""
Finalize and load the config as an instance of the `schema`
dataclass.

Returns:
C: Dataclass instance holding config
"""

try:
return self._adapter.validate_python(self._values)
except ValidationError as exc:
error_details = "\n".join(str(e) for e in exc.errors())
raise InvalidConfigError(
f"Something is wrong with the configuration file: \n {error_details}"
) from exc
model_config = SettingsConfigDict(
env_nested_delimiter="__", yaml_file=DEFAULT_PATH, yaml_file_encoding="utf-8"
)

DiamondJoseph marked this conversation as resolved.
Show resolved Hide resolved
@classmethod
def settings_customize_sources(
cls, init_settings, env_settings, file_secret_settings
):
path = cls.model_config.get("yaml_file")
return (
init_settings,
YamlConfigSettingsSource(settings_cls=cls, yaml_file=path),
env_settings,
file_secret_settings,
)
26 changes: 26 additions & 0 deletions src/blueapi/service/config_manager.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
# config_manager.py

from blueapi.config import ApplicationConfig


class ConfigManager:
"""Manages application configuration in a way that’s easy to test and mock."""

_config: ApplicationConfig

def __init__(self, config: ApplicationConfig = None):
if config is None:
ApplicationConfig.model_config["yaml_file"] = None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should: Again, I think you should find an alternative to mutating global state, if you fix the instance in CLI then you shouldn't need this line.

config = ApplicationConfig()
self._config = config

def get_config(self) -> ApplicationConfig:
"""Retrieve the current configuration."""
return self._config

def set_config(self, new_config: ApplicationConfig):
"""
This is a setter function that the main process uses
to pass the config into the subprocess
"""
self._config = new_config
24 changes: 11 additions & 13 deletions src/blueapi/service/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,20 +9,18 @@
from blueapi.config import ApplicationConfig, StompConfig
from blueapi.core.context import BlueskyContext
from blueapi.core.event import EventStream
from blueapi.service.config_manager import ConfigManager
from blueapi.service.model import DeviceModel, PlanModel, WorkerTask
from blueapi.worker.event import TaskStatusEnum, WorkerState
from blueapi.worker.task import Task
from blueapi.worker.task_worker import TaskWorker, TrackableTask

"""This module provides interface between web application and underlying Bluesky
context and worker"""
context and worker

"""

_CONFIG: ApplicationConfig = ApplicationConfig()


def config() -> ApplicationConfig:
return _CONFIG
config_manager = ConfigManager()


def set_config(new_config: ApplicationConfig):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could: Can we remove this now?

Expand All @@ -34,23 +32,23 @@ def set_config(new_config: ApplicationConfig):
@cache
def context() -> BlueskyContext:
ctx = BlueskyContext()
ctx.with_config(config().env)
env_config = config_manager.get_config().env
ctx.with_config(env_config)
return ctx


@cache
def worker() -> TaskWorker:
worker = TaskWorker(
context(),
broadcast_statuses=config().env.events.broadcast_status_events,
)
env_config = config_manager.get_config().env
should_broadcast_status_events: bool = env_config.events.broadcast_status_events
worker = TaskWorker(context(), broadcast_statuses=should_broadcast_status_events)
worker.start()
return worker


@cache
def stomp_client() -> StompClient | None:
stomp_config: StompConfig | None = config().stomp
stomp_config: StompConfig | None = config_manager.get_config().stomp
if stomp_config is not None:
client = StompClient.for_broker(
broker=Broker(
Expand Down Expand Up @@ -79,7 +77,7 @@ def stomp_client() -> StompClient | None:
def setup(config: ApplicationConfig) -> None:
"""Creates and starts a worker with supplied config"""

set_config(config)
config_manager.set_config(config)

# Eagerly initialize worker and messaging connection

Expand Down
2 changes: 0 additions & 2 deletions src/blueapi/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from .base_model import BlueapiBaseModel, BlueapiModelConfig, BlueapiPlanModelConfig
from .invalid_config_error import InvalidConfigError
from .modules import load_module_all
from .serialization import serialize
from .thread_exception import handle_all_exceptions
Expand All @@ -11,5 +10,4 @@
"BlueapiBaseModel",
"BlueapiModelConfig",
"BlueapiPlanModelConfig",
"InvalidConfigError",
]
3 changes: 0 additions & 3 deletions src/blueapi/utils/invalid_config_error.py

This file was deleted.

Loading
Loading