v2.30.3

CHANGELOG.md

-# v2.20.3
+# v2.20.4

satella/__init__.py

-__version__ = '2.20.3'
+__version__ = '2.20.4a1'

satella/db.py

@@ -41,8 +41,14 @@ class transaction:
     def __call__(self, fun):
         @wraps(fun)
         def inner(*args, **kwargs):
-            with self:
-                return fun(*args, **kwargs)
+            with self as cur:
+                if fun.__name__ == fun.__qualname__:
+                    return fun(cur, *args, **kwargs)
+                else:
+                    if not len(args):
+                        return fun(*args)
+                    else:
+                        return fun(*args[0], cur, *args[1:], **kwargs)
         return inner
 
     def __enter__(self):
@@ -51,7 +57,7 @@ class transaction:
 
     @property
     def connection(self):
-        if inspect.isfunction(self.connection):
+        if inspect.isfunction(self._connection):
             return self._connection()
         else:
             return self._connection

tests/test_db.py

@@ -4,28 +4,53 @@ from unittest.mock import Mock
 from satella.db import transaction
 
 
+class RealConnection:
+    def __init__(self):
+        self.cursor_called = 0
+        self.commit_called = 0
+        self.rollback_called = 0
+        self.close_called = 0
+
+    def cursor(self):
+        self.cursor_called += 1
+        return Mock()
+
+    def commit(self):
+        self.commit_called += 1
+
+    def rollback(self):
+        self.rollback_called += 1
+
+    def close(self):
+        self.close_called += 1
+
+
 class TestDB(unittest.TestCase):
-    def test_something(self):
-        class RealConnection:
-            def __init__(self):
-                self.cursor_called = 0
-                self.commit_called = 0
-                self.rollback_called = 0
-                self.close_called = 0
+    def test_db(self):
+        conn = RealConnection()
+
+        @transaction(conn)
+        def test(cur):
+            pass
+
+        test()
 
-            def cursor(self):
-                self.cursor_called += 1
-                return Mock()
+        self.assertEqual(conn.cursor_called, 1)
+        self.assertEqual(conn.commit_called, 1)
+
+    def test_db_imba(self):
+        conn = RealConnection()
 
-            def commit(self):
-                self.commit_called += 1
+        @transaction(conn)
+        def test(cur):
+            raise ValueError()
 
-            def rollback(self):
-                self.rollback_called += 1
+        test()
 
-            def close(self):
-                self.close_called += 1
+        self.assertEqual(conn.cursor_called, 1)
+        self.assertEqual(conn.rollback_called, 1)
 
+    def test_db2(self):
         conn = RealConnection()
         a = transaction(conn)
         with a as cur: