From 3e3b88831248817dd41408ab98cbf2c660c3ad2a Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:17:28 +0200 Subject: [PATCH 1/7] Add context manager for managing tools --- tests/functions/test_tools.py | 79 +++++++++++++++++++++++++++++++++++ zospy/functions/__init__.py | 5 ++- zospy/functions/tools.py | 63 ++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 2 deletions(-) create mode 100644 tests/functions/test_tools.py create mode 100644 zospy/functions/tools.py diff --git a/tests/functions/test_tools.py b/tests/functions/test_tools.py new file mode 100644 index 00000000..93a40047 --- /dev/null +++ b/tests/functions/test_tools.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from zospy.functions.tools import open_tool + +if TYPE_CHECKING: + from pytest_mock import MockerFixture + + from zospy.zpcore import OpticStudioSystem + + +class TestOpenTool: + def test_open_tool(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + assert oss.Tools.CurrentTool == tool + + assert oss.Tools.CurrentTool is None + + def test_close_current(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + oss.Tools.OpenLocalOptimization() + + with open_tool(oss, oss.Tools.OpenBatchRayTrace, close_current=True) as tool: + # Testing that the previous tool was closed messes up the OpticStudio connection, + # but since only one tool can be open at a time, this should be sufficient to verify that the previous tool was closed. + assert oss.Tools.CurrentTool == tool + + assert oss.Tools.CurrentTool is None + + def test_no_close_current(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + local_optimization = oss.Tools.OpenLocalOptimization() + + with ( + pytest.raises(RuntimeError, match="Cannot open tool because another tool is already open"), + open_tool(oss, oss.Tools.OpenBatchRayTrace, close_current=False), + ): + pass + + assert oss.Tools.CurrentTool == local_optimization + + def test_closes_after_exception(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with ( + pytest.raises(ValueError, match="An error occurred while using the tool"), + open_tool(oss, oss.Tools.OpenBatchRayTrace), + ): + raise ValueError("An error occurred while using the tool") + + assert oss.Tools.CurrentTool is None + + def test_does_not_close_if_already_closed(self, oss: OpticStudioSystem, mocker: MockerFixture): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + tool.Close() + + spy = mocker.spy(tool, "Close") + + assert oss.Tools.CurrentTool is None + spy.assert_not_called() + + def test_does_not_close_other_tool(self, oss: OpticStudioSystem): + assert oss.Tools.CurrentTool is None + + with open_tool(oss, oss.Tools.OpenBatchRayTrace) as tool: + tool.Close() + + tool2 = oss.Tools.OpenLocalOptimization() + + assert oss.Tools.CurrentTool == tool2 diff --git a/zospy/functions/__init__.py b/zospy/functions/__init__.py index 29ed5907..3046a76c 100644 --- a/zospy/functions/__init__.py +++ b/zospy/functions/__init__.py @@ -4,10 +4,11 @@ - **`zospy.functions.lde`** provides helper functions for the Lens Data Editor (LDE); - **`zospy.functions.nce`** provides helper functions for the Non-sequential Component Editor (NCE). +- **`zospy.functions.tools`** provides helper functions for Tools in OpticStudio. """ from __future__ import annotations -from zospy.functions import lde, nce +from zospy.functions import lde, nce, tools -__all__ = ("lde", "nce") +__all__ = ("lde", "nce", "tools") diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py new file mode 100644 index 00000000..17f537ea --- /dev/null +++ b/zospy/functions/tools.py @@ -0,0 +1,63 @@ +"""Utility functions for Tools in OpticStudio.""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from zospy.api import _ZOSAPI + from zospy.zpcore import OpticStudioSystem + +logger = logging.getLogger(__name__) + + +@contextmanager +def open_tool( + oss: OpticStudioSystem, tool: Callable[..., _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False +) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: + """Context manager for opening a tool in OpticStudio. + + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it can + will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system to use for opening the tool. + tool : Callable[..., _ZOSAPI.Tools.ISystemTool] + A callable that returns an instance of the tool to be opened. + close_current : bool + Whether to close the currently open tool before opening the new one. Defaults to False. + + Yields + ------ + _ZOSAPI.Tools.ISystemTool + An instance of the opened tool. + + Raises + ------ + RuntimeError + If there is already an open tool and `close_current` is False. + """ + if oss.Tools.CurrentTool is not None: + if close_current: + logger.warning( + "A tool is already open. Closing the currently open tool (%s) before opening the new one.", + oss.Tools.CurrentTool.__class__.__name__, + ) + oss.Tools.CurrentTool.Close() + else: + raise RuntimeError("Cannot open tool because another tool is already open.") + + new_tool = tool() + + try: + yield new_tool + finally: + if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: + logger.info("Closing tool %s.", new_tool.__class__.__name__) + new_tool.Close() From e91317c273b51d1b64c9cc2f0bc1a0a366906419 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:22:36 +0200 Subject: [PATCH 2/7] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 080d916e..82c1ab27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ ZOS-API can also be added in patch releases. ### Added - Support for Python 3.14 (#221) +- `zospy.functions.tools.open_tool`, a context manager to open a tool and close it automatically after use (#226) ### Changed From a00fc924e3a7b43fbf4b5559be8df0c9dccc202c Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:26:40 +0200 Subject: [PATCH 3/7] Fix Copilot review comments --- zospy/functions/tools.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py index 17f537ea..72e2e49a 100644 --- a/zospy/functions/tools.py +++ b/zospy/functions/tools.py @@ -12,16 +12,18 @@ from zospy.api import _ZOSAPI from zospy.zpcore import OpticStudioSystem +__all__ = ("open_tool",) + logger = logging.getLogger(__name__) @contextmanager def open_tool( - oss: OpticStudioSystem, tool: Callable[..., _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False + oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False ) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: """Context manager for opening a tool in OpticStudio. - Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it can + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. Parameters From 142bceaca6e803b250238385ba487afb836c1b72 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:29:12 +0200 Subject: [PATCH 4/7] Add base classes for tool wrappers and implement quick focus tool --- tests/tools/__init__.py | 0 .../test_tools.py => tools/test_base.py} | 2 +- zospy/__init__.py | 3 +- zospy/analyses/__init__.py | 2 +- zospy/functions/tools.py | 65 ---- zospy/tools/__init__.py | 14 + zospy/tools/base.py | 353 ++++++++++++++++++ zospy/tools/quick_focus.py | 65 ++++ 8 files changed, 436 insertions(+), 68 deletions(-) create mode 100644 tests/tools/__init__.py rename tests/{functions/test_tools.py => tools/test_base.py} (98%) delete mode 100644 zospy/functions/tools.py create mode 100644 zospy/tools/__init__.py create mode 100644 zospy/tools/base.py create mode 100644 zospy/tools/quick_focus.py diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/functions/test_tools.py b/tests/tools/test_base.py similarity index 98% rename from tests/functions/test_tools.py rename to tests/tools/test_base.py index 93a40047..acb12892 100644 --- a/tests/functions/test_tools.py +++ b/tests/tools/test_base.py @@ -4,7 +4,7 @@ import pytest -from zospy.functions.tools import open_tool +from zospy.tools import open_tool if TYPE_CHECKING: from pytest_mock import MockerFixture diff --git a/zospy/__init__.py b/zospy/__init__.py index 6d5be33a..aba7cc1d 100644 --- a/zospy/__init__.py +++ b/zospy/__init__.py @@ -5,7 +5,7 @@ import logging from importlib.metadata import version -from zospy import analyses, functions, solvers +from zospy import analyses, functions, solvers, tools from zospy.api import config, constants from zospy.zpcore import ZOS @@ -17,6 +17,7 @@ "constants", "functions", "solvers", + "tools", ) config.set_decimal_point_and_thousands_separator() diff --git a/zospy/analyses/__init__.py b/zospy/analyses/__init__.py index 2853daa3..c04e995a 100644 --- a/zospy/analyses/__init__.py +++ b/zospy/analyses/__init__.py @@ -2,7 +2,7 @@ Analyses in OpticStudio are available through `zospy.analyses`. This module provides basic classes and functions for interacting with -analyses (in `zospy.analyses.base`), as well as a procedural interface +analyses (in `zospy.analyses.base`), as well as an object-oriented interface to several analyses. Examples diff --git a/zospy/functions/tools.py b/zospy/functions/tools.py deleted file mode 100644 index 72e2e49a..00000000 --- a/zospy/functions/tools.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Utility functions for Tools in OpticStudio.""" - -from __future__ import annotations - -import logging -from contextlib import contextmanager -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable, Generator - - from zospy.api import _ZOSAPI - from zospy.zpcore import OpticStudioSystem - -__all__ = ("open_tool",) - -logger = logging.getLogger(__name__) - - -@contextmanager -def open_tool( - oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False -) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: - """Context manager for opening a tool in OpticStudio. - - Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it - will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. - - Parameters - ---------- - oss : OpticStudioSystem - The OpticStudio system to use for opening the tool. - tool : Callable[..., _ZOSAPI.Tools.ISystemTool] - A callable that returns an instance of the tool to be opened. - close_current : bool - Whether to close the currently open tool before opening the new one. Defaults to False. - - Yields - ------ - _ZOSAPI.Tools.ISystemTool - An instance of the opened tool. - - Raises - ------ - RuntimeError - If there is already an open tool and `close_current` is False. - """ - if oss.Tools.CurrentTool is not None: - if close_current: - logger.warning( - "A tool is already open. Closing the currently open tool (%s) before opening the new one.", - oss.Tools.CurrentTool.__class__.__name__, - ) - oss.Tools.CurrentTool.Close() - else: - raise RuntimeError("Cannot open tool because another tool is already open.") - - new_tool = tool() - - try: - yield new_tool - finally: - if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: - logger.info("Closing tool %s.", new_tool.__class__.__name__) - new_tool.Close() diff --git a/zospy/tools/__init__.py b/zospy/tools/__init__.py new file mode 100644 index 00000000..10e37d35 --- /dev/null +++ b/zospy/tools/__init__.py @@ -0,0 +1,14 @@ +"""OpticStudio Tools. + +Tools in OpticStudio are available through `zospy.tools`. +This module provides basic classes and functions for interacting with +tools (in `zospy.tools.base`), as well as an object-oriented interface +to several tools. +""" + +from __future__ import annotations + +from zospy.tools.base import open_tool +from zospy.tools.quick_focus import QuickFocus, QuickFocusSettings + +__all__ = ("QuickFocus", "QuickFocusSettings", "open_tool") diff --git a/zospy/tools/base.py b/zospy/tools/base.py new file mode 100644 index 00000000..c9f88861 --- /dev/null +++ b/zospy/tools/base.py @@ -0,0 +1,353 @@ +"""Base functionality for tools.""" + +from __future__ import annotations + +import dataclasses +import logging +import weakref +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import is_dataclass +from types import NoneType +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, get_args + +import numpy as np +import pandas as pd +import pydantic + +from zospy.analyses.base import _deserialize_analysis_data, _deserialize_zospy_class, _serialize_analysis_data_type +from zospy.analyses.parsers.types import ValidatedDataFrame + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from zospy.api import _ZOSAPI + from zospy.zpcore import OpticStudioSystem + + +__all__ = ("BaseToolWrapper", "ToolResult", "open_tool") + +logger = logging.getLogger(__name__) + + +@contextmanager +def open_tool( + oss: OpticStudioSystem, tool: Callable[[], _ZOSAPI.Tools.ISystemTool], *, close_current: bool = False +) -> Generator[_ZOSAPI.Tools.ISystemTool, Any, None]: + """Context manager for opening a tool in OpticStudio. + + Opens a tool in OpticStudio and ensures that it is properly closed after use. If there is already an open tool, it + will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system to use for opening the tool. + tool : Callable[..., _ZOSAPI.Tools.ISystemTool] + A callable that returns an instance of the tool to be opened. + close_current : bool + Whether to close the currently open tool before opening the new one. Defaults to False. + + Yields + ------ + _ZOSAPI.Tools.ISystemTool + An instance of the opened tool. + + Raises + ------ + RuntimeError + If there is already an open tool and `close_current` is False. + """ + if oss.Tools.CurrentTool is not None: + if close_current: + logger.warning( + "A tool is already open. Closing the currently open tool (%s) before opening the new one.", + oss.Tools.CurrentTool.__class__.__name__, + ) + oss.Tools.CurrentTool.Close() + else: + raise RuntimeError("Cannot open tool because another tool is already open.") + + new_tool = tool() + + try: + yield new_tool + finally: + if oss.Tools.CurrentTool is not None and oss.Tools.CurrentTool == new_tool: + logger.info("Closing tool %s.", new_tool.__class__.__name__) + new_tool.Close() + + +ToolOutputData = TypeVar("ToolOutputData") +ToolSettings = TypeVar("ToolSettings") + + +@pydantic.dataclasses.dataclass(frozen=True, config=pydantic.ConfigDict(ser_json_inf_nan="constants")) +class ToolResult(Generic[ToolOutputData, ToolSettings]): + """Zemax OpticStudio tool result. + + Attributes + ---------- + data : ToolOutputData + The data of the tool. Can be a `pandas.DataFrame`, `numpy.ndarray`, or a tool-specific dataclass. + settings : ToolSettings | None + The settings of the tool. + error_message : str | None + Error message from the analysis. If the tool ran successfully, this will be None. + """ + + data: ToolOutputData + settings: ToolSettings | None + error_message: str | None + + def to_json(self): + """Convert the result to a JSON string.""" + return pydantic.RootModel(self).model_dump_json(indent=4) + + @classmethod + def from_json(cls, data: str): + """Create a result from a JSON string.""" + return pydantic.TypeAdapter(cls).validate_json(data) + + @pydantic.field_serializer("data", mode="wrap", when_used="json") + @staticmethod + def _serialize_data( + value: ToolOutputData, + nxt: pydantic.SerializerFunctionWrapHandler, + info, # noqa: ARG004 + ): + if isinstance(value, pd.DataFrame): + return pydantic.TypeAdapter( + ValidatedDataFrame, config=pydantic.ConfigDict(ser_json_inf_nan="constants") + ).dump_python(value, mode="json") + + if isinstance(value, np.ndarray): + return value.tolist() + + return nxt(value) + + @pydantic.model_serializer(mode="wrap", when_used="json") + def _serialize_types(self, nxt: pydantic.SerializerFunctionWrapHandler): + data = nxt(self) + data["__tool_data__"] = _serialize_analysis_data_type(self.data) + data["__tool_settings__"] = { + "data_type": "zospy_class", + "name": type(self.settings).__name__, + "module": type(self.settings).__module__, + } + + return data + + @pydantic.model_validator(mode="wrap") + @classmethod + def _deserialize_types(cls, data: Any, handler): + if isinstance(data, dict): + if "__tool_data__" in data: + data["data"] = _deserialize_analysis_data(data["data"], data.pop("__tool_data__")) + if "__tool_settings__" in data: + data["settings"] = _deserialize_zospy_class(data["settings"], data.pop("__tool_settings__")) + + return handler(data) + + +class BaseToolWrapper(ABC, Generic[ToolOutputData, ToolSettings]): + """Base class for tool wrappers. + + This class provides a common interface for all tool wrappers. It defines the methods and properties that all + tool wrappers should implement. + + Attributes + ---------- + CONNECTION_MODE : {"standalone", "extension"} | None + The connection mode required for the tool. If None, the tool can be used in any mode. + """ + + CONNECTION_MODE: Literal["standalone", "extension"] | None = None + + def __init__(self, *, settings_kws: dict[str, Any] | None = None): + """Create a new tool wrapper. + + Settings can be changed by passing the settings as keyword arguments. Use the `with_settings` method to specify + the settings using a settings object. + + Parameters + ---------- + settings_kws : dict[str, Any] + Arguments to set the settings of the tool. + + Raises + ------ + ValueError + If `settings` is not a dataclass. + """ + self._settings = self._default_settings() + self.update_settings(settings_kws=settings_kws) + + self._oss = None + + def __init_subclass__( + cls, + *, + connection_mode: Literal["standalone", "extension"] | None = None, + **kwargs, + ): + """Determine the settings type and class-level configuration of the tool.""" + cls.CONNECTION_MODE = connection_mode + + if not hasattr(cls, "_settings_type"): + if hasattr(cls, "__orig_bases__"): + base = cls.__orig_bases__[0] + cls._settings_type: type[ToolSettings] = get_args(base)[1] + else: + cls._settings_type = NoneType + + super().__init_subclass__(**kwargs) + + @abstractmethod + def _get_tool_opener(self, oss: OpticStudioSystem) -> Callable[[], _ZOSAPI.Tools.ISystemTool]: + """Callable that opens the tool in OpticStudio and returns the tool object.""" + raise NotImplementedError("Subclasses of BaseToolWrapper must implement the `tool_opener` property.") + + def update_settings( + self, + *, + settings: ToolSettings | None = None, + settings_kws: dict[str, Any] | None = None, + ) -> None: + """Update the settings of the tool using a settings object or keyword arguments. + + Settings can be specified as an object and as keyword arguments. If both are specified, the keyword arguments + take precedence. If no settings are specified, the default settings are used. Furthermore, instead of using + a reference to the settings object, a new settings object is created with the specified parameters. This is done + to avoid modifying the original settings object. + + Parameters + ---------- + settings : ToolSettings + Tool settings object. + settings_kws + Dictionary with the settings parameters. + + Raises + ------ + ValueError + If `settings` is not a dataclass. + """ + # Use the existing settings if no settings are specified + settings = settings or self.settings + + if settings is None: + # Analysis does not have settings + return + + if not is_dataclass(settings): + raise TypeError("settings should be a dataclass.") + + # Create a new settings object with the specified parameters. If no parameters are specified, this creates a + # copy of the settings object. This is done to avoid modifying the original settings object. + self._settings = dataclasses.replace(settings, **(settings_kws or {})) + + @classmethod + def _default_settings(cls) -> ToolSettings: + """Get the default settings of the tool. + + Returns + ------- + ToolSettings + The default settings. + """ + return cls._settings_type() + + @classmethod + def with_settings(cls, settings: ToolSettings): + """Create a new tool with the specified settings. + + Parameters + ---------- + settings : ToolSettings + Settings of the tool. + + Returns + ------- + BaseAnalysisWrapper + The tool wrapper. + """ + instance = cls() + instance.update_settings(settings=settings) + + return instance + + @property + def settings(self) -> ToolSettings: + """Settings of the tool.""" + return self._settings + + @property + def oss(self) -> OpticStudioSystem: + """The OpticStudio system. This property is set when the tool is run.""" + if self._oss is None: + raise ValueError("OpticStudioSystem has not been set.") + + return self._oss + + @abstractmethod + def _run_tool(self, tool: _ZOSAPI.Tools.ISystemTool, *args, **kwargs) -> ToolOutputData: + """Run the tool and return the results.""" + + def _check_mode(self): + connection_modes: dict[str, str] = { + "Server": "standalone", + "Plugin": "extension", + } + + if self.CONNECTION_MODE is None: + return + + if str(self.oss.ZOS.Application.Mode) not in connection_modes: + raise ValueError(f"Unknown connection mode: {self.oss.ZOS.Application.Mode}") + + current_mode = connection_modes[str(self.oss.ZOS.Application.Mode)] + + if current_mode != self.CONNECTION_MODE: + raise ValueError(f"The tool requires {self.CONNECTION_MODE} mode, got {current_mode}.") + + def run( + self, + oss: OpticStudioSystem, + *, + close_current: bool = False, + ) -> ToolResult[ToolOutputData, ToolSettings]: + """Run the tool and return the results. + + This method opens and runs the tool in OpticStudio with the specified settings. + If another tool is already open, it will be closed if `close_current` is True. Otherwise, a RuntimeError is raised. + + Parameters + ---------- + oss : OpticStudioSystem + The OpticStudio system. + close_current : bool + Whether to close the current tool if one is already open. + + Returns + ------- + ToolResult + The tool result. + """ + self._oss = weakref.proxy(oss) + self._check_mode() + + with open_tool(oss, self._get_tool_opener(oss), close_current=close_current) as tool: + data = self._run_tool(tool) + + error_message = tool.ErrorMessage + + return ToolResult( + data=data, + settings=self.settings, + error_message=error_message, + ) + + def __call__(self, oss: OpticStudioSystem, *args, **kwargs): + """Run the tool and return the results.""" + return self.run(oss, *args, **kwargs) diff --git a/zospy/tools/quick_focus.py b/zospy/tools/quick_focus.py new file mode 100644 index 00000000..580965a5 --- /dev/null +++ b/zospy/tools/quick_focus.py @@ -0,0 +1,65 @@ +"""Quick Focus Tool.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pydantic import Field + +from zospy.analyses.decorators import analysis_settings +from zospy.analyses.parsers.types import ZOSAPIConstant # noqa: TC001 +from zospy.api import constants +from zospy.api.constants import process_constant +from zospy.tools.base import BaseToolWrapper + +if TYPE_CHECKING: + from collections.abc import Callable + + from zospy.api import _ZOSAPI + +__all__ = ("QuickFocus", "QuickFocusSettings") + + +@analysis_settings +class QuickFocusSettings: + """Settings for the Quick Focus tool. + + Attributes + ---------- + criterion : constants.Tools.General.QuickFocusCriterion | str + The criterion to use for quick focusing. Defaults to 'SpotSizeRadial'. + use_centroid : bool + Reference all calculations to the image centroid rather than the chief ray. Defaults to True. + """ + + criterion: ZOSAPIConstant("Tools.General.QuickFocusCriterion") = Field(default="SpotSizeRadial") + use_centroid: bool = Field(default=True) + + +class QuickFocus(BaseToolWrapper[None, QuickFocusSettings]): + """Wrapper for the Quick Focus tool.""" + + def __init__( + self, + *, + criterion: str | constants.Tools.General.QuickFocusCriterion = "SpotSizeRadial", + use_centroid: bool = True, + ) -> None: + """Initialize the Quick Focus tool. + + See Also + -------- + QuickFocusSettings : Settings for the Quick Focus tool. + """ + super().__init__(settings_kws=locals()) + + def _get_tool_opener(self, oss) -> Callable[[], _ZOSAPI.Tools.General.IQuickFocus]: + """Get a callable that opens the Quick Focus tool in OpticStudio and returns the tool object.""" + return oss.Tools.OpenQuickFocus + + def _run_tool(self, tool: _ZOSAPI.Tools.General.IQuickFocus) -> None: + """Run the Quick Focus tool.""" + tool.Criterion = process_constant(constants.Tools.General.QuickFocusCriterion, self.settings.criterion) + tool.UseCentroid = self.settings.use_centroid + + tool.RunAndWaitForCompletion() From 162fbec53f89c8771442a09a3d2f65873f8a7021 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:34:01 +0200 Subject: [PATCH 5/7] Add test for quick focus tool --- tests/tools/test_quick_focus.py | 19 +++++++++++++++++++ zospy/functions/__init__.py | 4 ++-- 2 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 tests/tools/test_quick_focus.py diff --git a/tests/tools/test_quick_focus.py b/tests/tools/test_quick_focus.py new file mode 100644 index 00000000..69b57f67 --- /dev/null +++ b/tests/tools/test_quick_focus.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from zospy.tools import QuickFocus + +if TYPE_CHECKING: + from zospy.zpcore import OpticStudioSystem + + +def test_quick_focus(simple_system: OpticStudioSystem): + """Test the Quick Focus tool.""" + simple_system.LDE.GetSurfaceAt(3).Thickness = 30 + + QuickFocus().run(simple_system) + + assert simple_system.LDE.GetSurfaceAt(3).Thickness == pytest.approx(19.777975) diff --git a/zospy/functions/__init__.py b/zospy/functions/__init__.py index 3046a76c..e974a3bb 100644 --- a/zospy/functions/__init__.py +++ b/zospy/functions/__init__.py @@ -9,6 +9,6 @@ from __future__ import annotations -from zospy.functions import lde, nce, tools +from zospy.functions import lde, nce -__all__ = ("lde", "nce", "tools") +__all__ = ("lde", "nce") From 1d28bbb42641c0de14f90da0115cec2965038307 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 15:18:30 +0200 Subject: [PATCH 6/7] Add tests for tools base functionality --- tests/analyses/conftest.py | 5 +- tests/helpers.py | 4 + tests/tools/test_base.py | 219 ++++++++++++++++++++++++++++++++++++- zospy/analyses/base.py | 8 +- zospy/tools/base.py | 8 +- 5 files changed, 232 insertions(+), 12 deletions(-) diff --git a/tests/analyses/conftest.py b/tests/analyses/conftest.py index 1fe2ea16..cbd5989e 100644 --- a/tests/analyses/conftest.py +++ b/tests/analyses/conftest.py @@ -7,6 +7,7 @@ import pytest from tests.config import CONFIG_DATA_FOLDER, REFERENCE_DATA_FOLDER, REFERENCE_VERSION +from tests.helpers import _all_subclasses from zospy.analyses.base import AnalysisResult, BaseAnalysisWrapper from zospy.analyses.psf.huygens_psf import BaseHuygensPSF from zospy.analyses.systemviewers.base import SystemViewerWrapper @@ -56,10 +57,6 @@ def reference_data(request) -> AnalysisResult: return AnalysisResult.from_json(data_file.read_text(encoding="utf-8")) -def _all_subclasses(cls): - return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in _all_subclasses(c)]) - - _analysis_wrapper_classes = [ c for c in _all_subclasses(BaseAnalysisWrapper) if c not in {BaseHuygensPSF, SystemViewerWrapper} ] diff --git a/tests/helpers.py b/tests/helpers.py index 01281c5c..abcb02f1 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -66,3 +66,7 @@ def assert_dataclass_equal( return True return None + + +def _all_subclasses(cls): + return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in _all_subclasses(c)]) diff --git a/tests/tools/test_base.py b/tests/tools/test_base.py index acb12892..135a8342 100644 --- a/tests/tools/test_base.py +++ b/tests/tools/test_base.py @@ -1,12 +1,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import inspect +import json +from dataclasses import dataclass, fields +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +import numpy as np import pytest +from pandas import DataFrame +from pydantic import Field +from pydantic.fields import FieldInfo +from tests.helpers import _all_subclasses +from zospy.analyses.decorators import analysis_settings +from zospy.analyses.parsers.types import ValidatedDataFrame from zospy.tools import open_tool +from zospy.tools.base import BaseToolWrapper, ToolResult, ToolSettings +from zospy.tools.quick_focus import QuickFocusSettings if TYPE_CHECKING: + from collections.abc import Callable + from pytest_mock import MockerFixture from zospy.zpcore import OpticStudioSystem @@ -77,3 +92,205 @@ def test_does_not_close_other_tool(self, oss: OpticStudioSystem): tool2 = oss.Tools.OpenLocalOptimization() assert oss.Tools.CurrentTool == tool2 + + +@dataclass +class MockToolOutputData: + int_data: int = 1 + string_data: str = "a" + + +@analysis_settings +class MockToolSettings: + int_setting: int = Field(default=1, description="An integer setting") + string_setting: str = Field(default="a", description="A string setting") + + +class MockTool(BaseToolWrapper[MockToolOutputData, MockToolSettings]): + def __init__( + self, + *, + int_setting: int = 1, + string_setting: str = "a", + ): + super().__init__(settings_kws=locals()) + + def _get_tool_opener(self, oss: OpticStudioSystem) -> Callable[[], Any]: # noqa: ARG002 + return lambda: SimpleNamespace( + Close=lambda: None, ErrorMessage=None + ) # This tool does not actually open anything in OpticStudio + + def _run_tool(self) -> MockToolOutputData: + return MockToolOutputData() + + +_tool_wrapper_classes = [c for c in _all_subclasses(BaseToolWrapper) if c not in {MockTool}] + + +@pytest.fixture(scope="module", params=_tool_wrapper_classes) +def tool_wrapper_class(request): + return request.param + + +class TestToolWrapper: + @staticmethod + def get_settings_defaults(settings_class): + result = {} + + for field in fields(settings_class): + if isinstance(field.default, FieldInfo): + result[field.name] = field.default.default + else: + result[field.name] = field.default + + return result + + def test_get_settings_type(self): + assert MockTool._settings_type == MockToolSettings # noqa: SLF001 + + def test_settings_type_is_specified(self): + assert MockTool._settings_type is not ToolSettings # noqa: SLF001 + + def test_tools_correct_tool_opener(self, oss: OpticStudioSystem, tool_wrapper_class): + instance = tool_wrapper_class() + + try: + instance._get_tool_opener(oss) # noqa: SLF001 + except Exception: # noqa: BLE001 + pytest.fail(f"{tool_wrapper_class.__name__} does not implement _get_tool_opener correctly.") + + def test_init_all_keyword_only_parameters(self, tool_wrapper_class): + assert all(p.kind.name == "KEYWORD_ONLY" for p in inspect.signature(tool_wrapper_class).parameters.values()) + + def test_init_contains_all_settings(self, tool_wrapper_class): + if tool_wrapper_class().settings is None: + return + + init_signature = inspect.signature(tool_wrapper_class.__init__) + settings_fields = fields(tool_wrapper_class().settings) + + assert all(field.name in init_signature.parameters for field in settings_fields) + + def test_tools_default_values(self, tool_wrapper_class): + if tool_wrapper_class().settings is None: + return + + settings_defaults = self.get_settings_defaults(type(tool_wrapper_class().settings)) + init_signature = inspect.signature(tool_wrapper_class.__init__) + + for field_name, default_value in settings_defaults.items(): + assert field_name in init_signature.parameters + assert init_signature.parameters[field_name].default == default_value + + def test_change_settings_from_parameters(self): + tool = MockTool(int_setting=2, string_setting="b") + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_change_settings_from_object(self): + settings = MockToolSettings(int_setting=2, string_setting="b") + tool = MockTool().with_settings(settings) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_settings_object_is_copied(self): + settings = MockToolSettings(int_setting=2, string_setting="b") + tool = MockTool().with_settings(settings) + + assert tool.settings is not settings + assert tool.settings == settings + + def test_update_settings_object(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings(settings=MockToolSettings(int_setting=2, string_setting="b")) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_dictionary(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings(settings_kws={"int_setting": 2, "string_setting": "b"}) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_object_and_dictionary(self): + tool = MockTool(int_setting=1, string_setting="a") + + tool.update_settings( + settings=MockToolSettings(int_setting=2, string_setting="a"), settings_kws={"string_setting": "b"} + ) + + assert tool.settings.int_setting == 2 + assert tool.settings.string_setting == "b" + + def test_update_settings_no_dataclass_raises_type_error(self): + with pytest.raises(TypeError, match="settings should be a dataclass"): + MockTool().update_settings(settings=123) + + +class TestAnalysisResultJSONConversion: + # Only test for non-dataclass results, because dataclass results are tested separately in the corresponding + # analysis' tests. + @pytest.mark.parametrize( + "result_type,result_value,type_info", + [ + (DataFrame, DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), {"data_type": "dataframe"}), + (np.ndarray, np.array([1, 2, 3]), {"data_type": "ndarray"}), + ], + ) + def test_data_to_json(self, result_type, result_value, type_info): + result = ToolResult[result_type, MockToolSettings]( + data=result_value, + settings=MockToolSettings(), + error_message=None, + ) + + result_json = result.to_json() + result_dict = json.loads(result_json) + + assert "__tool_data__" in result_dict + assert result_dict["__tool_data__"] == type_info + + def test_settings_to_json(self): + result = ToolResult[ValidatedDataFrame, MockToolSettings]( + data=DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), + settings=MockToolSettings(), + error_message=None, + ) + + result_json = result.to_json() + result_dict = json.loads(result_json) + + assert "__tool_settings__" in result_dict + assert result_dict["__tool_settings__"] == { + "data_type": "zospy_class", + "name": "MockToolSettings", + "module": "tests.tools.test_base", + } + + @pytest.mark.parametrize( + "result_type,result_value", + [ + (ValidatedDataFrame, DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})), + (np.ndarray, np.array([1, 2, 3])), + ], + ) + def test_roundtrip(self, result_type, result_value): + result = ToolResult[result_type, QuickFocusSettings]( + data=result_value, + settings=QuickFocusSettings(), + error_message=None, + ) + + result_json = result.to_json() + + result_roundtrip = ToolResult.from_json(result_json) + assert all(result_roundtrip.data == result.data) + assert str(result_roundtrip.settings.criterion) == str(result.settings.criterion) + assert result_roundtrip.settings.use_centroid == result.settings.use_centroid + assert result_roundtrip.error_message == result.error_message diff --git a/zospy/analyses/base.py b/zospy/analyses/base.py index 4a6507e3..63e42d6b 100644 --- a/zospy/analyses/base.py +++ b/zospy/analyses/base.py @@ -137,8 +137,8 @@ def _serialize_analysis_data_type(data: AnalysisData) -> _TypeInfo: raise ValueError(f"Cannot serialize data type: {type(data)}") -def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo) -> AnalysisData: - if typeinfo["module"].startswith("zospy.analyses"): +def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo, module: str) -> object: + if typeinfo["module"].startswith(module): try: m = import_module(typeinfo["module"]) t = getattr(m, typeinfo["name"]) @@ -150,7 +150,7 @@ def _deserialize_zospy_class(data: dict, typeinfo: _TypeInfo) -> AnalysisData: return data -def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> AnalysisData: +def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> object: if typeinfo["data_type"] == "none": return None @@ -161,7 +161,7 @@ def _deserialize_analysis_data(data: dict | list, typeinfo: _TypeInfo) -> Analys return np.array(data) if typeinfo["data_type"] == "zospy_class": - return _deserialize_zospy_class(data, typeinfo) + return _deserialize_zospy_class(data, typeinfo, module="zospy.analyses") raise ValueError(f"Cannot deserialize data type: {typeinfo['data_type']}") diff --git a/zospy/tools/base.py b/zospy/tools/base.py index c9f88861..c99dfba7 100644 --- a/zospy/tools/base.py +++ b/zospy/tools/base.py @@ -90,14 +90,14 @@ class ToolResult(Generic[ToolOutputData, ToolSettings]): ---------- data : ToolOutputData The data of the tool. Can be a `pandas.DataFrame`, `numpy.ndarray`, or a tool-specific dataclass. - settings : ToolSettings | None + settings : ToolSettings The settings of the tool. error_message : str | None Error message from the analysis. If the tool ran successfully, this will be None. """ data: ToolOutputData - settings: ToolSettings | None + settings: ToolSettings error_message: str | None def to_json(self): @@ -145,7 +145,9 @@ def _deserialize_types(cls, data: Any, handler): if "__tool_data__" in data: data["data"] = _deserialize_analysis_data(data["data"], data.pop("__tool_data__")) if "__tool_settings__" in data: - data["settings"] = _deserialize_zospy_class(data["settings"], data.pop("__tool_settings__")) + data["settings"] = _deserialize_zospy_class( + data["settings"], data.pop("__tool_settings__"), module="zospy.tools" + ) return handler(data) From aa1d3e9ef8fda29b9d52379c71da833e0c00d434 Mon Sep 17 00:00:00 2001 From: crnh <30109443+crnh@users.noreply.github.com> Date: Fri, 5 Jun 2026 19:38:06 +0200 Subject: [PATCH 7/7] Update changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 82c1ab27..b5f12997 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,7 @@ ZOS-API can also be added in patch releases. ### Added - Support for Python 3.14 (#221) -- `zospy.functions.tools.open_tool`, a context manager to open a tool and close it automatically after use (#226) +- New `zospy.tools` submodule with `open_tool` (a context manager to open a tool and close it automatically after use) and tool wrappers (#226) ### Changed