1 files changed, 97 insertions, 0 deletions
diff --git a/tools/tosa.py b/tools/tosa.py
new file mode 100644
index 0000000..87b4f1a
--- /dev/null
+++ b/tools/tosa.py
@@ -0,0 +1,97 @@
+import re
+import xml.etree.ElementTree as ET
+
+
+class TOSAOperatorArgumentCategory:
+    def __init__(self, name, profiles=None):
+        self.name = name
+        self.profiles = profiles
+
+
+class TOSAOperatorArgument:
+    def __init__(self, name, description, categories, ty, shape):
+        self.name = name
+        self.description = description
+        self.categories = categories
+        self.type = ty
+        self.shape = shape
+
+
+class TOSAOperatorDataTypeSupport:
+    def __init__(self, mode, tymap, profiles=None):
+        self.mode = mode
+        self.tymap = tymap
+        self.profiles = profiles
+
+
+class TOSAOperator:
+    def __init__(self, name, arguments, types, typesupports):
+        self.name = name
+        self.arguments = arguments
+        self.types = types
+        self.typesupports = typesupports
+
+
+class TOSAOperatorGroup:
+    def __init__(self, name, operators):
+        self.name = name
+        self.operators = operators
+
+
+class TOSASpec:
+    def __init__(self, xmlpath):
+        tree = ET.parse(xmlpath)
+        self.xmlroot = tree.getroot()
+        self.operatorgroups = []
+        self.__load_spec()
+
+    def __load_spec(self):
+        for group in self.xmlroot.findall("./operators/operatorgroup"):
+            self.operatorgroups.append(self.__load_operator_group(group))
+
+    def __load_operator_group(self, group):
+        name = group.get("name")
+        operators = []
+        for op in group.findall("operator"):
+            operators.append(self.__load_operator(op))
+        return TOSAOperatorGroup(name, operators)
+
+    def __load_operator(self, op):
+        name = op.find("name").text
+        args = []
+        types = []
+        typesupports = []
+        for arg in op.findall("arguments/argument"):
+            args.append(self.__load_operator_argument(arg))
+
+        # TODO add pseudo-code to operator object?
+
+        for ty in op.findall("types/type"):
+            types.append(ty.get("name"))
+
+        for tysup in op.findall("typesupport"):
+            tsmode = tysup.get("mode")
+            tsmap = {}
+            profiles = tysup.findall("profile")
+            tsprofiles = []
+            for p in profiles:
+                tsprofiles.append(p.get("name"))
+            for ty in types:
+                tsmap[ty] = tysup.get(ty)
+            typesupports.append(TOSAOperatorDataTypeSupport(tsmode, tsmap, tsprofiles))
+        return TOSAOperator(name, args, types, typesupports)
+
+    def __load_operator_argument(self, arg):
+        name = arg.get("name")
+        desc = arg.find("description").text.strip()
+        argcats = []
+        argtype = arg.get("type")
+        shape = arg.get("shape")
+
+        cats = re.findall(
+            r"(input|output|attribute)\(?([A-Z,]+)?\)?", arg.get("category")
+        )
+        for cat in cats:
+            argcats.append(TOSAOperatorArgumentCategory(cat[0], cat[1].split(",")))
+
+        return TOSAOperatorArgument(name, desc, argcats, argtype, shape)