aboutsummaryrefslogtreecommitdiff
path: root/src/mlia/backend/output_consumer.py
blob: 3c3b1329d3d43b12c2308c2c27623b50f1385a74 (plain)
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
# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Output consumers module."""
from __future__ import annotations

import base64
import json
import re
from typing import Protocol
from typing import runtime_checkable


@runtime_checkable
class OutputConsumer(Protocol):
    """Protocol to consume output."""

    def feed(self, line: str) -> bool:
        """
        Feed a new line to be parsed.

        Return True if the line should be removed from the output.
        """


class Base64OutputConsumer(OutputConsumer):
    """
    Parser to extract base64-encoded JSON from tagged standard output.

    Example of the tagged output:
    ```
        # Encoded JSON: {"test": 1234}
        <metrics>eyJ0ZXN0IjogMTIzNH0</metrics>
    ```
    """

    TAG_NAME = "metrics"

    def __init__(self) -> None:
        """Set up the regular expression to extract tagged strings."""
        self._regex = re.compile(rf"<{self.TAG_NAME}>(.*)</{self.TAG_NAME}>")
        self.parsed_output: list = []

    def feed(self, line: str) -> bool:
        """
        Parse the output line and save the decoded output.

        Returns True if the line contains tagged output.

        Example:
        Using the tagged output from the class docs the parser should collect
        the following:
        ```
            [
                {"test": 1234}
            ]
        ```
        """
        res_b64 = self._regex.search(line)
        if res_b64:
            res_json = base64.b64decode(res_b64.group(1), validate=True)
            res = json.loads(res_json)
            self.parsed_output.append(res)
            # Remove this line from the output, i.e. consume it, as it
            # does not contain any human readable content.
            return True

        return False