fix buffer

buffer.py

@@ -13,10 +13,11 @@ class FileBuffer:
         f = open(filename, "rb")
         segments = []
         while 1:
-            data = f.read(DATA_SIZE)
+            data = f.read(DATA_SIZE-1)
             if not data:
                 break
-            segments.append(Segment(data=data))
+            seg = Segment(data=data)
+            segments.append(seg)
         return segments
 
     def write(self, segments:List[Segment], filename):
@@ -24,5 +25,17 @@ class FileBuffer:
         # write array of segment to file
         with open(filename, 'wb') as f:
             for seg in segments:
-                f.write(seg.get_data)
+                data = seg.get_data()
+                f.write(data)
         return
+
+if __name__=="__main__":
+    fb = FileBuffer()
+    segments = fb.read("data/test")
+
+    print(len(segments))
+
+    fb.write(segments, "data/out")
+
+    buff = [10, 90, 99]
+    byte = bytes(buff)

segment.py

@@ -3,14 +3,15 @@ import os
 
 # constant
 SEG_SIZE = 32782
-DATA_SIZE = SEG_SIZE - 12
+HEADER_SIZE = 12
+DATA_SIZE = SEG_SIZE - HEADER_SIZE
 
 class Segment():
 
     def __init__(self, seq_no=0, ack_no=0, data=bytes(), checksum=0):
         self.DATA_PTR = 12
 
-        self.bytes = [0 for _ in range(SEG_SIZE)]
+        self.bytes = [0 for _ in range(HEADER_SIZE)]
         self.set_seq_no(seq_no)
         self.set_ack_no(ack_no)
         self.set_data(data)
@@ -42,14 +43,17 @@ class Segment():
 
     def set_data(self, data):
         ptr = self.DATA_PTR
-        for byte in data:
-            self.bytes[ptr] = byte
-            ptr += 1
+        bound = len(data)
+        if bound > DATA_SIZE:
+            bound = DATA_SIZE
+        for i in range(bound):
+            self.bytes.append(data[i])
+            
 
     def get_data(self):
         data=[]
         ptr = self.DATA_PTR   # pointer to data in segment
-        while self.bytes[ptr] != 0:
+        for b in self.bytes[HEADER_SIZE:]:
             data.append(self.bytes[ptr])
             ptr += 1
         return bytes(data)