1#!/usr/bin/env python3
2#
3#   Copyright 2018 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import os
18
19from acts.libs.proto.proto_utils import parse_proto_to_ascii
20from acts.libs.proto.proto_utils import to_descriptor_proto
21from acts.utils import dump_string_to_file
22
23
24class ProtoMetric(object):
25    """A wrapper around a protobuf containing metrics data.
26
27    This is the primary data structure used as the output of MetricLoggers. It
28    is generally intended to be used as-is as a simple wrapper structure, but
29    can be extended to provide self-populating APIs.
30
31    Attributes:
32        name: The name of the metric.
33        data: The data of the metric.
34    """
35
36    def __init__(self, name=None, data=None):
37        """Initializes a metric with given name and data.
38
39        Args:
40            name: The name of the metric. Used in identifiers such as filename.
41            data: The data of the metric. Should be a protobuf object.
42        """
43        self.name = name
44        if not data:
45            raise ValueError("Parameter 'data' cannot be None.")
46        self.data = data
47
48    def get_binary(self):
49        """Gets the binary representation of the protobuf data."""
50        return self.data.SerializeToString()
51
52    def get_ascii(self):
53        """Gets the ascii representation of the protobuf data."""
54        return parse_proto_to_ascii(self.data)
55
56    def get_descriptor_binary(self):
57        """Gets the binary representation of the descriptor protobuf."""
58        return to_descriptor_proto(self.data).SerializeToString()
59
60    def get_descriptor_ascii(self):
61        """Gets the ascii representation of the descriptor protobuf."""
62        return parse_proto_to_ascii(to_descriptor_proto(self.data))
63
64
65class MetricPublisher(object):
66    """A helper object for publishing metric data.
67
68    This is a base class intended to be implemented to accommodate specific
69    metric types and output formats.
70
71    Attributes:
72        context: The context in which the metrics are being published.
73    """
74
75    def __init__(self, context):
76        """Initializes a publisher for the give context.
77
78        Args:
79            context: The context in which the metrics are being published.
80                     Typically matches that of a containing MetricLogger.
81        """
82        if not context:
83            raise ValueError("Parameter 'context' cannot be None.")
84        self.context = context
85
86    def publish(self, metrics):
87        """Publishes a list of metrics.
88
89        Args:
90            metrics: A list of metrics to publish. The requirements on the
91            object type of these metrics is up to the implementing class.
92        """
93        raise NotImplementedError()
94
95
96class ProtoMetricPublisher(MetricPublisher):
97    """A MetricPublisher that will publish ProtoMetrics to files.
98
99    Attributes:
100        publishes_binary: Whether to publish the binary proto.
101        publishes_ascii: Whether to publish the ascii proto.
102        publishes_descriptor_binary: Whether to publish the binary descriptor.
103        publishes_descriptor_ascii: Whether to publish the ascii descriptor.
104    """
105
106    ASCII_EXTENSION = 'proto.data'
107    BINARY_EXTENSION = 'proto.bin'
108    ASCII_DESCRIPTOR_EXTENSION = 'proto.desc'
109    BINARY_DESCRIPTOR_EXTENSION = 'proto.desc.bin'
110
111    METRICS_DIR = 'metrics'
112
113    def __init__(self,
114                 context,
115                 publishes_binary=True,
116                 publishes_ascii=True,
117                 publishes_descriptor_binary=True,
118                 publishes_descriptor_ascii=True):
119        """Initializes a ProtoMetricPublisher.
120
121        Args:
122            context: The context in which the metrics are being published.
123            publishes_binary: Whether to publish the binary proto.
124            publishes_ascii: Whether to publish the ascii proto.
125            publishes_descriptor_binary: Whether to publish the binary
126                                         descriptor.
127            publishes_descriptor_ascii: Whether to publish the ascii
128                                        descriptor.
129        """
130        super().__init__(context)
131        self.publishes_binary = publishes_binary
132        self.publishes_ascii = publishes_ascii
133        self.publishes_descriptor_binary = publishes_descriptor_binary
134        self.publishes_descriptor_ascii = publishes_descriptor_ascii
135
136    def get_output_path(self):
137        """Gets the output directory path of the metrics."""
138        return os.path.join(self.context.get_full_output_path(),
139                            self.METRICS_DIR)
140
141    def publish(self, metrics):
142        """Publishes the given list of metrics.
143
144        Based on the publish_* attributes of this class, this will publish
145        the varying data formats provided by the metric object. Data is written
146        to files on disk named according to the name of the metric.
147
148        Args:
149            metrics: The list metric to publish. Assumed to be a list of
150                     ProtoMetric objects.
151        """
152        if isinstance(metrics, list):
153            for metric in metrics:
154                self._publish_single(metric)
155        else:
156            self._publish_single(metrics)
157
158    def _publish_single(self, metric):
159        """Publishes a single metric.
160
161        Based on the publish_* attributes of this class, this will publish
162        the varying data formats provided by the metric object. Data is written
163        to files on disk named according to the name of the metric.
164
165        Args:
166            metric: The metric to publish. Assumed to be a ProtoMetric object.
167        """
168        output_path = self.get_output_path()
169
170        os.makedirs(output_path, exist_ok=True)
171
172        if self.publishes_binary:
173            self.write_binary(metric, output_path)
174        if self.publishes_ascii:
175            self.write_ascii(metric, output_path)
176        if self.publishes_descriptor_binary:
177            self.write_descriptor_binary(metric, output_path)
178        if self.publishes_descriptor_ascii:
179            self.write_descriptor_ascii(metric, output_path)
180
181    def write_binary(self, metric, output_path):
182        """Writes the binary format of the protobuf to file.
183
184        Args:
185            metric: The metric to write.
186            output_path: The output directory path to write the file to.
187        """
188        filename = self._get_output_file(
189            output_path, metric.name, self.BINARY_EXTENSION)
190        dump_string_to_file(metric.get_binary(), filename, mode='wb')
191
192    def write_ascii(self, metric, output_path):
193        """Writes the ascii format of the protobuf to file.
194
195        Args:
196            metric: The metric to write.
197            output_path: The output directory path to write the file to.
198        """
199        filename = self._get_output_file(
200            output_path, metric.name, self.ASCII_EXTENSION)
201        dump_string_to_file(metric.get_ascii(), filename)
202
203    def write_descriptor_binary(self, metric, output_path):
204        """Writes the binary format of the protobuf descriptor to file.
205
206        Args:
207            metric: The metric to write.
208            output_path: The output directory path to write the file to.
209        """
210        filename = self._get_output_file(
211            output_path, metric.name, self.BINARY_DESCRIPTOR_EXTENSION)
212        dump_string_to_file(metric.get_descriptor_binary(), filename, mode='wb')
213
214    def write_descriptor_ascii(self, metric, output_path):
215        """Writes the ascii format of the protobuf descriptor to file.
216
217        Args:
218            metric: The metric to write.
219            output_path: The output directory path to write the file to.
220        """
221        filename = self._get_output_file(
222            output_path, metric.name, self.ASCII_DESCRIPTOR_EXTENSION)
223        dump_string_to_file(metric.get_descriptor_ascii(), filename)
224
225    def _get_output_file(self, output_path, filename, extension):
226        """Gets the full output file path.
227
228        Args:
229            output_path: The output directory path.
230            filename: The base filename of the file.
231            extension: The extension of the file, without the leading '.'
232
233        Returns:
234            The full file path.
235        """
236        return os.path.join(output_path, "%s.%s" % (filename, extension))
237