From f4df9adf3ee2cb7e3b9848cfbd7078bd1532f9f3 Mon Sep 17 00:00:00 2001
From: siyansusan <siyan.liu435@duke.edu>
Date: Tue, 1 Nov 2022 00:23:22 -0400
Subject: [PATCH] fix matrix market code

---
 MatrixMarket.py | 106 +++++++++++++++++++++++++++++-------------------
 1 file changed, 64 insertions(+), 42 deletions(-)

diff --git a/MatrixMarket.py b/MatrixMarket.py
index d5bf8ea..56c901a 100755
--- a/MatrixMarket.py
+++ b/MatrixMarket.py
@@ -3,6 +3,7 @@
 # This is OPEN SOURCE SOFTWARE governed by the Gnu General Public
 # License (GPL) version 3, as described at www.opensource.org.
 # Author: William H. Majoros (bmajoros@alumni.duke.edu)
+# Author(10/31/2022): Susan Liu
 #=========================================================================
 from __future__ import (absolute_import, division, print_function, 
    unicode_literals, generators, nested_scopes, with_statement)
@@ -18,62 +19,83 @@
 
 #=========================================================================
 # Attributes:
-#    FH : file handle
+#    filename: string
 #    header : array of int
-#    nextLine : string
+#    col_index: int
+#    groups_generator: generator of array
 # Instance Methods:
 #    MatrixMarket(filename)
+#    get_groups(self)
 #    nextGroup(self,colIndex)
 #    getHeader()
 # Class Methods:
 #    allGroups=loadFile(filename,colIndex)
 #=========================================================================
 class MatrixMarket:
-    def __init__(self,filename):
-        if(rex.find("\.gz$",filename)):
-            self.FH=gzip.open(filename,"rt")
+    def __init__(self, filename):
+        self.filename = filename
+
+        self.header = None
+
+        self.col_index = None
+        self.groups_generator = None
+
+    def get_groups(self):
+
+        #find file opening function
+        if self.filename.endswith(".gz"):
+            open_func = gzip.open
         else:
-            self.FH=open(filename,"rt")
-        self.header=None
-        self.nextLine=None
+            open_func = open
+
+        with open_func(self.filename, "rt") as fh:
+
+            #skip first line of header
+            fh.readline()
+
+            #read totals
+            line = fh.readline()
+            self.header = [int(x) for x in line.rstrip().split()]
+
+            #initialize current ID and group
+            cur_ID = None
+            group = []
+
+            #continue reading groups
+            for line in fh:
+
+                fields = line.rstrip().split()
+                
+                if cur_ID is None:
+                    cur_ID = fields[self.col_index]
+
+                if cur_ID != fields[self.col_index]:
+                    yield group
+
+                    cur_ID = fields[self.col_index]
+                    group = [fields]
+
+                else:
+                    group.append(fields)
+
+            #yield the last group
+            yield group
+
+    def nextGroup(self, colIndex):
+
+        #initialize group generator
+        if self.groups_generator is None:
+            self.col_index = colIndex
+            self.groups_generator = self.get_groups()
+        
+        try:
+            return next(self.groups_generator)
+        except StopIteration:
+            return None
 
     def getHeader(self):
         return self.header
 
-    def nextGroup(self,colIndex):
-        line=None
-        # First, see if the header needs to be parsed
-        while(True):
-            line=self.nextLine
-            if(line is None): line=self.FH.readline()
-            if(line is None): return None
-            L=len(line)
-            if(L>0 and line[0]=="%"): continue
-            break
-        # The first non-comment line contains the totals
-        if(self.header is None):
-            self.header=line.rstrip().split()
-            self.header=[int(x) for x in self.header]
-            line=self.FH.readline()
-        # Now we can read in the next group of lines
-        prevID=None
-        group=[]
-        if(self.nextLine is not None): # buffered from previous call
-            prevID=int(self.nextLine.rstrip().split()[colIndex])
-        while(True):
-            fields=line.rstrip().split()
-            if(len(fields)==0): return None
-            if(prevID is None): prevID=int(fields[colIndex])
-            if(colIndex>len(fields)-1):
-                raise Exception("colIndex=",colIndex,"len(fields)=",len(fields))
-            thisID=int(fields[colIndex])
-            if(thisID!=prevID):
-                self.nextLine=line
-                return group
-            group.append(fields)
-            line=self.FH.readline()
-            if(line is None): return None
-
     @classmethod
     def loadFile(self,filename,colIndex):
         reader=MatrixMarket(filename)