Add iterator and context manager methods to GrampsCursor object and use them in GrampsDbBase and GrampsDBDir.

svn: r12660

src/gen/db/base.py

@@ -24,7 +24,7 @@
 Base class for the GRAMPS databases. All database interfaces should inherit
 from this class.
 """
-
+from __future__ import with_statement
 #-------------------------------------------------------------------------
 #
 # libraries
@@ -1347,7 +1347,7 @@ class GrampsDbBase(Callback):
 
     def all_handles(self, table):
         return table.keys()
-
+        
     def get_person_handles(self, sort_handles=True):
         """
         Return a list of database handles, one handle for each Person in
@@ -1357,19 +1357,21 @@ class GrampsDbBase(Callback):
         """
         if self.db_is_open:
             if sort_handles:
-                slist = []
+                with self.get_person_cursor() as cursor:
-                cursor = self.get_person_cursor()
+                    slist = sorted((data[1][3][3], data[0]) for data in cursor)
-                data = cursor.first()
-                while data:
-                    slist.append((data[1][3][3], data[0]))
-                    data = cursor.next()
-                cursor.close()
-                slist.sort()
                 return [x[1] for x in slist]
             else:
                 return self.all_handles(self.person_map)
         return []
 
+    def iter_person_handles(self):
+        """
+        Return an iterator over handles for Persons in the database
+        """
+        with self.get_person_cursor() as cursor:
+            for data in cursor:
+                yield data[0]
+
     def get_place_handles(self, sort_handles=True):
         """
         Return a list of database handles, one handle for each Place in
@@ -1377,21 +1379,24 @@ class GrampsDbBase(Callback):
         
         If sort_handles is True, the list is sorted by Place title.
         """
+        print "base.py: get_place_handles"
         if self.db_is_open:
             if sort_handles:
-                slist = []
+                with self.get_place_cursor() as cursor:
-                cursor = self.get_place_cursor()
+                    slist = sorted(((data[1][2], data[0])) for data in cursor)
-                data = cursor.first()
-                while data:
-                    slist.append((data[1][2], data[0]))
-                    data = cursor.next()
-                cursor.close()
-                slist.sort()
                 return [x[1] for x in slist]
             else:
                 return self.all_handles(self.place_map)
         return []
-
+        
+    def iter_place_handles(self):
+        """
+        Return an iterator over handles for Places in the database
+        """
+        with self.get_place_cursor() as cursor:
+            for data in cursor:
+                yield data[0]
+                
     def get_source_handles(self, sort_handles=True):
         """
         Return a list of database handles, one handle for each Source in
@@ -1399,13 +1404,22 @@ class GrampsDbBase(Callback):
         
          If sort_handles is True, the list is sorted by Source title.
         """
+        print "base.py: get_source_handles"
         if self.db_is_open:
             handle_list = self.all_handles(self.source_map)
             if sort_handles:
-                handle_list.sort(self.__sortbysource)
+                handle_list.sort(key=self.__sortbysource_key)
             return handle_list
         return []
-
+        
+    def iter_source_handles(self):
+        """
+        Return an iterator over handles for Sources in the database
+        """
+        with self.get_source_cursor() as cursor:
+            for data in cursor:
+                yield data[0]
+                
     def get_media_object_handles(self, sort_handles=True):
         """
         Return a list of database handles, one handle for each MediaObject in
@@ -1413,12 +1427,21 @@ class GrampsDbBase(Callback):
         
         If sort_handles is True, the list is sorted by title.
         """
+        print "base.py: get_media_object_handles"
         if self.db_is_open:
             handle_list = self.all_handles(self.media_map)
             if sort_handles:
-                handle_list.sort(self.__sortbymedia)
+                handle_list.sort(key=self.__sortbymedia_key)
             return handle_list
         return []
+        
+    def iter_media_handles(self):
+        """
+        Return an iterator over handles for Media in the database
+        """
+        with self.get_media_cursor() as cursor:
+            for data in cursor:
+                yield data[0]
 
     def get_event_handles(self):
         """
@@ -1428,6 +1451,14 @@ class GrampsDbBase(Callback):
         if self.db_is_open:
             return self.all_handles(self.event_map)
         return []
+        
+    def iter_event_handles(self):
+        """
+        Return an iterator over handles for Events in the database
+        """
+        with self.get_event_cursor() as cursor:
+            for data in cursor:
+                yield data[0]        
 
     def get_family_handles(self):
         """
@@ -1437,6 +1468,14 @@ class GrampsDbBase(Callback):
         if self.db_is_open:
             return self.all_handles(self.family_map)
         return []
+        
+    def iter_family_handles(self):
+        """
+        Return an iterator over handles for Families in the database
+        """
+        with self.get_family_cursor() as cursor:
+            for data in cursor:
+                yield data[0]        
 
     def get_repository_handles(self):
         """
@@ -1446,6 +1485,14 @@ class GrampsDbBase(Callback):
         if self.db_is_open:
             return self.all_handles(self.repository_map)
         return []
