Added checks for callbacks.

FSX · Dec 19, 2012 · 4f1f404 · 4f1f404
1 parent d7f83fe
commit 4f1f404
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/momoko/connection.py b/momoko/connection.py
@@ -185,7 +185,9 @@ def __init__(self,
         self._transaction_status = self.connection.get_transaction_status
         self.ioloop = ioloop or IOLoop.instance()
 
-        self.callback = partial(callback, self)
+        if callback:
+            self.callback = partial(callback, self)
+
         self.ioloop.add_handler(self.fileno, self.io_callback, IOLoop.WRITE)
 
     def io_callback(self, fd=None, events=None):
@@ -197,11 +199,14 @@ def io_callback(self, fd=None, events=None):
             if not isinstance(error, psycopg2.DatabaseError):
                 self.ioloop.update_handler(self.fileno, 0)
 
-            self.callback(error)
+            if self.callback:
+                self.callback(error)
         else:
             if state == POLL_OK:
                 self.ioloop.update_handler(self.fileno, 0)
-                self.callback(None)
+
+                if self.callback:
+                    self.callback(None)
             elif state == POLL_READ:
                 self.ioloop.update_handler(self.fileno, IOLoop.READ)
             elif state == POLL_WRITE:

diff --git a/momoko/utils.py b/momoko/utils.py
@@ -9,7 +9,6 @@
 MIT, see LICENSE for more details.
 """
 
-
 import sys
 from tornado import gen
 from functools import partial