fix: update type annotations

browniebroke · Apr 2, 2021 · 24f9093 · 24f9093
1 parent ca35714
commit 24f9093
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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,