+        
+    def iter_repository_handles(self):
+        """
+        Return an iterator over handles for Repositories in the database
+        """
+        with self.get_repository_cursor() as cursor:
+            for data in cursor:
+                yield data[0]        
 
     def get_note_handles(self):
         """
@@ -1455,6 +1502,14 @@ class GrampsDbBase(Callback):
         if self.db_is_open:
             return self.all_handles(self.note_map)
         return []
+        
+    def iter_note_handles(self):
+        """
+        Return an iterator over handles for Notes in the database
+        """
+        with self.get_note_cursor() as cursor:
+            for data in cursor:
+                yield data[0]
 
     def get_gramps_ids(self, obj_key):
         key2table = {
@@ -2236,16 +2291,27 @@ class GrampsDbBase(Callback):
         return locale.strcoll(self.place_map.get(str(first))[2], 
                               self.place_map.get(str(second))[2])
 
+    def __sortbyplace_key(self, place):
+        return locale.strxfrm(self.place_map.get(str(place))[2])
+
     def __sortbysource(self, first, second):
         source1 = unicode(self.source_map[str(first)][2])
         source2 = unicode(self.source_map[str(second)][2])
         return locale.strcoll(source1, source2)
+        
+    def __sortbysource_key(self, key):
+        source = unicode(self.source_map[str(key)][2])
+        return locale.strxfrm(source)
 
     def __sortbymedia(self, first, second):
         media1 = self.media_map[str(first)][4]
         media2 = self.media_map[str(second)][4]
         return locale.strcoll(media1, media2)
 
+    def __sortbymedia_key(self, key):
+        media = self.media_map[str(key)][4]
+        return locale.strxfrm(media)
+
     def set_mediapath(self, path):
         """Set the default media path for database, path should be utf-8."""
         if (self.metadata is not None) and (not self.readonly):

src/gen/db/cursor.py

@@ -18,6 +18,19 @@
 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 #
 
+#-------------------------------------------------------------------------
+#
+# Standard python modules
+#
+#-------------------------------------------------------------------------
+import cPickle as pickle
+
+#-------------------------------------------------------------------------
+#
+# GrampsCursor class
+#
+#-------------------------------------------------------------------------
+
 class GrampsCursor(object):
     """
     Provide a basic iterator that allows the user to cycle through
@@ -30,6 +43,13 @@ class GrampsCursor(object):
     database. If multiple passes are needed, multiple cursors
     should be used.
     """
+    
+    def __init__(self):
+        """
+        Instantiate the object. Note, this method should be overridden in
+        derived classes that properly set self.cursor and self.source
+        """
+        self.cursor = self.source = None
 
     def first(self):
         """
@@ -42,6 +62,10 @@ class GrampsCursor(object):
 
         If no data is available, None is returned.
         """
+        
+        data = self.cursor.first()
+        if data:
+            return (data[0], pickle.loads(data[1]))
         return None
 
     def next(self):
@@ -55,7 +79,17 @@ class GrampsCursor(object):
 
         None is returned when no more data is available.
         """
+        
+        data = self.cursor.next()
+        if data:
+            return (data[0], pickle.loads(data[1]))
         return None
+        
+    def delete(self):
+        """
+        Delete the data at the current cursor position
+        """
+        self.cursor.delete()        
 
     def close(self):
         """
@@ -64,11 +98,34 @@ class GrampsCursor(object):
         This should be called when the user is finished using the cursor, 
         freeing up the cursor's resources.
         """
-        raise NotImplementedError
+        self.cursor.close()
     
     def get_length(self):
         """
         Return the number of records in the table referenced by the cursor.
         """
-        raise NotImplementedError
+        return self.source.stat()['ndata']
+        
+    def __iter__(self):
+        """
+        Iterator
+        """
+        
+        data = self.first()
+        while data:
+            yield data
+            data = self.next()
+        
+    def __enter__(self):
+        """
+        Context manager enter method
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Context manager exit method
+        """
+        self.close()
+        return exc_type is None
 

src/gen/db/dbdir.py

@@ -117,27 +117,6 @@ class GrampsDBDirCursor(GrampsCursor):
         self.cursor = source.db.cursor(txn)
         self.source = source
         
-    def first(self):
-        d = self.cursor.first()
-        if d:
-            return (d[0], pickle.loads(d[1]))
-        return None
-
-    def next(self):
-        d = self.cursor.next()
-        if d:
-            return (d[0], pickle.loads(d[1]))
-        return None
-
-    def close(self):
-        self.cursor.close()
-
-    def delete(self):
-        self.cursor.delete()
-        
-    def get_length(self):
-        return self.source.stat()['ndata']
-
 #-------------------------------------------------------------------------
 #
 # GrampsDBDirAssocCursor
@@ -149,27 +128,6 @@ class GrampsDBDirAssocCursor(GrampsCursor):
         self.cursor = source.cursor(txn)
         self.source = source
         
-    def first(self):
-        d = self.cursor.first()
-        if d:
-            return (d[0], pickle.loads(d[1]))
-        return None
-
-    def next(self):
-        d = self.cursor.next()
-        if d:
-            return (d[0], pickle.loads(d[1]))
-        return None
-
-    def close(self):
-        self.cursor.close()
-
-    def delete(self):
-        self.cursor.delete()
-        
-    def get_length(self):
-        return self.source.stat()['ndata']
-
 #-------------------------------------------------------------------------
 #
 # GrampsDBDirDupCursor