From 67c99434e5bad1108081addd1b9b65c078a96ff5 Mon Sep 17 00:00:00 2001
From: Seth Yastrov <syastrov@gmail.com>
Date: Tue, 19 Feb 2019 07:55:50 +0100
Subject: [PATCH] Add better typings plus test for transaction.atomic.

- All cases are handled, including bare decorator (@transaction.atomic).
- Decorated function's signature is preserved when type-checking.
---
 django-stubs/db/transaction.pyi      | 21 +++++++-----
 test-data/typecheck/transaction.test | 50 ++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 test-data/typecheck/transaction.test

diff --git a/django-stubs/db/transaction.pyi b/django-stubs/db/transaction.pyi
index 65172cb8e..d8f9cb078 100644
--- a/django-stubs/db/transaction.pyi
+++ b/django-stubs/db/transaction.pyi
@@ -1,5 +1,4 @@
-from contextlib import ContextDecorator
-from typing import Any, Callable, Optional, Union, Iterator, overload, ContextManager
+from typing import Any, Callable, Optional, overload, TypeVar
 
 from django.db import ProgrammingError
 
@@ -18,19 +17,23 @@ def get_rollback(using: None = ...) -> bool: ...
 def set_rollback(rollback: bool, using: Optional[str] = ...) -> None: ...
 def on_commit(func: Callable, using: None = ...) -> None: ...
 
-class Atomic(ContextDecorator):
+_C = TypeVar("_C", bound=Callable)  # Any callable
+
+# Don't inherit from ContextDecorator, so we can provide a more specific signature for __call__
+class Atomic:
     using: Optional[str] = ...
     savepoint: bool = ...
     def __init__(self, using: Optional[str], savepoint: bool) -> None: ...
+    # When decorating, return the decorated function as-is, rather than clobbering it as ContextDecorator does.
+    def __call__(self, func: _C) -> _C: ...
     def __enter__(self) -> None: ...
     def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None: ...
 
+# Bare decorator
 @overload
-def atomic() -> Atomic: ...
-@overload
-def atomic(using: Optional[str] = ...,) -> ContextManager[Atomic]: ...
-@overload
-def atomic(using: Callable = ...) -> Callable: ...
+def atomic(using: _C) -> _C: ...
+
+# Decorator or context-manager with parameters
 @overload
-def atomic(using: Optional[str] = ..., savepoint: bool = ...) -> ContextManager[Atomic]: ...
+def atomic(using: Optional[str] = None, savepoint: bool = True) -> Atomic: ...
 def non_atomic_requests(using: Callable = ...) -> Callable: ...
diff --git a/test-data/typecheck/transaction.test b/test-data/typecheck/transaction.test
new file mode 100644
index 000000000..9fb83ca7f
--- /dev/null
+++ b/test-data/typecheck/transaction.test
@@ -0,0 +1,50 @@
+[CASE test_transaction_atomic]
+
+from django.db import transaction
+
+with transaction.atomic():
+    pass
+
+with transaction.atomic(using="mydb"):
+    pass
+
+with transaction.atomic(using="mydb", savepoint=False):
+    pass
+
+@transaction.atomic()
+def decorated_func(param1: str, param2: int) -> bool:
+    pass
+
+# Ensure that the function's type is preserved
+reveal_type(decorated_func)  # E: Revealed type is 'def (param1: builtins.str, param2: builtins.int) -> builtins.bool'
+
+@transaction.atomic(using="mydb")
+def decorated_func_using(param1: str, param2: int) -> bool:
+    pass
+
+# Ensure that the function's type is preserved
+reveal_type(decorated_func_using)  # E: Revealed type is 'def (param1: builtins.str, param2: builtins.int) -> builtins.bool'
+
+class ClassWithAtomicMethod:
+    # Bare decorator
+    @transaction.atomic
+    def atomic_method1(self, abc: int) -> str:
+        pass
+
+    @transaction.atomic(savepoint=True)
+    def atomic_method2(self):
+        pass
+
+    @transaction.atomic(using="db", savepoint=True)
+    def atomic_method3(self, myparam: str) -> int:
+        pass
+
+ClassWithAtomicMethod().atomic_method1("abc") # E: Argument 1 to "atomic_method1" of "ClassWithAtomicMethod" has incompatible type "str"; expected "int"
+
+# Ensure that the method's type is preserved
+reveal_type(ClassWithAtomicMethod().atomic_method1) # E: Revealed type is 'def (abc: builtins.int) -> builtins.str'
+
+# Ensure that the method's type is preserved
+reveal_type(ClassWithAtomicMethod().atomic_method3) # E: Revealed type is 'def (myparam: builtins.str) -> builtins.int'
+
+[out]
\ No newline at end of file