Skip to content

Consider IPython's traitlets as an application framework #10

Open

Description

@eslavich

The traitlets package is used by IPython/Jupyter for application configuration. It handles type checking and other validation and mapping to CLI arguments and config files. It also comes with logger configuration support built-in. We may be able replace configobj and a bunch of custom code with traitlets, and usage would be familiar to users who have previously configured ipython, the notebook server, etc.

Here's an example of how traitlets work:

#!/usr/bin/env python3
from traitlets import Integer, Enum, Bool, validate, TraitError
from traitlets.config import Application, LoggingConfigurable


class Step(LoggingConfigurable):
    save_results = Bool(
        default_value=False,
        help="""
            Force save results for an intermediate step.
        """,
        config=True,
    )


class Extract1dStep(Step):
    smoothing_length = Integer(
        default_value=None,
        allow_none=True,
        help="""
            If set, the background regions (if any) will be smoothed
            with a boxcar function of this width along the dispersion
            direction.  This must be an odd integer.
        """,
        config=True,
    )

    @validate("smoothing_length")
    def _validate_odd(self, proposal):
        value = proposal["value"]
        trait = proposal["trait"]
        if value % 2 == 0:
            raise TraitError(f"{trait.name} must be an odd integer")
        return value

    bkg_fit = Enum(
        ["poly", "mean", "median"],
        default_value="poly",
        help="""
            A string indicating the type of fitting to be applied to
            background values in each column (or row, if the dispersion is
            vertical).
            """,
        config=True,
    )

    bkg_order = Integer(
        default_value=None,
        min=0,
        allow_none=True,
        help="""
            If present, a polynomial with order 'bkg_order' will be fit to
            each column (or row, if the dispersion direction is vertical)
            of the background region or regions.  For a given column (row),
            one polynomial will be fit to all background regions.  The
            polynomial will be evaluated at each pixel of the source
            extraction region(s) along the column (row), and the fitted value
            will be subtracted from the data value at that pixel.
            If both 'smoothing_length' and 'bkg_order' are not None, the
            boxcar smoothing will be done first.
        """,
        config=True,
    )

    log_increment = Integer(
        default_value=50,
        help="""
            If log_increment is greater than 0 and the input data are multi-integration
            (which can be CubeModel or SlitModel), a message will be written to the log
            with log level INFO every log_increment integrations.  This is intended to
            provide progress information when invoking the step interactively.
        """,
        config=True,
    )

    subtract_background = Bool(
        default_value=None,
        allow_none=True,
        help="""
            A flag which indicates whether the background should be subtracted.
            If absent, the value in the extract_1d reference file will be used.
            If present, this parameter overrides the value in the
            extract_1d reference file.
        """,
        config=True,
    )

    use_source_posn = Bool(
        default_value=None,
        allow_none=True,
        help="""
            If True, the source and background extraction positions specified in
            the extract1d reference file (or the default position, if there is no
            reference file) will be shifted to account for the computed position
            of the source in the data.  If absent, the values in the
            reference file will be used. Aperture offset is determined by computing
            the pixel location of the source based on its RA and Dec. It does not
            make sense to apply aperture offsets for extended sources, so this
            parameter can be overriden (set to False) internally by the step.
        """,
        config=True,
    )

    apply_apcorr = Bool(
        default_value=True,
        help="""
            Switch to select whether or not to apply an APERTURE correction during
            the Extract1dStep.
        """,
        config=True,
    )

    def __call__(self):
        self.log.debug(f"Called Extract1dStep.__call__")
        self.log.debug(f"smoothing_length={self.smoothing_length}")
        self.log.debug(f"save_results={self.save_results}")
        pass


class Stpipe(Application):
    name = "stpipe"

    description = "Calibration pipeline CLI"

    classes = [Extract1dStep]

    def start(self):
        self.log.debug("Called Stpipe.start")

        step = Extract1dStep(config=self.config)
        step()


def main():
    Stpipe.launch_instance()


