add sigsegments field to MultiRecord objects

cx1111 · cx1111 · commit 6ae258348ab5 · 2017-06-13T11:08:05.000-04:00
diff --git a/wfdb/readwrite/_headers.py b/wfdb/readwrite/_headers.py
@@ -362,9 +362,9 @@ def wrheaderfile(self, writefields):
         linestofile(self.recordname+'.hea', headerlines)
 
     # Get a list of the segment numbers that contain a particular signal
-    # (or a list of segment numbers for a list of signals)
+    # (or a dictionary of segment numbers for a list of signals)
     # Only works if information about the segments has been read in
-    def signalsegments(self, signame=None):
+    def getsigsegments(self, signame=None):
         if self.segments is None:
             raise Exception("The MultiRecord's segments must be read in before this method is called. ie. Call rdheader() with rdsegments=True")
         
@@ -373,7 +373,10 @@ def signalsegments(self, signame=None):
             signame = self.getsignames()
 
         if type(signame) == list:
-            return [self.signalsegments(sig) for sig in signame]
+            sigdict = {}
+            for sig in signame:
+                sigdict[sig] = self.getsigsegments(sig)
+            return sigdict
         elif type(signame) == str:
             sigsegs = []
             for i in range(self.nseg):
diff --git a/wfdb/readwrite/records.py b/wfdb/readwrite/records.py
@@ -475,12 +475,12 @@ class MultiRecord(BaseRecord, _headers.MultiHeadersMixin):
     of the record as a Record object. The resulting Record object will have its 'p_signals' field set.
 
     Contructor function:
-    def __init__(self, segments = None, layout = None,
+    def __init__(self, segments=None, layout=None,
                  recordname=None, nsig=None, fs=None,
                  counterfreq=None, basecounter=None,
                  siglen=None, basetime=None, basedate=None,
-                 segname = None, seglen = None, comments=None,
-                 signame = None)
+                 segname=None, seglen=None, comments=None,
+                 signame=None, sigsegments=None)
 
     Example Usage:
     import wfdb
@@ -492,12 +492,12 @@ def __init__(self, segments = None, layout = None,
     """
 
     # Constructor
-    def __init__(self, segments = None, layout = None,
+    def __init__(self, segments=None, layout=None,
                  recordname=None, nsig=None, fs=None,
                  counterfreq=None, basecounter=None,
                  siglen=None, basetime=None, basedate=None,
-                 segname = None, seglen = None, comments=None,
-                 signame = None):
+                 segname=None, seglen=None, comments=None,
+                 signame=None, sigsegments=None):
 
 
         super(MultiRecord, self).__init__(recordname, nsig,
@@ -508,6 +508,7 @@ def __init__(self, segments = None, layout = None,
         self.segments = segments
         self.segname = segname
         self.seglen = seglen
+        self.sigsegments=sigsegments
 
     # Write a multi-segment header, along with headers and dat files for all segments
     def wrsamp(self):
@@ -585,21 +586,12 @@ def requiredsegments(self, sampfrom, sampto, channels):
             readsegs = list(range(readsegs[0], readsegs[1]+1))
             readsamps = [[0, self.seglen[s]] for s in readsegs]
 
-
-            #print('readsegs: ', readsegs)
-            #print('cumsumlengths: ', cumsumlengths)
-            #print('startseg: ', startseg)
-
             # Starting sample for first segment.
             readsamps[0][0] = sampfrom - ([0] + cumsumlengths)[readsegs[0]-startseg]
 
             # End sample for last segment
             readsamps[-1][1] = sampto - ([0] + cumsumlengths)[readsegs[-1]-startseg]
 
-        #print('\n\nEnd of requiredsegments.')
-        #print('readsegs: ', readsegs)
-        #print('readsamps: ', readsamps)
-
         return (readsegs, readsamps)
 
     # Get the channel numbers to be read from each segment
@@ -981,6 +973,8 @@ def rdheader(recordname, pbdir = None, rdsegments = False):
                     record.segments.append(rdheader(os.path.join(dirname,s), pbdir))
             # Fill in the signame attribute
             record.signame = record.getsignames()
+            # Fill in the sigsegments attribute
+            record.sigsegments = record.getsigsegments()
 
     # Set the comments field
     record.comments = []
