From 24f9093673cc7241e77966d9cdbeca12b8f7b4f2 Mon Sep 17 00:00:00 2001
From: Bruno Alla <bruno.alla@festicket.com>
Date: Fri, 2 Apr 2021 11:34:40 +0100
Subject: [PATCH] fix: update type annotations

---
 django_codemod/visitors/admin.py         |  9 ++++++---
 django_codemod/visitors/models.py        | 11 ++++++++---
 django_codemod/visitors/template_tags.py |  3 ++-
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/django_codemod/visitors/admin.py b/django_codemod/visitors/admin.py
index 3afbf0e3..8f6d2900 100644
--- a/django_codemod/visitors/admin.py
+++ b/django_codemod/visitors/admin.py
@@ -7,6 +7,7 @@
     BaseStatement,
     Call,
     ClassDef,
+    FlattenSentinel,
     FunctionDef,
     ImportFrom,
     ImportStar,
@@ -32,7 +33,9 @@ class InlineHasAddPermissionsTransformer(BaseDjCodemodTransformer):
 
     def leave_ImportFrom(
         self, original_node: ImportFrom, updated_node: ImportFrom
-    ) -> Union[BaseSmallStatement, RemovalSentinel]:
+    ) -> Union[
+        BaseSmallStatement, FlattenSentinel[BaseSmallStatement], RemovalSentinel
+    ]:
         if isinstance(updated_node.names, ImportStar):
             return super().leave_ImportFrom(original_node, updated_node)
         base_cls_matcher = []
@@ -94,7 +97,7 @@ def visit_ClassDef_bases(self, node: ClassDef) -> None:
 
     def leave_ClassDef(
         self, original_node: ClassDef, updated_node: ClassDef
-    ) -> Union[BaseStatement, RemovalSentinel]:
+    ) -> Union[BaseStatement, FlattenSentinel[BaseStatement], RemovalSentinel]:
         self.context.scratch.pop(self.ctx_key_visiting_subclass, None)
         return super().leave_ClassDef(original_node, updated_node)
 
@@ -104,7 +107,7 @@ def is_visiting_subclass(self):
 
     def leave_FunctionDef(
         self, original_node: FunctionDef, updated_node: FunctionDef
-    ) -> Union[BaseStatement, RemovalSentinel]:
+    ) -> Union[BaseStatement, FlattenSentinel[BaseStatement], RemovalSentinel]:
         if (
             self.is_visiting_subclass
             and m.matches(
diff --git a/django_codemod/visitors/models.py b/django_codemod/visitors/models.py
index 31a85386..2112bfcd 100644
--- a/django_codemod/visitors/models.py
+++ b/django_codemod/visitors/models.py
@@ -7,6 +7,7 @@
     BaseSmallStatement,
     BaseStatement,
     Call,
+    FlattenSentinel,
     FunctionDef,
     ImportFrom,
     ImportStar,
@@ -35,7 +36,9 @@ class ModelsPermalinkTransformer(BaseDjCodemodTransformer):
 
     def leave_ImportFrom(
         self, original_node: ImportFrom, updated_node: ImportFrom
-    ) -> Union[BaseSmallStatement, RemovalSentinel]:
+    ) -> Union[
+        BaseSmallStatement, FlattenSentinel[BaseSmallStatement], RemovalSentinel
+    ]:
         if isinstance(updated_node.names, ImportStar):
             return super().leave_ImportFrom(original_node, updated_node)
         if m.matches(
@@ -99,7 +102,7 @@ def visit_FunctionDef(self, node: FunctionDef) -> Optional[bool]:
 
     def leave_FunctionDef(
         self, original_node: FunctionDef, updated_node: FunctionDef
-    ) -> Union[BaseStatement, RemovalSentinel]:
+    ) -> Union[BaseStatement, FlattenSentinel[BaseStatement], RemovalSentinel]:
         if self.visiting_permalink_method:
             for decorator in updated_node.decorators:
                 if m.matches(decorator, self.decorator_matcher):
@@ -122,7 +125,9 @@ def visiting_permalink_method(self):
 
     def leave_Return(
         self, original_node: Return, updated_node: Return
-    ) -> Union[BaseSmallStatement, RemovalSentinel]:
+    ) -> Union[
+        BaseSmallStatement, FlattenSentinel[BaseSmallStatement], RemovalSentinel
+    ]:
         if self.visiting_permalink_method and m.matches(
             updated_node.value, m.Tuple()  # type: ignore
         ):
diff --git a/django_codemod/visitors/template_tags.py b/django_codemod/visitors/template_tags.py
index ad54fa53..82c32011 100644
--- a/django_codemod/visitors/template_tags.py
+++ b/django_codemod/visitors/template_tags.py
@@ -4,6 +4,7 @@
     Assign,
     AssignTarget,
     Decorator,
+    FlattenSentinel,
     ImportFrom,
     ImportStar,
     Module,
@@ -125,7 +126,7 @@ def _gen_decorator_matchers(
 
     def leave_Decorator(
         self, original_node: Decorator, updated_node: Decorator
-    ) -> Union[Decorator, RemovalSentinel]:
+    ) -> Union[Decorator, FlattenSentinel[Decorator], RemovalSentinel]:
         """Update decorator call if all conditions are met."""
         if self.decorators_matcher and m.matches(updated_node, self.decorators_matcher):
             # If we have a decorator matcher, and it matches,