1# Copyright 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""This module provides standard functions for working with Autotest labels.
6
7There are two types of labels, plain ("webcam") or keyval
8("pool:suites").  Most of this module's functions work with keyval
9labels.
10
11Most users should use LabelsMapping, which provides a dict-like
12interface for working with keyval labels.
13
14This module also provides functions for working with cros version
15strings, which are common keyval label values.
16"""
17
18import collections
19import re
20
21
22class Key(object):
23    """Enum for keyval label keys."""
24    CROS_VERSION = 'cros-version'
25    CROS_ANDROID_VERSION = 'cheets-version'
26    FIRMWARE_RW_VERSION = 'fwrw-version'
27    FIRMWARE_RO_VERSION = 'fwro-version'
28    FIRMWARE_CR50_RW_VERSION = 'cr50-rw-version'
29
30
31class LabelsMapping(collections.MutableMapping):
32    """dict-like interface for working with labels.
33
34    The constructor takes an iterable of labels, either plain or keyval.
35    Plain labels are saved internally and ignored except for converting
36    back to string labels.  Keyval labels are exposed through a
37    dict-like interface (pop(), keys(), items(), etc. are all
38    supported).
39
40    When multiple keyval labels share the same key, the first one wins.
41
42    The one difference from a dict is that setting a key to None will
43    delete the corresponding keyval label, since it does not make sense
44    for a keyval label to have a None value.  Prefer using del or pop()
45    instead of setting a key to None.
46
47    LabelsMapping has one method getlabels() for converting back to
48    string labels.
49    """
50
51    def __init__(self, str_labels=()):
52        self._plain_labels = []
53        self._keyval_map = collections.OrderedDict()
54        for str_label in str_labels:
55            self._add_label(str_label)
56
57    def _add_label(self, str_label):
58        """Add a label string to the internal map or plain labels list."""
59        try:
60            keyval_label = parse_keyval_label(str_label)
61        except ValueError:
62            self._plain_labels.append(str_label)
63        else:
64            if keyval_label.key not in self._keyval_map:
65                self._keyval_map[keyval_label.key] = keyval_label.value
66
67    def __getitem__(self, key):
68        return self._keyval_map[key]
69
70    def __setitem__(self, key, val):
71        if val is None:
72            self.pop(key, None)
73        else:
74            self._keyval_map[key] = val
75
76    def __delitem__(self, key):
77        del self._keyval_map[key]
78
79    def __iter__(self):
80        return iter(self._keyval_map)
81
82    def __len__(self):
83        return len(self._keyval_map)
84
85    def getlabels(self):
86        """Return labels as a list of strings."""
87        str_labels = self._plain_labels[:]
88        keyval_labels = (KeyvalLabel(key, value)
89                         for key, value in self.iteritems())
90        str_labels.extend(format_keyval_label(label)
91                          for label in keyval_labels)
92        return str_labels
93
94
95_KEYVAL_LABEL_SEP = ':'
96
97
98KeyvalLabel = collections.namedtuple('KeyvalLabel', 'key, value')
99
100
101def parse_keyval_label(str_label):
102    """Parse a string as a KeyvalLabel.
103
104    If the argument is not a valid keyval label, ValueError is raised.
105    """
106    key, value = str_label.split(_KEYVAL_LABEL_SEP, 1)
107    return KeyvalLabel(key, value)
108
109
110def format_keyval_label(keyval_label):
111    """Format a KeyvalLabel as a string."""
112    return _KEYVAL_LABEL_SEP.join(keyval_label)
113
114
115CrosVersion = collections.namedtuple(
116        'CrosVersion', 'group, board, milestone, version, rc')
117
118
119_CROS_VERSION_REGEX = (
120        r'^'
121        r'(?P<group>[a-z0-9_-]+)'
122        r'/'
123        r'(?P<milestone>R[0-9]+)'
124        r'-'
125        r'(?P<version>[0-9.]+)'
126        r'(-(?P<rc>rc[0-9]+))?'
127        r'$'
128)
129
130_CROS_BOARD_FROM_VERSION_REGEX = (
131        r'^'
132        r'(trybot-)?'
133        r'(?P<board>[a-z_-]+)-(release|paladin|pre-cq|test-ap|toolchain)'
134        r'/R.*'
135        r'$'
136)
137
138
139def parse_cros_version(version_string):
140    """Parse a string as a CrosVersion.
141
142    If the argument is not a valid cros version, ValueError is raised.
143    Example cros version string: 'lumpy-release/R27-3773.0.0-rc1'
144    """
145    match = re.search(_CROS_VERSION_REGEX, version_string)
146    if match is None:
147        raise ValueError('Invalid cros version string: %r' % version_string)
148    parts = match.groupdict()
149    match = re.search(_CROS_BOARD_FROM_VERSION_REGEX, version_string)
150    if match is None:
151        raise ValueError('Invalid cros version string: %r. Failed to parse '
152                         'board.' % version_string)
153    parts['board'] = match.group('board')
154    return CrosVersion(**parts)
155
156
157def format_cros_version(cros_version):
158    """Format a CrosVersion as a string."""
159    if cros_version.rc is not None:
160        return '{group}/{milestone}-{version}-{rc}'.format(
161                **cros_version._asdict())
162    else:
163        return '{group}/{milestone}-{version}'.format(**cros_version._asdict())
164