add the 'signame' field to MultiRecord objects, and the signalsegments() method to get segments that contain a signal name

cx1111 · cx1111 · commit c9add72e495b · 2017-06-12T17:59:08.000-04:00
diff --git a/wfdb/readwrite/_headers.py b/wfdb/readwrite/_headers.py
@@ -345,7 +345,7 @@ def wrheaderfile(self, writefields):
                 recordline = recordline + recfieldspecs[field].delimiter + str(getattr(self, field))
         headerlines.append(recordline)
 
-        # Create segment specification lines 
+        # Create segment specification lines
         segmentlines = self.nseg*['']
         # For both fields, add each of its elements with the delimiter to the appropriate line 
         for field in ['segname', 'segname']:
@@ -361,6 +361,43 @@ 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)
+    # Only works if information about the segments has been read in
+    def signalsegments(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")
+        
+        # Default value = all signal names.
+        if signame is None:
+            signame = self.getsignames()
+
+        if type(signame) == list:
+            return [self.signalsegments(sig) for sig in signame]
+        elif type(signame) == str:
+            sigsegs = []
+            for i in range(self.nseg):
+                if self.segname[i] != '~' and signame in self.segments[i].signame:
+                    sigsegs.append(i)
+            return sigsegs
+        else:
+            raise TypeError('signame must be a string or a list of strings')
+
+    # Get the signal names for the entire record
+    def getsignames(self):
+        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")
+        
+        if self.layout == 'Fixed':
+            for i in range(self.nseg):
+                if self.segname[i] != '~':
+                    signame = self.segments[i].signame
+                    break
+        else:
+            signame = self.segments[0].signame
+        
+        return signame
+
 
 # Regexp objects for reading headers
 
diff --git a/wfdb/readwrite/records.py b/wfdb/readwrite/records.py
@@ -24,7 +24,7 @@ class BaseRecord(object):
     def __init__(self, recordname=None, nsig=None,
                  fs=None, counterfreq=None, basecounter = None,
                  siglen = None, basetime = None, basedate = None,
-                 comments = None):
+                 comments = None, signame=None):
         self.recordname = recordname
         self.nsig = nsig
         self.fs = fs
@@ -34,7 +34,7 @@ def __init__(self, recordname=None, nsig=None,
         self.basetime = basetime
         self.basedate = basedate
         self.comments = comments
-
+        self.signame = signame
 
     # Check whether a single field is valid in its basic form. Does not check compatibility with other fields.
     # ch is only used for signal specification fields, specifying the channels to check. Other channels
@@ -359,7 +359,7 @@ def __init__(self, p_signals=None, d_signals=None,
 
         super(Record, self).__init__(recordname, nsig,
                     fs, counterfreq, basecounter, siglen,
-                    basetime, basedate, comments)
+                    basetime, basedate, comments, signame)
 
         self.p_signals = p_signals
         self.d_signals = d_signals
@@ -380,7 +380,6 @@ def __init__(self, p_signals=None, d_signals=None,
         self.initvalue=initvalue
         self.checksum=checksum
         self.blocksize=blocksize
-        self.signame=signame
 
     # Equal comparison operator for objects of this type
     def __eq__(self, other):
@@ -480,7 +479,8 @@ 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)
+                 segname = None, seglen = None, comments=None,
+                 signame = None)
 
     Example Usage:
     import wfdb
@@ -496,12 +496,13 @@ 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):
+                 segname = None, seglen = None, comments=None,
+                 signame = None):
 
 
         super(MultiRecord, self).__init__(recordname, nsig,
                     fs, counterfreq, basecounter, siglen,
-                    basetime, basedate, comments)
+                    basetime, basedate, comments, signame)
 
         self.layout = layout
         self.segments = segments
@@ -978,6 +979,8 @@ def rdheader(recordname, pbdir = None, rdsegments = False):
                     record.segments.append(None)
                 else:
                     record.segments.append(rdheader(os.path.join(dirname,s), pbdir))
+            # Fill in the signame attribute
+            record.signame = record.getsignames()
 
     # Set the comments field
     record.comments = []
