TOOLS/matroska.py: allow using as module

Don't force CLI usage. It can be imported, and generate_C_header() and
generate_C_definitions() can be called with a file argument instead of
writing to stdout always.

1 files changed, 49 insertions, 46 deletions

diff --git a/TOOLS/matroska.py b/TOOLS/matroska.py
index 91e65a26b3..0ff0d46862 100755
--- a/TOOLS/matroska.py
+++ b/TOOLS/matroska.py
@@ -284,56 +284,58 @@ def parse_elems(l, namespace):
 parse_elems(elements_ebml, 'EBML')
 parse_elems(elements_matroska, 'MATROSKA')
 
-def generate_C_header():
-    print('// Generated by TOOLS/matroska.py, do not edit manually')
-    print()
+def printf(out, *args):
+    out.write(' '.join([str(x) for x in args]))
+    out.write('\n')
+
+def generate_C_header(out):
+    printf(out, '// Generated by TOOLS/matroska.py, do not edit manually')
+    printf(out)
 
     for el in elementlist:
-        print('#define {0.definename:40} 0x{0.elid}'.format(el))
+        printf(out, '#define {0.definename:40} 0x{0.elid}'.format(el))
 
-    print()
+    printf(out)
 
     for el in reversed(elementlist):
         if not el.subelements:
             continue
-        print()
-        print('struct {0.structname} {{'.format(el))
+        printf(out)
+        printf(out, 'struct {0.structname} {{'.format(el))
         l = max(len(subel.valname) for subel, multiple in el.subelements)+1
         for subel, multiple in el.subelements:
-            print('    {e.valname:{l}} {star}{e.fieldname};'.format(
-                    e=subel, l=l, star=' *'[multiple]))
-        print()
+            printf(out, '    {e.valname:{l}} {star}{e.fieldname};'.format(
+                        e=subel, l=l, star=' *'[multiple]))
+        printf(out)
         for subel, multiple in el.subelements:
-            print('    int  n_{0.fieldname};'.format(subel))
-        print('};')
+            printf(out, '    int  n_{0.fieldname};'.format(subel))
+        printf(out, '};')
 
     for el in elementlist:
         if not el.subelements:
             continue
-        print('extern const struct ebml_elem_desc {0.structname}_desc;'.format(
-                el))
-
-    print()
-    print('#define MAX_EBML_SUBELEMENTS', max(len(el.subelements)
-                                              for el in elementlist))
+        printf(out, 'extern const struct ebml_elem_desc {0.structname}_desc;'.format(el))
 
+    printf(out)
+    printf(out, '#define MAX_EBML_SUBELEMENTS', max(len(el.subelements)
+                                                    for el in elementlist))
 
 
-def generate_C_definitions():
-    print('// Generated by TOOLS/matroska.py, do not edit manually')
-    print()
+def generate_C_definitions(out):
+    printf(out, '// Generated by TOOLS/matroska.py, do not edit manually')
+    printf(out)
     for el in reversed(elementlist):
-        print()
+        printf(out)
         if el.subelements:
-            print('#define N', el.fieldname)
-            print('E_S("{0}", {1})'.format(el.name, len(el.subelements)))
+            printf(out, '#define N', el.fieldname)
+            printf(out, 'E_S("{0}", {1})'.format(el.name, len(el.subelements)))
             for subel, multiple in el.subelements:
-                print('F({0.definename}, {0.fieldname}, {1})'.format(
-                        subel, int(multiple)))
-            print('}};')
-            print('#undef N')
+                printf(out, 'F({0.definename}, {0.fieldname}, {1})'.format(
+                            subel, int(multiple)))
+            printf(out, '}};')
+            printf(out, '#undef N')
         else:
-            print('E("{0.name}", {0.fieldname}, {0.ebmltype})'.format(el))
+            printf(out, 'E("{0.name}", {0.fieldname}, {0.ebmltype})'.format(el))
 
 def read(s, length):
     t = s.read(length)
@@ -444,20 +446,21 @@ def parse_one(s, depth, parent, maxlen):
         read(s, length)
     return this_length
 
-def parse_toplevel(s):
-    parse_one(s, 0, None, 1 << 63)
-
-if sys.argv[1] == '--generate-header':
-    generate_C_header()
-elif sys.argv[1] == '--generate-definitions':
-    generate_C_definitions()
-else:
-    s = open(sys.argv[1], "rb")
-    while 1:
-        start = s.tell()
-        try:
-            parse_toplevel(s)
-        except EOF:
-            if s.tell() != start:
-                raise Exception("Unexpected end of file")
-            break
+if __name__ == "__main__":
+    def parse_toplevel(s):
+        parse_one(s, 0, None, 1 << 63)
+
+    if sys.argv[1] == '--generate-header':
+        generate_C_header(sys.stdout)
+    elif sys.argv[1] == '--generate-definitions':
+        generate_C_definitions(sys.stdout)
+    else:
+        s = open(sys.argv[1], "rb")
+        while 1:
+            start = s.tell()
+            try:
+                parse_toplevel(s)
+            except EOF:
+                if s.tell() != start:
+                    raise Exception("Unexpected end of file")
+                break