🔥 Fix import added too early (#17)

* Update tests to reproduce the issue * Fix bug with import added too early in the file * Add refactoring from Sourcery * Made code more readable Co-authored-by: luttik <[email protected]>
Luttik · Dec 5, 2021 · cd11824 · cd11824
1 parent af7358c
commit cd11824
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 19 deletions.
diff --git a/src/auto_optional/code_processing.py b/src/auto_optional/code_processing.py
@@ -99,27 +99,66 @@ def leave_Module(
 
         Is only used to add the ``typing.Optional`` import statement when missing.
         """
-        if self.optional_import_checker.import_names and self.optional_import_needed:
-            import_statement = cst.SimpleStatementLine(
-                body=[
-                    cst.ImportFrom(
-                        module=cst.Name(
-                            value="typing",
-                        ),
-                        names=[
-                            cst.ImportAlias(
-                                name=cst.Name(
-                                    value="Optional",
-                                ),
+        if (
+            not self.optional_import_checker.import_names
+            or not self.optional_import_needed
+        ):
+            return updated_node
+
+        import_statement = self.__build_import_statement()
+        import_statement_location = self.__find_line_for_optional_import_insertion(
+            updated_node
+        )
+        return updated_node.with_changes(
+            body=[
+                *updated_node.body[:import_statement_location],
+                import_statement,
+                *updated_node.body[import_statement_location:],
+            ]
+        )
+
+    def __build_import_statement(self) -> cst.SimpleStatementLine:
+        return cst.SimpleStatementLine(
+            body=[
+                cst.ImportFrom(
+                    module=cst.Name(
+                        value="typing",
+                    ),
+                    names=[
+                        cst.ImportAlias(
+                            name=cst.Name(
+                                value="Optional",
                             ),
-                        ],
+                        ),
+                    ],
+                )
+            ]
+        )
+
+    def __find_line_for_optional_import_insertion(
+        self, updated_node: cst.Module
+    ) -> int:
+        insert_at = 0
+        for index, body_part in enumerate(updated_node.body):
+            if (
+                isinstance(body_part, cst.SimpleStatementLine)
+                and len(body_part.children) > 0
+                and (
+                    (  # module level docstring
+                        index == 0
+                        and m.matches(body_part.children[0], m.Expr(m.SimpleString()))
                     )
-                ]
-            )
-            return updated_node.with_changes(
-                body=[import_statement, *updated_node.body]
-            )
-        return updated_node
+                    or m.matches(  # future import
+                        body_part.children[0],
+                        m.ImportFrom(module=m.Name("__future__")),
+                    )
+                )
+            ):
+                # module docstring
+                insert_at += 1
+            elif index > 1:
+                break
+        return insert_at
 
     def __check_if__union_none(self, annotation: cst.BaseExpression) -> bool:
         #  Return if it is a Union[..., None, ...] statement

diff --git a/test/simple-tests.yaml b/test/simple-tests.yaml
@@ -94,6 +94,41 @@ tests:
       def bla(foo: Optional[str] = None):
           ...
 
+  - name: test_with_module_docstring
+    before: |
+      """Something"""
+      def f(s: str = None):
+          ...
+    after: |
+      """Something"""
+      from typing import Optional
+      def f(s: Optional[str] = None):
+          ...
+
+  - name: test_with_future_import
+    before: |
+      from __future__ import annotations
+      def f(s: str = None):
+          ...
+    after: |
+      from __future__ import annotations
+      from typing import Optional
+      def f(s: Optional[str] = None):
+          ...
+
+  - name: test_docstring_with_future_import
+    before: |
+      """Something"""
+      from __future__ import annotations
+      def f(s: str = None):
+          ...
+    after: |
+      """Something"""
+      from __future__ import annotations
+      from typing import Optional
+      def f(s: Optional[str] = None):
+          ...
+
   - name: test_nochange_typing_as_import
     unchanged: |
       import typing as t