1
    2
    3
    4
    5
    6
    7
    8
    9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20
   21
   22
   23
   24
   25
   26
   27
   28
   29
   30
   31
   32
   33
   34
   35
   36
   37
   38
   39
   40
   41
   42
   43
   44
   45
   46
   47
   48
   49
   50
   51
   52
   53
   54
   55
   56
   57
   58
   59
   60
   61
   62
   63
   64
   65
   66
   67
   68
   69
   70
   71
   72
   73
   74
   75
   76
   77
   78
   79
   80
   81
   82
   83
   84
   85
   86
   87
   88
   89
   90
   91
   92
   93
   94
   95
   96
   97
   98
   99
  100
  101
  102
  103
  104
  105
  106
  107
  108
  109
  110
  111
  112
  113
  114
  115
  116
  117
  118
  119
  120
  121
  122
  123
  124
  125
  126
  127
  128
  129
  130
  131
  132
  133
  134
  135
  136
  137
  138
  139
  140
  141
  142
  143
  144
  145
  146
  147
  148
  149
  150
  151
  152
  153
  154
  155
  156
  157
  158
  159
  160
  161
  162
  163
  164
  165
  166

infra / config / lib / structs.star [blame]

# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Library for working with starlark structs.

The functionality for this module can be used by loading the structs symbol,
which provides the following functions:
* to_properties - convert a struct to a dict that can represent a proto in
    properties
* evolve - set new values for struct attributes
* extend - extend values for struct attributes
* remove - remove values for struct attributes
"""

load("./args.star", "args")

def _convert_to_dict(s):
    """Convert a struct to a dict

    Due to starlark not supporting recursion, in order to enable a struct
    containing nested structs to be recursively converted, the conversion won't
    necessarily be complete so a partially converted dict will be returned along
    with information of the attributes that need to themselves be converted.

    Args:
        s: The struct to convert.

    Returns:
        A 2-tuple:
            * A dict of converted attributes.
            * A sequence of 3-tuples for each attribute that needs to be recursively
            converted:
                * The dictionary to add the converted value to.
                * The name of the key to set with the converted value.
                * The value to convert.
    """
    d = {}
    to_convert = []
    for a in dir(s):
        v = getattr(s, a)
        if v == None or v == []:
            continue
        if type(v) == type(struct()):
            to_convert.append((d, a, v))
            continue
        d[a] = v
    return d, to_convert

def _to_proto_properties(s):
    """Converts a struct to a properties dict that can represent a proto.

    Args:
        s: The struct to convert.

    Returns:
        A dict with items corresponding to the attributes of `s`. Attributes
        with a None value or empty list will be omitted from the dict since the
        corresponding jsonpb values are equivalent to the field not being set.
    """
    top_level_dict, to_convert = _convert_to_dict(s)

    # Since starlark doesn't support recursion, iterate over the depth of the
    # structure, recording the work to be done at the next depth. Since starlark
    # doesn't support while loops, iterate up to an aribtrary maximum depth and
    # break out of the loop if there's no more work to be done.
    for _ in range(15):
        if not to_convert:
            break
        next_to_convert = []
        for d, k, s in to_convert:
            converted_s, to_convert_for_s = _convert_to_dict(s)
            next_to_convert.extend(to_convert_for_s)
            d[k] = converted_s
        to_convert = next_to_convert

    if to_convert:
        fail("excessively nested struct")

    return top_level_dict

def _evolve(s, **kwargs):
    """Modify a struct's attributes.

    Args:
        s: The struct to modify.
        **kwargs: The attributes to update the struct with.

    Returns:
        A new struct with the value of each attribute specified in `kwargs`
        is set to the corresponding value.

    Fails:
        If `kwargs` contains an item for an attribute not present on `s`.
    """
    d = {a: getattr(s, a) for a in dir(s)}
    for k, v in kwargs.items():
        if k not in d:
            fail("attempting to modify unknown field {!r}".format(k))
        d[k] = v
    return struct(**d)

def _extend(s, **kwargs):
    """Extend a struct's list attributes.

    Args:
        s: The struct to modify.
        **kwargs: The attributes to update the struct with. The value of each
            element can either be a single value or a list of values.

    Returns:
        A new struct where the value of each attribute specified in `kwargs`
        is set to the combination of the existing list of values if any and the
        values specified for the attribute in `kwargs`.

    Fails:
        If `kwargs` contains an item for an attribute not present on `s`.
        If any of the attributes to modify do not have list values.
    """
    d = {a: getattr(s, a) for a in dir(s)}
    for k, v in kwargs.items():
        if k not in d:
            fail("attempting to modify unknown field {!r}".format(k))
        value = d[k]
        if type(value) != type([]):
            fail("attempting to extend a non-list field {!r}".format(k))
        d[k] = args.listify(value, v)
    return struct(**d)

def _remove(s, **kwargs):
    """Remove elements from a struct's list attributes.

    Args:
        s: The struct to modify.
        **kwargs: The attributes to update the struct with. The value of each
            element can either be a single value or a list of values.

    Returns:
        A new struct where the value of each attribute specified in `kwargs`
        has the given values removed.

    Fails:
        If `kwargs` contains an item for an attribute not present on `s`.
        If any of the attributes to modify do not have list values.
        If any of the elements to remove do not exist in the specified
            attribute.
    """
    d = {a: getattr(s, a) for a in dir(s)}
    for k, v in kwargs.items():
        if k not in d:
            fail("attempting to modify unknown field {!r}".format(k))
        value = d[k]
        if type(value) != type([]):
            fail("attempting to remove elements from a non-list field {!r}"
                .format(k))
        for e in args.listify(v):
            value.remove(e)
        d[k] = value
    return struct(**d)

structs = struct(
    to_proto_properties = _to_proto_properties,
    evolve = _evolve,
    extend = _extend,
    remove = _remove,
)