Standardized usage for subcapset.pyHEADmaster

1 files changed, 10 insertions, 7 deletions

diff --git a/subcapset.py b/subcapset.py
index fe8868f..375dc1b 100755
--- a/subcapset.py
+++ b/subcapset.py
@@ -3,27 +3,30 @@
 
 [00 01 10] is a subcapset of [00 01 10 22], and thus would be filtered."""
 
-from sys import argv
+from sys import argv, stdout
 
 def main():
-    if len(argv) == 2:
-        infile = argv[1]
-    else:
-        print("Usage: %s INFILE" % argv[0])
-        exit(1)
+    if len(argv) <= 1: infile = "capset.out"
+    if len(argv) >  1: infile = argv[1]
+    if len(argv) <= 2: outfile = "capset.filtered"
+    if len(argv) >  2: outfile = argv[2]
+
     capsets = []
     capset_file = open(infile, "r")
     for line in capset_file:
         capsets.append(tuple(line[1:-2].split(" ")))
     capset_file.close()
 
+    if outfile:
+        stdout = open(outfile, "w")
+
     capsetSet = set(capsets)
     for capset in capsets:
         for i in range(len(capset)):
             capsetSet.discard(capset[:i] + capset[i+1:])
     for capset in capsets:
         if capset in capsetSet:
-            print("[" + " ".join(capset) + "]")
+            print("[" + " ".join(capset) + "]", file=stdout)
 
 if __name__ == "__main__":
     main()