Decorator for analyses

[crnh]

Currently, a lot of code is duplicated between all analyses. This has been improved by the Analysis.complete method, but still many analyses contain the same code for handling input and output files and retrieving header data, metadata and error messages.
I briefly considered using analysis classes instead of analysis functions, but this would change the API, which I would like to avoid. Furthermore, I prefer a procedural interface over an object-oriented interface for analyses.

Based on a suggestion by @andibarg, I propose to write a decorator for analyses. This decorator creates the analysis, handles txtoutfile and cfgoutfile, runs the decorated analysis function, adds the header data, metadata and error messages to the AnalysisResult, and then handles oncomplete.

@opticstudio_analysis("ZernikeStandardCoefficients")
def zernike_standard_coefficients(
    analysis: zospy.api._ZOSAPI.Analyses.IA_,
    sampling: str = "64x64",
    maximum_term: int = 37,
    wavelength: str | int = 1,
    field: str | int = 1,
    reference_opd_to_vertex: bool = False,
    surface: str | int = "Image",
    sx: float = 0.0,
    sy: float = 0.0,
    sr: float = 0.0,
) -> zospy.api._ZOSAPI.Analysis.Data.IAR_:
    # Set analysis settings, e.g.
    analysis.Settings.SampleSize = sampling

    # Run analysis and retrieve result
    ...

    return result

It should be noted that this decorator will change the signature of the analysis function, i.e. it will add the oss, txtoutfile, cfgoutfile and oncomplete parameters, and remove the analysis parameter, since this is injected by the decorator. I think this will break autocomplete. I know functools.wraps is often used to propagate docstrings and annotations to the decorated function, but this won't help in our case, due to the changed signature.
However, it will still be possible to provide correct autocompletions by adding .pyi files with the signatures and docstrings for the decorated functions.

I would like to hear your opinion on this! Is there anything that can be improved in this design, or do you know about different methods to preserve autocomplete without having to add .pyi files?

[crnh]

Here is a basic example of how I imagine the decorator:

import zospy as zp
import zospy.constants as constants
import zospy.api._ZOSAPI as _ZOSAPI

from zospy.analyses.base import new_analysis, AnalysisResult, OnComplete

def opticstudio_analysis(analysis_type: constants.Analysis.AnalysisIDM | str, txtoutfile: bool = False, cfgoutfile: bool = False):
    def analysis_decorator(f):
        def wrapper(oss: zp.zpcore.OpticStudioSystem, *args, **kwargs):
            analysis_type = constants.process_constant(constants.Analysis.AnalysisIDM, analysis_type)

            analysis = new_analysis(oss, analysis_type)

            txtoutfile_path = kwargs.pop("txtoutfile", None)
            cfgoutfile_path = kwargs.pop("txtoutfile", None)
            oncomplete = kwargs.pop("oncomplete", OnComplete.Close)

            if txtoutfile:
                # Handle txtoutfile
            if cfgoutfile:
                # Handle cfgoutfile

            result = f(analysis, *args, **kwargs)

            result.headerdata = analysis.get_headerdata()
            result.metadata = analysis.get_metadata()
            result.messages = analysis.get_messages()

            return analysis.complete(result, oncomplete)

        return wrapper

[andibarg]

I think this could be a good change and I like the implementation by @crnh.

Originally, my thinking was that there would be decorators e.g. for handling creation and cleanup of txtoutfile or cfgoutfile only. They could be added whenever these files are needed to the analysis function and thus keep the code of the analysis itself clean. I imagine using it something like this:

@textoutputfile
@cfgoutputfile
def polarization_pupil_map(
    oss: OpticStudioSystem,
    jx: float = 1,
    jy: float = 0,
    x_phase: float = 0,
    y_phase: float = 0,
    wavelength: int = 1,
    field: int = 1,
    surface: str | int = "Image",
    sampling: str | int = "11x11",
    add_configs: str = "",
    sub_configs: str = "",
    oncomplete: OnComplete | str = OnComplete.Close,
    cfgoutfile: str | None = None,
    txtoutfile: str | None = None,
) -> AnalysisResult:

    # Run analysis and retrieve result
    ...

    return analysis.complete(oncomplete, result)

I was not quite sure how to get the arguments from the analysis function correctly, even if they are 'None' but I think @crnh solved that :)

[crnh]

This feature is planned for the next major release, which will streamline the development of new analysis wrappers.

[crnh]

The implementation discussed above uses a decorator that changes the function signature. This is considered a code smell. Besides, the proposed use of .pyi stub files to solve problems with editor autosuggestions caused by the signature change can lead to an increased maintenance burden. The final implementation will therefore be slightly different from this proposal.

[crnh]

It's been a while, but we are currently working on something related to this; see PR #78 .