Ass 3 q2 basically done

1 files changed, 21 insertions, 11 deletions

diff --git a/ass3/q2/header.py b/ass3/q2/header.py
index 8d98fb4..0db542f 100644
--- a/ass3/q2/header.py
+++ b/ass3/q2/header.py
@@ -1,4 +1,5 @@
 import heapq
+import sys
 
 
 class HuffNode:
@@ -79,28 +80,37 @@ def elias(n):
     binary = bin(n)[2:]
     result = binary
     while len(binary) != 1:
-        binary = "0" + bin(len(binary)-1)[3:]
+        binary = "0" + bin(len(binary) - 1)[3:]
         result = binary + result
     return result
 
+
 def header(string):
     huff_code = huffman(string)
     alphabet_size = len(huff_code.keys())
     output = f"{elias(alphabet_size)}"
-    output2 = f"{alphabet_size} -- "
     for char in huff_code:
         output += bin(ord(char))[2:]
         output += elias(len(huff_code[char]))
         output += huff_code[char]
+    return output
+
+
+def read_file(filename):
+    with open(filename, "r") as file:
+        contents = file.read()
+    return contents
+
+
+def write_file(filename, contents):
+    with open(filename, "w") as file:
+        file.write(contents)
 
-        output2 += "("
-        output2 += str(ord(char))
-        output2 += ": "
-        output2 += str(len(huff_code[char]))
-        output2 += " -> '"
-        output2 += huff_code[char]
-        output2 += "'), "
-    return output, output2
 
+def main():
+    assert len(sys.argv) >= 2
+    string = read_file(sys.argv[1])
+    write_file("output_header.txt", header(string)[0])
 
-print(header("aacaacabcaba"))
+if __name__ == "__main__":
+    main()