From a3cb4917b7abb7c578b549b9687a1039d847435d Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Fri, 27 Dec 2019 01:02:11 +0300
Subject: [PATCH] libcst-based script to condition code in stubs

---
 scripts/remove_django_2_2_code.py | 44 +++++++++++++++++++++++++++++++
 scripts/stub.pyi                  | 11 ++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 scripts/remove_django_2_2_code.py
 create mode 100644 scripts/stub.pyi

diff --git a/scripts/remove_django_2_2_code.py b/scripts/remove_django_2_2_code.py
new file mode 100644
index 000000000..e1101bcf7
--- /dev/null
+++ b/scripts/remove_django_2_2_code.py
@@ -0,0 +1,44 @@
+import libcst
+import libcst.matchers as m
+from libcst import If, IndentedBlock, BaseSuite
+from libcst.codemod import ContextAwareTransformer, CodemodContext
+
+STUB = 'stub.pyi'
+IS_3_0 = True
+
+
+def is_django_3_0(if_node: If) -> bool:
+    return m.matches(if_node.test, m.Name(value='DJANGO_3_0'))
+
+
+class DjangoVersionTransformer(ContextAwareTransformer):
+    def leave_IndentedBlock(
+            self, original_node: "IndentedBlock", updated_node: "IndentedBlock"
+    ) -> "BaseSuite":
+        modified_body = []
+        for statement in original_node.body:
+            if not m.matches(statement, m.If(test=m.Name(value='DJANGO_3_0'))):
+                modified_body.append(statement)
+                continue
+
+            if_statement = libcst.ensure_type(statement, If)
+            if IS_3_0:
+                modified_body.extend(if_statement.body.body)
+            else:
+                if if_statement.orelse is not None:
+                    modified_body.extend(if_statement.orelse.body.body)
+                else:
+                    continue
+
+        return updated_node.with_changes(body=modified_body)
+
+
+with open(STUB) as f:
+    contents = f.read()
+
+tree = libcst.parse_module(contents)
+context = CodemodContext()
+transformer = DjangoVersionTransformer(context)
+
+modified_tree = tree.visit(transformer)
+print(modified_tree.code)
diff --git a/scripts/stub.pyi b/scripts/stub.pyi
new file mode 100644
index 000000000..7f07d51c1
--- /dev/null
+++ b/scripts/stub.pyi
@@ -0,0 +1,11 @@
+DJANGO_2_2 = 1
+DJANGO_3_0 = 2
+
+class MyClass:
+    if DJANGO_3_0:
+        def django_3_0_func(self, a: int) -> int: ...
+    else:
+        def django_2_2_func(self, a: str) -> str: ...
+
+    if DJANGO_3_0:
+        def only_django_3_0_func(self, a: int) -> int: ...