1 files changed, 67 insertions, 0 deletions
diff --git a/src/mlia/backend/executor/output_consumer.py b/src/mlia/backend/executor/output_consumer.py
new file mode 100644
index 0000000..3c3b132
--- /dev/null
+++ b/src/mlia/backend/executor/output_consumer.py
@@ -0,0 +1,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