implementing superevents() method

ligo/gracedb/rest.py

@@ -876,31 +876,60 @@ class GraceDb(GsiRest):
         return self.get(self.templates['superevent-detail-template'].format(
             superevent_id=superevent_id))
 
-    def superevents(self, query=None, orderby=None, count=None, columns=None):
+    def superevents(self, query='', orderby=[], count=None, columns=[]):
         """
         Get an iterator of superevents in response to a query.
-        NOT YET IMPLEMENTED
+
+        Arguments:
+            query: query string for filtering superevents (same as on the
+                   web interface)
+            orderby: list of strings corresponding to attribute(s) of the
+                     superevents used to order the results (optional).
+                     Available options: created, t_start, t_0, t_end, is_gw,
+                     id, preferred_event. Default is ascending order, but
+                     prefix any option with "-" to apply descending order.
+            count: each generator iteration will yield this many objects
+                   (optional; default determined by the server)
+            columns: which attributes of the superevents to return
+                     (default: all).
+
+        Example:
+
+            >>> g = GraceDb()
+            >>> for s in g.superevents(query='is_gw=True', orderby=['-preferred_event'], columns='superevent_id,events'):
+            ...     print(s['superevent_id'])
         """
-        raise NotImplementedError("TBD")
-        #uri = self.links['superevents']
-        #qdict = {}
-        #if query:   qdict['query'] = query
-        #if count:   qdict['count'] = count
-        #if orderby: qdict['orderby'] = orderby
-        #if columns: qdict['columns'] = columns
-        #if qdict:
-        #    uri += "?" + urllib.urlencode(qdict)
-
-        #while uri:
-        #    response = self.get(uri).json()
-        #    print(uri)
-        #    print(response)
-        #    events = response.get('events',[])
-        #    print(events)
-        #    uri = response.get('links',{}).get('next')
-        #    print(uri)
-        #    for event in events:
-        #         yield event
+        # If columns is a comma-separated string, split it to a list
+        if isinstance(columns, six.string_types):
+            columns = columns.split(',')
+
+        # If orderby is a list (should be), convert it to a comma-separated
+        # string (that's what the server expects)
+        if isinstance(orderby, list):
+            orderby = ",".join(orderby)
+
+        # Get URI
+        uri = self.links['superevents']
+
+        # Compile URL parameters
+        qdict = {}
+        if query:   qdict['query'] = query
+        if count:   qdict['count'] = count
+        if orderby: qdict['sort'] = orderby
+        if qdict:
+            uri += "?" + urlencode(qdict)
+
+        # Get superevent information and construct a generator
+        while uri:
+            response = self.get(uri).json()
+            superevents = response.get('superevents',[])
+            uri = response.get('links',{}).get('next')
+            for superevent in superevents:
+                # If columns are specified, only return specific values
+                if columns:
+                    yield {k: superevent[k] for k in columns}
+                else:
+                    yield superevent
 
     @event_or_superevent
     def files(self, object_id, filename="", *args, **kwargs):