iterating over statement instance

kasun · Nov 19, 2016 · ddadd30 · ddadd30
1 parent ea40e67
commit ddadd30
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/beesql/__init__.py b/beesql/__init__.py
@@ -35,6 +35,9 @@ def __repr__(self):
     def __iter__(self):
         return iter(self.rows)
 
+    def __getitem__(self, key):
+        return self.all()[key]
+
     @property
     def count(self):
         return len(self.rows)

diff --git a/beesql/query/base.py b/beesql/query/base.py
@@ -240,13 +240,21 @@ def __init__(self, query, **kwargs):
         self.query = query
         self.secondary_keywords = []
 
+    def __repr__(self):
+        return '{}: {}'.format(self.__class__, self.get_sql())
+
+    def __iter__(self):
+        rows = self.execute().all()
+        return iter(rows)
+
+    def __getitem__(self, key):
+        rows = self.execute().all()
+        return rows[key]
+
     def add_secondary_keyword(self, keyword):
         heap = self.secondary_keywords
         heapq.heappush(heap, keyword)
 
-    def __repr__(self):
-        return '{}: {}'.format(self.__class__, self.get_sql())
-
     def execute(self):
         res = None
         with self.query.db.connect() as conn: