From 9dfc40e3f3f329965a8ff29b3f5b7becdc2b675e Mon Sep 17 00:00:00 2001
From: Jared Flatow <jflatow@gmail.com>
Date: Sun, 8 May 2011 11:09:26 -0700
Subject: [PATCH] added iterator to index objects

---
 pygit2.c           | 83 +++++++++++++++++++++++++++++++++++++++++++++-
 test/test_index.py |  4 +++
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/pygit2.c b/pygit2.c
index d409207..0f01da2 100644
--- a/pygit2.c
+++ b/pygit2.c
@@ -71,6 +71,12 @@ typedef struct {
     int own_obj:1;
 } Index;
 
+typedef struct {
+  PyObject_HEAD
+  Index *owner;
+  Py_ssize_t i;
+} IndexIter;
+
 typedef struct {
     PyObject_HEAD
     git_index_entry *entry;
@@ -89,12 +95,17 @@ static PyTypeObject TreeType;
 static PyTypeObject BlobType;
 static PyTypeObject TagType;
 static PyTypeObject IndexType;
+static PyTypeObject IndexIterType;
 static PyTypeObject IndexEntryType;
 static PyTypeObject WalkerType;
 static PyTypeObject ReferenceType;
 
 static PyObject *GitError;
 
+static PyObject * IndexIter_new(PyTypeObject *, Index *);
+static void       IndexIter_dealloc(IndexIter *);
+static PyObject * IndexIter_iternext(IndexIter *);
+
 static PyObject *
 Error_type(int err) {
     switch (err) {
@@ -1515,6 +1526,11 @@ Index_contains(Index *self, PyObject *value) {
     return 1;
 }
 
+static PyObject *
+Index_iter(Index *self) {
+  return IndexIter_new(&IndexIterType, self);
+}
+
 static Py_ssize_t
 Index_len(Index *self) {
     return (Py_ssize_t)git_index_entrycount(self->index);
@@ -1647,7 +1663,7 @@ static PyTypeObject IndexType = {
     0,                                         /* tp_clear */
     0,                                         /* tp_richcompare */
     0,                                         /* tp_weaklistoffset */
-    0,                                         /* tp_iter */
+    (getiterfunc)Index_iter,                   /* tp_iter */
     0,                                         /* tp_iternext */
     Index_methods,                             /* tp_methods */
     0,                                         /* tp_members */
@@ -1662,6 +1678,71 @@ static PyTypeObject IndexType = {
     0,                                         /* tp_new */
 };
 
+
+static PyObject *
+IndexIter_new(PyTypeObject *type, Index *owner) {
+  IndexIter *self = PyObject_New(IndexIter, type);
+  if (self != NULL) {
+    Py_INCREF(owner);
+    self->owner = owner;
+    self->i = 0;
+  }
+  return (PyObject *)self;
+}
+
+static void
+IndexIter_dealloc(IndexIter *self) {
+  Py_CLEAR(self->owner);
+  PyObject_Del(self);
+}
+
+static PyObject *
+IndexIter_iternext(IndexIter *self) {
+  PyObject *value = NULL, *entry = NULL;
+
+  if (self->i >= Index_len(self->owner))
+    return NULL;
+
+  value = PyInt_FromSsize_t(self->i++);
+  if (value == NULL)
+    return NULL;
+
+  entry = Index_getitem(self->owner, value);
+  Py_CLEAR(value);
+  return entry;
+}
+
+static PyTypeObject IndexIterType = {
+   PyVarObject_HEAD_INIT(&PyType_Type, 0)
+   "pygit2.IndexIter",                      /* tp_name           */
+   sizeof(IndexIter),                       /* tp_basicsize      */
+   0,                                       /* tp_itemsize       */
+   (destructor)IndexIter_dealloc ,          /* tp_dealloc        */
+   0,                                       /* tp_print          */
+   0,                                       /* tp_getattr        */
+   0,                                       /* tp_setattr        */
+   0,                                       /* tp_compare        */
+   0,                                       /* tp_repr           */
+   0,                                       /* tp_as_number      */
+   0,                                       /* tp_as_sequence    */
+   0,                                       /* tp_as_mapping     */
+   0,                                       /* tp_hash           */
+   0,                                       /* tp_call           */
+   0,                                       /* tp_str            */
+   PyObject_GenericGetAttr,                 /* tp_getattro       */
+   0,                                       /* tp_setattro       */
+   0,                                       /* tp_as_buffer      */
+   Py_TPFLAGS_DEFAULT |
+   Py_TPFLAGS_BASETYPE,                     /* tp_flags          */
+   0,                                       /* tp_doc            */
+   0,                                       /* tp_traverse       */
+   0,                                       /* tp_clear          */
+   0,                                       /* tp_richcompare    */
+   0,                                       /* tp_weaklistoffset */
+   PyObject_SelfIter,                       /* tp_iter           */
+   (iternextfunc)IndexIter_iternext,        /* tp_iternext       */
+ };
+
 static void
 IndexEntry_dealloc(IndexEntry *self) {
     self->ob_type->tp_free((PyObject*)self);
diff --git a/test/test_index.py b/test/test_index.py
index 1da7558..2d88f4d 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -89,6 +89,10 @@ class IndexTest(utils.RepoTestCase):
         sha = self.repo.index.create_tree()
         self.assertEqual(sha, 'fd937514cb799514d4b81bb24c5fcfeb6472b245')
 
+    def test_iter(self):
+        index = self.repo.index
+        entries = [index[x] for x in xrange(len(index))]
+        self.assertNotEqual(list(index), entries)
 
 if __name__ == '__main__':
     unittest.main()