if __name__ == "__main__":
    main()

Calling this script with no arguments:

$ ./stpipe
[Stpipe] smoothing_length=None
[Stpipe] save_results=False

Overriding save_results for all subclasses of Step:

$ ./stpipe --Step.save_results=true
[Stpipe] smoothing_length=None
[Stpipe] save_results=True

Overriding for just Extract1dStep:

$ ./stpipe --Extract1dStep.save_results=true
[Stpipe] smoothing_length=None
[Stpipe] save_results=True

Invalid smoothing_length caught by custom validator:

$ ./stpipe --Extract1dStep.smoothing_length=10
Traceback (most recent call last):
...
traitlets.traitlets.TraitError: smoothing_length must be an odd integer

Okay, fine, let's turn it up to 11:

$ ./stpipe --Extract1dStep.smoothing_length=11
[Stpipe] smoothing_length=11
[Stpipe] save_results=False

Built-in support for log configuration:

$ ./stpipe --Stpipe.log_format='%(asctime)s - %(levelname)s - %(message)s' --Stpipe.log_level=DEBUG
2021-01-21 22:55:22 - DEBUG - Called Stpipe.start
2021-01-21 22:55:22 - DEBUG - Called Extract1dStep.__call__
2021-01-21 22:55:22 - INFO - smoothing_length=None
2021-01-21 22:55:22 - INFO - save_results=False

Auto-generated help for Extract1dStep:

Extract1dStep(Step) options
---------------------------
--Extract1dStep.apply_apcorr=<Bool>
    Switch to select whether or not to apply an APERTURE correction during the
    Extract1dStep.
    Default: True
--Extract1dStep.bkg_fit=<Enum>
    A string indicating the type of fitting to be applied to background values
    in each column (or row, if the dispersion is vertical).
    Choices: any of ['poly', 'mean', 'median']
    Default: 'poly'
--Extract1dStep.bkg_order=<Int>
    If present, a polynomial with order 'bkg_order' will be fit to each column
    (or row, if the dispersion direction is vertical) of the background region
    or regions.  For a given column (row), one polynomial will be fit to all
    background regions.  The polynomial will be evaluated at each pixel of the
    source extraction region(s) along the column (row), and the fitted value
    will be subtracted from the data value at that pixel. If both
    'smoothing_length' and 'bkg_order' are not None, the boxcar smoothing will
    be done first.
    Default: None
--Extract1dStep.log_increment=<Int>
    If log_increment is greater than 0 and the input data are multi-integration
    (which can be CubeModel or SlitModel), a message will be written to the log
    with log level INFO every log_increment integrations.  This is intended to
    provide progress information when invoking the step interactively.
    Default: 50
--Extract1dStep.save_results=<Bool>
    Force save results for an intermediate step.
    Default: False
--Extract1dStep.smoothing_length=<Int>
    If set, the background regions (if any) will be smoothed with a boxcar
    function of this width along the dispersion direction.  This must be an odd
    integer.
    Default: None
--Extract1dStep.subtract_background=<Bool>
    A flag which indicates whether the background should be subtracted. If
    absent, the value in the extract_1d reference file will be used. If present,
    this parameter overrides the value in the extract_1d reference file.
    Default: None
--Extract1dStep.use_source_posn=<Bool>
    If True, the source and background extraction positions specified in the
    extract1d reference file (or the default position, if there is no reference
    file) will be shifted to account for the computed position of the source in
    the data.  If absent, the values in the reference file will be used.
    Aperture offset is determined by computing the pixel location of the source
    based on its RA and Dec. It does not make sense to apply aperture offsets
    for extended sources, so this parameter can be overriden (set to False)
    internally by the step.
    Default: None

Config file example (we wouldn't have to use this, but it's available if we need it):

c.Extract1dStep.save_results = True
c.Extract1dStep.smoothing_length = 11

Activity

  1. perrygreenfield commented on Jan 22, 2021

    @perrygreenfield

    Yes, I'd say we should take a look at it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions